update testing (#395)

Author: arogozhnikov

.github/workflows/run_tests.yml

@@ -8,13 +8,12 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        # TF didn't support 3.13, add when supported
-        python-version: ['3.9', '3.11', '3.12']
-        # currently there is conflict between tf, oneflow and paddle in protobuf versions.
+        python-version: ['3.10', '3.11', '3.13']
+        # there was (and likely still is) conflict between tf, oneflow and paddle in protobuf versions.
         # cupy is not tested because it demands gpu
         # oneflow testing is dropped, see details at https://github.com/Oneflow-Inc/oneflow/issues/10340
-        # paddle was switched off because of divergence with numpy in py3.10, paddle==2.6.1
-        frameworks: ['numpy pytorch tensorflow jax']
+        # paddle  testing is dropped because of divergence with numpy in py3.10, paddle==2.6.1
+        frameworks: ['numpy pytorch tensorflow jax mlx']
         # The last pytensor release for 3.9 does not have a pin to numpy < 2, but not compatible with numpy 2,
         # so we skip it in testing
         include:

README.md

@@ -311,16 +311,24 @@ Einops works with ...
 - [tinygrad](https://github.com/tinygrad/tinygrad) (community)
 - [pytensor](https://github.com/pymc-devs/pytensor) (community)
 
-Additionally, einops can be used with any framework that supports
+
+
+```python
+from einops import rearrange  
+=> from einops.array_api import rearrange
+```
+
+But actually it is even better: einops can be used with *any* framework that supports
 [Python array API standard](https://data-apis.org/array-api/latest/API_specification/index.html),
-which includes
+to name a few:
 
 - numpy >= 2.0
 - [MLX](https://github.com/ml-explore/mlx)  # yes, einops works with apple's framework
 - [pydata/sparse](https://github.com/pydata/sparse) >= 0.15 # and works with sparse tensors
 - [cubed](https://github.com/cubed-dev/cubed) # and with distributed tensors too
 - [quantco/ndonnx](https://github.com/Quantco/ndonnx)
-- recent releases of jax and cupy.
+- jax
+- cupy
 - dask is supported via [array-api-compat](https://github.com/data-apis/array-api-compat)
 
 

einops/tests/run_tests.py

@@ -34,8 +34,11 @@ def main():
         "paddle": ["paddlepaddle"],
         "oneflow": ["oneflow==0.9.0"],
         "pytensor": ["pytensor"],
-        "mlx": ["mlx"],
     }
+    if sys.platform == "darwin":
+        framework_name2installation["mlx"] = ["mlx"]
+    if sys.platform.startswith("linux"):
+        framework_name2installation["mlx"] = ["mlx[cpu]"]
 
     usage = f"""
     Usage:   python -m einops.tests.run_tests <frameworks> [--pip-install]