Moving tests into source tree (#338)

* delete unused code to detect nvidia-smi

* use corrent term flag -> envvar

* move tests to source tree

* exclude .devcontainer and scripts folders from pypi distribution

* set ruff's target-version to python 3.8

* remove compat with python < 3.6

* remove parametrized without part of tests

* catch Exception, not BaseException

* remove dependency on parameterized

* move test entrypoint to main codebase

* minor

* update jax installation instructions

* change testing instructions to python -m

* minor reformatting

* fix flag detection

* add testing of notebooks as a special script

* add pytest to requirements for notebook testing

.github/workflows/run_tests.yml

@@ -8,7 +8,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ['3.8', '3.10']
+        python-version: ['3.8', '3.10', '3.12']
         # currently there 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
@@ -27,4 +27,4 @@ jobs:
           pip install ruff==0.6.5 && ruff check . && ruff format . --check
       - name: Run tests
         run: |
-          python test.py ${{ matrix.frameworks }}
+          pip install -e . && python -m einops.tests.run_tests ${{ matrix.frameworks }} --pip-install

.github/workflows/test_notebooks.yml

@@ -0,0 +1,25 @@
+name: Test notebooks
+
+on: [push, pull_request]
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.8', '3.12']
+
+    steps:
+      - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: pip
+      - name: Install dependencies to run/test notebooks
+        run: |
+          pip install nbformat nbconvert jupyter pillow pytest numpy torch tensorflow
+      - name: Testing notebooks
+        run: |
+          pip install -e . && pytest scripts/test_notebooks.py

einops/array_api.py

@@ -116,7 +116,7 @@ def unpack(tensor: Tensor, packed_shapes: List[Shape], pattern: str) -> List[Ten
             )
             for i, element_shape in enumerate(packed_shapes)
         ]
-    except BaseException:
+    except Exception:
         # this hits if there is an error during reshapes, which means passed shapes were incorrect
         raise RuntimeError(
             f'Error during unpack(..., "{pattern}"): could not split axis of size {split_positions[-1]}'

einops/packing.py

@@ -182,7 +182,7 @@ def unpack(tensor: Tensor, packed_shapes: List[Shape], pattern: str) -> List[Ten
             )
             for i, element_shape in enumerate(packed_shapes)
         ]
-    except BaseException:
+    except Exception:
         # this hits if there is an error during reshapes, which means passed shapes were incorrect
         raise RuntimeError(
             f'Error during unpack(..., "{pattern}"): could not split axis of size {split_positions[-1]}'

tests/__init__.py → einops/tests/__init__.py

@@ -1,3 +1,8 @@
+"""
+Common utils for testing.
+These functions allow testing only some frameworks, not all.
+"""
+
 import logging
 import os
 from functools import lru_cache
@@ -26,14 +31,22 @@ def find_names_of_all_frameworks() -> List[str]:
     return [b.framework_name for b in backend_subclasses]
 
 
-FLAG_NAME = "EINOPS_TEST_BACKENDS"
+ENVVAR_NAME = "EINOPS_TEST_BACKENDS"
+
+
+def unparse_backends(backend_names: List[str]) -> Tuple[str, str]:
+    _known_backends = find_names_of_all_frameworks()
+    for backend_name in backend_names:
+        if backend_name not in _known_backends:
+            raise RuntimeError(f"Unknown framework: {backend_name}")
+    return ENVVAR_NAME, ",".join(backend_names)
 
 
 @lru_cache(maxsize=1)
 def parse_backends_to_test() -> List[str]:
-    if FLAG_NAME not in os.environ:
-        raise RuntimeError(f"Testing frameworks were not specified, flag {FLAG_NAME} not set")
-    parsed_backends = os.environ[FLAG_NAME].split(",")
+    if ENVVAR_NAME not in os.environ:
+        raise RuntimeError(f"Testing frameworks were not specified, env var {ENVVAR_NAME} not set")
+    parsed_backends = os.environ[ENVVAR_NAME].split(",")
     _known_backends = find_names_of_all_frameworks()
     for backend_name in parsed_backends:
         if backend_name not in _known_backends:
@@ -43,19 +56,12 @@ def parse_backends_to_test() -> List[str]:
 
 
 def is_backend_tested(backend: str) -> bool:
+    """Used to skip test if corresponding backend is not tested"""
     if backend not in find_names_of_all_frameworks():
         raise RuntimeError(f"Unknown framework {backend}")
     return backend in parse_backends_to_test()
 
 
-def unparse_backends(backend_names: List[str]) -> Tuple[str, str]:
-    _known_backends = find_names_of_all_frameworks()
-    for backend_name in backend_names:
-        if backend_name not in _known_backends:
-            raise RuntimeError(f"Unknown framework: {backend_name}")
-    return FLAG_NAME, ",".join(backend_names)
-
-
 def collect_test_backends(symbolic=False, layers=False) -> List[_backends.AbstractBackend]:
     """
     :param symbolic: symbolic or imperative frameworks?

test.py → einops/tests/run_tests.py

@@ -1,15 +1,10 @@
 """
-Usage: python test.py <frameworks>
-
-1. Installs part of dependencies (make sure `which pip` points to correct location)
-2. Installs current version of einops in editable mode
-3. Runs the tests
+Runs tests that are appropriate for framework.
 """
 
 import os
-import shutil
 import sys
-from subprocess import Popen, PIPE
+from subprocess import Popen
 from pathlib import Path
 
 __author__ = "Alex Rogozhnikov"
@@ -18,25 +13,20 @@
 def run(cmd, **env):
     # keeps printing output when testing
     cmd = cmd.split(" ") if isinstance(cmd, str) else cmd
+    print("running:", cmd)
     p = Popen(cmd, cwd=str(Path(__file__).parent), env={**os.environ, **env})
     p.communicate()
     return p.returncode
 
 
-# check we have nvidia-smi
-have_cuda = False
-if shutil.which("nvidia-smi") is not None:
-    output, _ = Popen("nvidia-smi".split(" "), stdout=PIPE).communicate()
-    if b"failed because" not in output:
-        have_cuda = True
-
-
 def main():
-    _executable, *frameworks = sys.argv
+    _executable, *args = sys.argv
+    frameworks = [x for x in args if x != "--pip-install"]
+    pip_install_is_set = "--pip-install" in args
     framework_name2installation = {
         "numpy": ["numpy"],
         "torch": ["torch --index-url https://download.pytorch.org/whl/cpu"],
-        "jax": ["jax[cpu]", "jaxlib", "flax"],
+        "jax": ["jax[cpu]", "flax"],
         "tensorflow": ["tensorflow"],
         "cupy": ["cupy"],
         # switch to stable paddlepaddle, because of https://github.com/PaddlePaddle/Paddle/issues/63927
@@ -46,10 +36,12 @@ def main():
     }
 
     usage = f"""
-    Usage:   python test.py <frameworks>
-    Example: python test.py numpy pytorch
+    Usage:   python -m einops.tests.run_tests <frameworks> [--pip-install]
+    Example: python -m einops.tests.run_tests numpy pytorch --pip-install
 
     Available frameworks: {list(framework_name2installation)}
+    When --pip-install is set, auto-installs requirements with pip.
+     (make sure which pip points to right pip)
     """
     if len(frameworks) == 0:
         print(usage)
@@ -66,33 +58,24 @@ def main():
             print(usage)
             raise RuntimeError(f"Unrecognized frameworks: {wrong_frameworks}")
 
-    other_dependencies = [
-        "nbformat",
-        "nbconvert",
-        "jupyter",
-        "parameterized",
-        "pillow",
-        "pytest",
-    ]
-    for framework in frameworks:
-        print(f"Installing {framework}")
-        pip_instructions = framework_name2installation[framework]
-        assert 0 == run("pip install {} --progress-bar off".format(" ".join(pip_instructions)))
-
-    print("Install testing infra")
-    assert 0 == run("pip install {} --progress-bar off".format(" ".join(other_dependencies)))
+    if pip_install_is_set:
+        print("Install testing infra")
+        other_dependencies = ["pytest"]
+        assert 0 == run("pip install {} --progress-bar off -q".format(" ".join(other_dependencies)))
 
-    # install einops
-    assert 0 == run("pip install -e .")
+        for framework in frameworks:
+            print(f"Installing {framework}")
+            pip_instructions = framework_name2installation[framework]
+            assert 0 == run("pip install {} --progress-bar off -q".format(" ".join(pip_instructions)))
 
     # we need to inform testing script which frameworks to use
-    # this is done by setting a flag EINOPS_TEST_BACKENDS
-    from tests import unparse_backends
+    # this is done by setting an envvar EINOPS_TEST_BACKENDS
+    from einops.tests import unparse_backends
 
-    flag_name, flag_value = unparse_backends(backend_names=frameworks)
+    envvar_name, envvar_value = unparse_backends(backend_names=frameworks)
     return_code = run(
-        "python -m pytest tests",
-        **{flag_name: flag_value},
+        "python -m pytest .",
+        **{envvar_name: envvar_value},
     )
     assert return_code == 0
 

tests/test_einsum.py → einops/tests/test_einsum.py

@@ -1,5 +1,5 @@
 from typing import Any, Callable
-from . import collect_test_backends
+from einops.tests import collect_test_backends
 from einops.einops import _compactify_pattern_for_einsum, einsum, EinopsError
 import numpy as np
 import pytest

tests/test_examples.py → einops/tests/test_examples.py

@@ -2,8 +2,8 @@
 import pytest
 
 from einops import rearrange, parse_shape, reduce
-from tests import is_backend_tested
-from tests.test_ops import imp_op_backends
+from einops.tests import is_backend_tested
+from einops.tests.test_ops import imp_op_backends
 
 
 def test_rearrange_examples():

tests/test_layers.py → einops/tests/test_layers.py

@@ -5,7 +5,7 @@
 import pytest
 
 from einops import rearrange, reduce
-from . import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS
+from einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS
 
 __author__ = "Alex Rogozhnikov"
 

tests/test_ops.py → einops/tests/test_ops.py

@@ -6,7 +6,7 @@
 
 from einops import EinopsError
 from einops.einops import rearrange, reduce, repeat, _enumerate_directions
-from . import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS
+from einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS
 
 imp_op_backends = collect_test_backends(symbolic=False, layers=False)
 sym_op_backends = collect_test_backends(symbolic=True, layers=False)