From b626f63f124d8d6d447592e9cf4b6fb7d941fcdb Mon Sep 17 00:00:00 2001 From: David Holtz Date: Fri, 21 Aug 2026 14:25:41 +0000 Subject: [PATCH 01/15] feat: prefer recipe ported einops Signed-off-by: David Holtz --- .gitignore | 4 + .ports/einops/overlay/flake.lock | 117 ++++ .ports/einops/overlay/flake.nix | 17 + .ports/einops/overlay/tests/__init__.py | 119 ++++ .ports/einops/overlay/tests/run_tests.py | 85 +++ .ports/einops/overlay/tests/test_einsum.py | 360 ++++++++++ .ports/einops/overlay/tests/test_examples.py | 303 ++++++++ .ports/einops/overlay/tests/test_layers.py | 484 +++++++++++++ .ports/einops/overlay/tests/test_ops.py | 684 +++++++++++++++++++ .ports/einops/overlay/tests/test_other.py | 307 +++++++++ .ports/einops/overlay/tests/test_packing.py | 318 +++++++++ .ports/einops/overlay/tests/test_parsing.py | 134 ++++ .ports/einops/port.kdl | 60 ++ einops/.port-provenance.json | 12 + einops/build.toml | 10 +- 15 files changed, 3012 insertions(+), 2 deletions(-) create mode 100644 .ports/einops/overlay/flake.lock create mode 100644 .ports/einops/overlay/flake.nix create mode 100644 .ports/einops/overlay/tests/__init__.py create mode 100644 .ports/einops/overlay/tests/run_tests.py create mode 100644 .ports/einops/overlay/tests/test_einsum.py create mode 100644 .ports/einops/overlay/tests/test_examples.py create mode 100644 .ports/einops/overlay/tests/test_layers.py create mode 100644 .ports/einops/overlay/tests/test_ops.py create mode 100644 .ports/einops/overlay/tests/test_other.py create mode 100644 .ports/einops/overlay/tests/test_packing.py create mode 100644 .ports/einops/overlay/tests/test_parsing.py create mode 100644 .ports/einops/port.kdl create mode 100644 einops/.port-provenance.json diff --git a/.gitignore b/.gitignore index 4dcaf8670..9a41b7635 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,10 @@ build .* !.github/ +# Port files and provenance +!.ports/ +!.port-provenance.json + # Python stuff __pycache__/ diff --git a/.ports/einops/overlay/flake.lock b/.ports/einops/overlay/flake.lock new file mode 100644 index 000000000..1446e6473 --- /dev/null +++ b/.ports/einops/overlay/flake.lock @@ -0,0 +1,117 @@ +{ + "nodes": { + "flake-compat": { + "locked": { + "lastModified": 1767039857, + "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=", + "owner": "edolstra", + "repo": "flake-compat", + "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab", + "type": "github" + }, + "original": { + "owner": "edolstra", + "repo": "flake-compat", + "type": "github" + } + }, + "flake-utils": { + "inputs": { + "systems": "systems" + }, + "locked": { + "lastModified": 1731533236, + "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", + "owner": "numtide", + "repo": "flake-utils", + "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", + "type": "github" + }, + "original": { + "owner": "numtide", + "repo": "flake-utils", + "type": "github" + } + }, + "kernel-builder": { + "inputs": { + "flake-compat": "flake-compat", + "flake-utils": "flake-utils", + "nixpkgs": "nixpkgs", + "rust-overlay": "rust-overlay" + }, + "locked": { + "lastModified": 1785427612, + "narHash": "sha256-hUFEwQKDjVHw9v6xgC6Lx402TYfUA8zXEi8hDE/F4y0=", + "owner": "huggingface", + "repo": "kernels", + "rev": "d0610aa58db33b142c86b59598a2a1c730f52996", + "type": "github" + }, + "original": { + "owner": "huggingface", + "repo": "kernels", + "type": "github" + } + }, + "nixpkgs": { + "locked": { + "lastModified": 1783284758, + "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=", + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", + "type": "github" + }, + "original": { + "owner": "NixOS", + "repo": "nixpkgs", + "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", + "type": "github" + } + }, + "root": { + "inputs": { + "kernel-builder": "kernel-builder" + } + }, + "rust-overlay": { + "inputs": { + "nixpkgs": [ + "kernel-builder", + "nixpkgs" + ] + }, + "locked": { + "lastModified": 1783320166, + "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=", + "owner": "oxalica", + "repo": "rust-overlay", + "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673", + "type": "github" + }, + "original": { + "owner": "oxalica", + "repo": "rust-overlay", + "type": "github" + } + }, + "systems": { + "locked": { + "lastModified": 1681028828, + "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", + "owner": "nix-systems", + "repo": "default", + "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", + "type": "github" + }, + "original": { + "owner": "nix-systems", + "repo": "default", + "type": "github" + } + } + }, + "root": "root", + "version": 7 +} diff --git a/.ports/einops/overlay/flake.nix b/.ports/einops/overlay/flake.nix new file mode 100644 index 000000000..0b2a92f1e --- /dev/null +++ b/.ports/einops/overlay/flake.nix @@ -0,0 +1,17 @@ +{ + description = "Flake for eigen kernel"; + + inputs = { + kernel-builder.url = "github:huggingface/kernels"; + }; + + outputs = + { + self, + kernel-builder, + }: + kernel-builder.lib.genKernelFlakeOutputs { + inherit self; + path = ./.; + }; +} diff --git a/.ports/einops/overlay/tests/__init__.py b/.ports/einops/overlay/tests/__init__.py new file mode 100644 index 000000000..b8c629217 --- /dev/null +++ b/.ports/einops/overlay/tests/__init__.py @@ -0,0 +1,119 @@ +""" +Common utils for testing. +These functions allow testing only some frameworks, not all. +""" + +import logging +import os +from functools import lru_cache +from typing import List, Tuple + +import kernels +_einops = kernels.get_kernel("kernels-community/einops", version=1) +_backends = _einops._backends +import warnings + +__author__ = "Alex Rogozhnikov" + + +# minimize noise in tests logging +logging.getLogger("tensorflow").disabled = True +logging.getLogger("matplotlib").disabled = True + +FLOAT_REDUCTIONS = ("min", "max", "sum", "mean", "prod") # not includes any/all + + +def find_names_of_all_frameworks() -> List[str]: + backend_subclasses = [] + backends = _backends.AbstractBackend.__subclasses__() + while backends: + backend = backends.pop() + backends += backend.__subclasses__() + backend_subclasses.append(backend) + return [b.framework_name for b in backend_subclasses] + + +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 ENVVAR_NAME in os.environ: + 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: + raise RuntimeError(f"Unknown framework: {backend_name}") + return parsed_backends + + # Auto-detect available frameworks. + available = [] + for module, name in [("torch", "torch"), ("jax", "jax"), ("numpy", "numpy"), ("tensorflow", "tensorflow")]: + try: + __import__(module) + available.append(name) + except ImportError: + pass + return available + + +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 collect_test_backends(symbolic=False, layers=False) -> List[_backends.AbstractBackend]: + """ + :param symbolic: symbolic or imperative frameworks? + :param layers: layers or operations? + :return: list of backends satisfying set conditions + """ + if not symbolic: + if not layers: + backend_types = [ + _backends.NumpyBackend, + _backends.JaxBackend, + _backends.TorchBackend, + _backends.TensorflowBackend, + _backends.OneFlowBackend, + _backends.PaddleBackend, + _backends.CupyBackend, + ] + else: + backend_types = [ + _backends.TorchBackend, + _backends.OneFlowBackend, + _backends.PaddleBackend, + ] + else: + if not layers: + backend_types = [ + _backends.PyTensorBackend, + ] + else: + backend_types = [ + _backends.TFKerasBackend, + ] + + backend_names_to_test = parse_backends_to_test() + result = [] + for backend_type in backend_types: + if backend_type.framework_name not in backend_names_to_test: + continue + try: + result.append(backend_type()) + except ImportError: + # problem with backend installation fails a specific test function, + # but will be skipped in all other test cases + warnings.warn("backend could not be initialized for tests: {}".format(backend_type)) + return result diff --git a/.ports/einops/overlay/tests/run_tests.py b/.ports/einops/overlay/tests/run_tests.py new file mode 100644 index 000000000..cf7a4eb6e --- /dev/null +++ b/.ports/einops/overlay/tests/run_tests.py @@ -0,0 +1,85 @@ +""" +Runs tests that are appropriate for framework. +""" + +import os +import sys +from subprocess import Popen +from pathlib import Path + +__author__ = "Alex Rogozhnikov" + + +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 + + +def main(): + _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]", "flax"], + "tensorflow": ["tensorflow"], + "cupy": ["cupy"], + # switch to stable paddlepaddle, because of https://github.com/PaddlePaddle/Paddle/issues/63927 + # "paddle": ["paddlepaddle==0.0.0 -f https://www.paddlepaddle.org.cn/whl/linux/cpu-mkl/develop.html"], + "paddle": ["paddlepaddle"], + "oneflow": ["oneflow==0.9.0"], + "pytensor": ["pytensor"], + } + + usage = f""" + Usage: python -m einops.tests.run_tests [--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) + return + else: + synonyms = { + "tf": "tensorflow", + "pytorch": "torch", + "paddlepaddle": "paddle", + } + frameworks = [synonyms.get(f, f) for f in frameworks] + wrong_frameworks = [f for f in frameworks if f not in framework_name2installation] + if wrong_frameworks: + print(usage) + raise RuntimeError(f"Unrecognized frameworks: {wrong_frameworks}") + + 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))) + + 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 an envvar EINOPS_TEST_BACKENDS + from einops.tests import unparse_backends + + envvar_name, envvar_value = unparse_backends(backend_names=frameworks) + return_code = run( + "python -m pytest .", + **{envvar_name: envvar_value}, + ) + assert return_code == 0 + + +if __name__ == "__main__": + main() diff --git a/.ports/einops/overlay/tests/test_einsum.py b/.ports/einops/overlay/tests/test_einsum.py new file mode 100644 index 000000000..2e05d939d --- /dev/null +++ b/.ports/einops/overlay/tests/test_einsum.py @@ -0,0 +1,360 @@ +from typing import Any, Callable +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +_compactify_pattern_for_einsum = einops.einops._compactify_pattern_for_einsum +einsum = einops.einops.einsum +EinopsError = einops.EinopsError +from . import collect_test_backends +import numpy as np +import pytest +import string + + +class Arguments: + def __init__(self, *args: Any, **kargs: Any): + self.args = args + self.kwargs = kargs + + def __call__(self, function: Callable): + return function(*self.args, **self.kwargs) + + +test_layer_cases = [ + ( + Arguments("b c_in h w -> w c_out h b", "c_in c_out", bias_shape=None, c_out=13, c_in=12), + (2, 12, 3, 4), + (4, 13, 3, 2), + ), + ( + Arguments("b c_in h w -> w c_out h b", "c_in c_out", bias_shape="c_out", c_out=13, c_in=12), + (2, 12, 3, 4), + (4, 13, 3, 2), + ), + ( + Arguments("b c_in h w -> w c_in h b", "", bias_shape=None, c_in=12), + (2, 12, 3, 4), + (4, 12, 3, 2), + ), + ( + Arguments("b c_in h w -> b c_out", "c_in h w c_out", bias_shape=None, c_in=12, h=3, w=4, c_out=5), + (2, 12, 3, 4), + (2, 5), + ), + ( + Arguments("b t head c_in -> b t head c_out", "head c_in c_out", bias_shape=None, head=4, c_in=5, c_out=6), + (2, 3, 4, 5), + (2, 3, 4, 6), + ), +] + + +# Each of the form: +# (Arguments, true_einsum_pattern, in_shapes, out_shape) +test_functional_cases = [ + ( + # Basic: + "b c h w, b w -> b h", + "abcd,ad->ac", + ((2, 3, 4, 5), (2, 5)), + (2, 4), + ), + ( + # Three tensors: + "b c h w, b w, b c -> b h", + "abcd,ad,ab->ac", + ((2, 3, 40, 5), (2, 5), (2, 3)), + (2, 40), + ), + ( + # Ellipsis, and full names: + "... one two three, three four five -> ... two five", + "...abc,cde->...be", + ((32, 5, 2, 3, 4), (4, 5, 6)), + (32, 5, 3, 6), + ), + ( + # Ellipsis at the end: + "one two three ..., three four five -> two five ...", + "abc...,cde->be...", + ((2, 3, 4, 32, 5), (4, 5, 6)), + (3, 6, 32, 5), + ), + ( + # Ellipsis on multiple tensors: + "... one two three, ... three four five -> ... two five", + "...abc,...cde->...be", + ((32, 5, 2, 3, 4), (32, 5, 4, 5, 6)), + (32, 5, 3, 6), + ), + ( + # One tensor, and underscores: + "first_tensor second_tensor -> first_tensor", + "ab->a", + ((5, 4),), + (5,), + ), + ( + # Trace (repeated index) + "i i -> ", + "aa->", + ((5, 5),), + (), + ), + ( + # Too many spaces in string: + " one two , three four->two four ", + "ab,cd->bd", + ((2, 3), (4, 5)), + (3, 5), + ), + # The following tests were inspired by numpy's einsum tests + # https://github.com/numpy/numpy/blob/v1.23.0/numpy/core/tests/test_einsum.py + ( + # Trace with other indices + "i middle i -> middle", + "aba->b", + ((5, 10, 5),), + (10,), + ), + ( + # Ellipsis in the middle: + "i ... i -> ...", + "a...a->...", + ((5, 3, 2, 1, 4, 5),), + (3, 2, 1, 4), + ), + ( + # Product of first and last axes: + "i ... i -> i ...", + "a...a->a...", + ((5, 3, 2, 1, 4, 5),), + (5, 3, 2, 1, 4), + ), + ( + # Triple diagonal + "one one one -> one", + "aaa->a", + ((5, 5, 5),), + (5,), + ), + ( + # Axis swap: + "i j k -> j i k", + "abc->bac", + ((1, 2, 3),), + (2, 1, 3), + ), + ( + # Identity: + "... -> ...", + "...->...", + ((5, 4, 3, 2, 1),), + (5, 4, 3, 2, 1), + ), + ( + # Elementwise product of three tensors + "..., ..., ... -> ...", + "...,...,...->...", + ((3, 2), (3, 2), (3, 2)), + (3, 2), + ), + ( + # Basic summation: + "index ->", + "a->", + ((10,)), + (()), + ), +] + + +@pytest.mark.kernels_ci +def test_layer(): + for backend in collect_test_backends(layers=True, symbolic=False): + if backend.framework_name in ["tensorflow", "torch", "oneflow", "paddle"]: + layer_type = backend.layers().EinMix + for args, in_shape, out_shape in test_layer_cases: + layer = args(layer_type) + print("Running", layer.einsum_pattern, "for", backend.framework_name) + input = np.random.uniform(size=in_shape).astype("float32") + input_framework = backend.from_numpy(input) + output_framework = layer(input_framework) + output = backend.to_numpy(output_framework) + assert output.shape == out_shape + + +valid_backends_functional = [ + "tensorflow", + "torch", + "jax", + "numpy", + "oneflow", + "cupy", + "tensorflow.keras", + "paddle", + "pytensor", +] + + +@pytest.mark.kernels_ci +def test_functional(): + # Functional tests: + backends = filter(lambda x: x.framework_name in valid_backends_functional, collect_test_backends()) + for backend in backends: + for einops_pattern, true_pattern, in_shapes, out_shape in test_functional_cases: + print(f"Running '{einops_pattern}' for {backend.framework_name}") + + # Create pattern: + predicted_pattern = _compactify_pattern_for_einsum(einops_pattern) + assert predicted_pattern == true_pattern + + # Generate example data: + rstate = np.random.RandomState(0) + in_arrays = [rstate.uniform(size=shape).astype("float32") for shape in in_shapes] + in_arrays_framework = [backend.from_numpy(array) for array in in_arrays] + + # Loop over whether we call it manually with the backend, + # or whether we use `einops.einsum`. + for do_manual_call in [True, False]: + # Actually run einsum: + if do_manual_call: + out_array = backend.einsum(predicted_pattern, *in_arrays_framework) + else: + out_array = einsum(*in_arrays_framework, einops_pattern) + + # Check shape: + if tuple(out_array.shape) != out_shape: + raise ValueError(f"Expected output shape {out_shape} but got {out_array.shape}") + + # Check values: + true_out_array = np.einsum(true_pattern, *in_arrays) + predicted_out_array = backend.to_numpy(out_array) + np.testing.assert_array_almost_equal(predicted_out_array, true_out_array, decimal=5) + + +@pytest.mark.kernels_ci +def test_functional_symbolic(): + backends = filter( + lambda x: x.framework_name in valid_backends_functional, collect_test_backends(symbolic=True, layers=False) + ) + for backend in backends: + for einops_pattern, true_pattern, in_shapes, out_shape in test_functional_cases: + print(f"Running '{einops_pattern}' for symbolic {backend.framework_name}") + # Create pattern: + predicted_pattern = _compactify_pattern_for_einsum(einops_pattern) + assert predicted_pattern == true_pattern + + rstate = np.random.RandomState(0) + in_syms = [backend.create_symbol(in_shape) for in_shape in in_shapes] + in_data = [rstate.uniform(size=in_shape).astype("float32") for in_shape in in_shapes] + + expected_out_data = np.einsum(true_pattern, *in_data) + + for do_manual_call in [True, False]: + if do_manual_call: + predicted_out_symbol = backend.einsum(predicted_pattern, *in_syms) + else: + predicted_out_symbol = einsum(*in_syms, einops_pattern) + + predicted_out_data = backend.eval_symbol( + predicted_out_symbol, + list(zip(in_syms, in_data)), + ) + if predicted_out_data.shape != out_shape: + raise ValueError(f"Expected output shape {out_shape} but got {predicted_out_data.shape}") + np.testing.assert_array_almost_equal(predicted_out_data, expected_out_data, decimal=5) + + +@pytest.mark.kernels_ci +def test_functional_errors(): + # Specific backend does not matter, as errors are raised + # during the pattern creation. + + rstate = np.random.RandomState(0) + + def create_tensor(*shape): + return rstate.uniform(size=shape).astype("float32") + + # raise NotImplementedError("Singleton () axes are not yet supported in einsum.") + with pytest.raises(NotImplementedError, match="^Singleton"): + einsum( + create_tensor(5, 1), + "i () -> i", + ) + + # raise NotImplementedError("Shape rearrangement is not yet supported in einsum.") + with pytest.raises(NotImplementedError, match="^Shape rearrangement"): + einsum( + create_tensor(5, 1), + "a b -> (a b)", + ) + + with pytest.raises(NotImplementedError, match="^Shape rearrangement"): + einsum( + create_tensor(10, 1), + "(a b) -> a b", + ) + + # raise RuntimeError("Encountered empty axis name in einsum.") + # raise RuntimeError("Axis name in einsum must be a string.") + # ^ Not tested, these are just a failsafe in case an unexpected error occurs. + + # raise NotImplementedError("Anonymous axes are not yet supported in einsum.") + with pytest.raises(NotImplementedError, match="^Anonymous axes"): + einsum( + create_tensor(5, 1), + "i 2 -> i", + ) + + # ParsedExpression error: + with pytest.raises(EinopsError, match="^Invalid axis identifier"): + einsum( + create_tensor(5, 1), + "i 2j -> i", + ) + + # raise ValueError("Einsum pattern must contain '->'.") + with pytest.raises(ValueError, match="^Einsum pattern"): + einsum( + create_tensor(5, 3, 2), + "i j k", + ) + + # raise RuntimeError("Too many axes in einsum.") + with pytest.raises(RuntimeError, match="^Too many axes"): + einsum( + create_tensor(1), + " ".join(string.ascii_letters) + " extra ->", + ) + + # raise RuntimeError("Unknown axis on right side of einsum.") + with pytest.raises(RuntimeError, match="^Unknown axis"): + einsum( + create_tensor(5, 1), + "i j -> k", + ) + + # raise ValueError( + # "The last argument passed to `einops.einsum` must be a string," + # " representing the einsum pattern." + # ) + with pytest.raises(ValueError, match="^The last argument"): + einsum( + "i j k -> i", + create_tensor(5, 4, 3), + ) + + # raise ValueError( + # "`einops.einsum` takes at minimum two arguments: the tensors," + # " followed by the pattern." + # ) + with pytest.raises(ValueError, match="^`einops.einsum` takes"): + einsum( + "i j k -> i", + ) + with pytest.raises(ValueError, match="^`einops.einsum` takes"): + einsum( + create_tensor(5, 1), + ) + + # TODO: Include check for giving normal einsum pattern rather than einops. diff --git a/.ports/einops/overlay/tests/test_examples.py b/.ports/einops/overlay/tests/test_examples.py new file mode 100644 index 000000000..4a66431bc --- /dev/null +++ b/.ports/einops/overlay/tests/test_examples.py @@ -0,0 +1,303 @@ +import numpy +import pytest + +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +rearrange = einops.rearrange +parse_shape = einops.parse_shape +reduce = einops.reduce +from . import is_backend_tested +from .test_ops import imp_op_backends + + +@pytest.mark.kernels_ci +def test_rearrange_examples(): + def test1(x): + # transpose + y = rearrange(x, "b c h w -> b h w c") + assert tuple(y.shape) == (10, 30, 40, 20) + return y + + def test2(x): + # view / reshape + y = rearrange(x, "b c h w -> b (c h w)") + assert tuple(y.shape) == (10, 20 * 30 * 40) + return y + + def test3(x): + # depth-to-space + y = rearrange(x, "b (c h1 w1) h w -> b c (h h1) (w w1)", h1=2, w1=2) + assert tuple(y.shape) == (10, 5, 30 * 2, 40 * 2) + return y + + def test4(x): + # space-to-depth + y = rearrange(x, "b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=2, w1=2) + assert tuple(y.shape) == (10, 20 * 4, 30 // 2, 40 // 2) + return y + + def test5(x): + # simple transposition + y = rearrange(x, "b1 sound b2 letter -> b1 b2 sound letter") + assert tuple(y.shape) == (10, 30, 20, 40) + return y + + def test6(x): + # parsing parameters + t = rearrange(x, "b c h w -> (b h w) c") + t = t[:, ::2] # replacement for dot-product, just changes size of second axis + assert tuple(t.shape) == (10 * 30 * 40, 10) + + y = rearrange(t, "(b h w) c2 -> b c2 h w", **parse_shape(x, "b _ h w")) + assert tuple(y.shape) == (10, 10, 30, 40) + return y + + def test7(x): + # split of embedding into groups + y1, y2 = rearrange(x, "b (c g) h w -> g b c h w", g=2) + assert tuple(y1.shape) == (10, 10, 30, 40) + assert tuple(y2.shape) == (10, 10, 30, 40) + return y1 + y2 # only one tensor is expected in output + + def test8(x): + # max-pooling + y = reduce(x, "b c (h h1) (w w1) -> b c h w", reduction="max", h1=2, w1=2) + assert tuple(y.shape) == (10, 20, 30 // 2, 40 // 2) + return y + + def test9(x): + # squeeze - unsqueeze + y = reduce(x, "b c h w -> b c () ()", reduction="max") + assert tuple(y.shape) == (10, 20, 1, 1) + y = rearrange(y, "b c () () -> c b") + assert tuple(y.shape) == (20, 10) + return y + + def test10(x): + # stack + tensors = list(x + 0) # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185 + tensors = rearrange(tensors, "b c h w -> b h w c") + assert tuple(tensors.shape) == (10, 30, 40, 20) + return tensors + + def test11(x): + # concatenate + tensors = list(x + 0) # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185 + tensors = rearrange(tensors, "b c h w -> h (b w) c") + assert tuple(tensors.shape) == (30, 10 * 40, 20) + return tensors + + def shufflenet(x, convolve, c1, c2): + # shufflenet reordering example + x = convolve(x) + x = rearrange(x, "b (c1 c2) h w-> b (c2 c1) h w", c1=c1, c2=c2) + x = convolve(x) + return x + + def convolve_strided_1d(x, stride, usual_convolution): + x = rearrange(x, "b c t1 t2 -> b c (t1 t2)") # reduce dimensionality + x = rearrange(x, "b c (t stride) -> (stride b) c t", stride=stride) + x = usual_convolution(x) + x = rearrange(x, "(stride b) c t -> b c (t stride)", stride=stride) + return x + + def convolve_strided_2d(x, h_stride, w_stride, usual_convolution): + x = rearrange(x, "b c (h hs) (w ws) -> (hs ws b) c h w", hs=h_stride, ws=w_stride) + x = usual_convolution(x) + x = rearrange(x, "(hs ws b) c h w -> b c (h hs) (w ws)", hs=h_stride, ws=w_stride) + return x + + def unet_like_1d(x, usual_convolution): + # u-net like steps for increasing / reducing dimensionality + x = rearrange(x, "b c t1 t2 -> b c (t1 t2)") # reduce dimensionality + y = rearrange(x, "b c (t dt) -> b (dt c) t", dt=2) + y = usual_convolution(y) + x = x + rearrange(y, "b (dt c) t -> b c (t dt)", dt=2) + return x + + # mock for convolution (works for all backends) + def convolve_mock(x): + return x + + tests = [ + test1, + test2, + test3, + test4, + test5, + test6, + test7, + test8, + test9, + test10, + test11, + lambda x: shufflenet(x, convolve=convolve_mock, c1=4, c2=5), + lambda x: convolve_strided_1d(x, stride=2, usual_convolution=convolve_mock), + lambda x: convolve_strided_2d(x, h_stride=2, w_stride=2, usual_convolution=convolve_mock), + lambda x: unet_like_1d(x, usual_convolution=convolve_mock), + ] + + for backend in imp_op_backends: + print("testing source_examples for ", backend.framework_name) + for test in tests: + x = numpy.arange(10 * 20 * 30 * 40).reshape([10, 20, 30, 40]) + result1 = test(x) + result2 = backend.to_numpy(test(backend.from_numpy(x))) + assert numpy.array_equal(result1, result2) + + # now with strides + x = numpy.arange(10 * 2 * 20 * 3 * 30 * 1 * 40).reshape([10 * 2, 20 * 3, 30 * 1, 40 * 1]) + # known torch bug - torch doesn't support negative steps + last_step = -1 if (backend.framework_name != "torch" and backend.framework_name != "oneflow") else 1 + indexing_expression = numpy.index_exp[::2, ::3, ::1, ::last_step] + result1 = test(x[indexing_expression]) + result2 = backend.to_numpy(test(backend.from_numpy(x)[indexing_expression])) + assert numpy.array_equal(result1, result2) + + +def tensor_train_example_numpy(): + # kept here just for a collection, only tested for numpy + # https://arxiv.org/pdf/1509.06569.pdf, (5) + x = numpy.ones([3, 4, 5, 6]) + rank = 4 + if numpy.__version__ < "1.15.0": + # numpy.einsum fails here, skip test + return + # creating appropriate Gs + Gs = [numpy.ones([d, d, rank, rank]) for d in x.shape] + Gs[0] = Gs[0][:, :, :1, :] + Gs[-1] = Gs[-1][:, :, :, :1] + + # einsum way + y = x.reshape((1,) + x.shape) + for G in Gs: + # taking partial results left-to-right + # y = numpy.einsum('i j alpha beta, alpha i ... -> beta ... j', G, y) + y = numpy.einsum("i j a b, a i ... -> b ... j", G, y) + y1 = y.reshape(-1) + + # alternative way + y = x.reshape(-1) + for G in Gs: + i, j, alpha, beta = G.shape + y = rearrange(y, "(i rest alpha) -> rest (alpha i)", alpha=alpha, i=i) + y = y @ rearrange(G, "i j alpha beta -> (alpha i) (j beta)") + y = rearrange(y, "rest (beta j) -> (beta rest j)", beta=beta, j=j) + y2 = y + assert numpy.allclose(y1, y2) + + # yet another way + y = x + for G in Gs: + i, j, alpha, beta = G.shape + y = rearrange(y, "i ... (j alpha) -> ... j (alpha i)", alpha=alpha, i=i) + y = y @ rearrange(G, "i j alpha beta -> (alpha i) (j beta)") + y3 = y.reshape(-1) + assert numpy.allclose(y1, y3) + + +@pytest.mark.kernels_ci +def test_pytorch_yolo_fragment(): + if not is_backend_tested("torch"): + pytest.skip() + + import torch + + def old_way(input, num_classes, num_anchors, anchors, stride_h, stride_w): + # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L28-L44 + bs = input.size(0) + in_h = input.size(2) + in_w = input.size(3) + scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors] + + prediction = input.view(bs, num_anchors, 5 + num_classes, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous() + # Get outputs + x = torch.sigmoid(prediction[..., 0]) # Center x + y = torch.sigmoid(prediction[..., 1]) # Center y + w = prediction[..., 2] # Width + h = prediction[..., 3] # Height + conf = torch.sigmoid(prediction[..., 4]) # Conf + pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred. + + # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L70-L92 + FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor + LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor + # Calculate offsets for each grid + grid_x = ( + torch.linspace(0, in_w - 1, in_w) + .repeat(in_w, 1) + .repeat(bs * num_anchors, 1, 1) + .view(x.shape) + .type(FloatTensor) + ) + grid_y = ( + torch.linspace(0, in_h - 1, in_h) + .repeat(in_h, 1) + .t() + .repeat(bs * num_anchors, 1, 1) + .view(y.shape) + .type(FloatTensor) + ) + # Calculate anchor w, h + anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0])) + anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1])) + anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape) + anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape) + # Add offset and scale with anchors + pred_boxes = FloatTensor(prediction[..., :4].shape) + pred_boxes[..., 0] = x.data + grid_x + pred_boxes[..., 1] = y.data + grid_y + pred_boxes[..., 2] = torch.exp(w.data) * anchor_w + pred_boxes[..., 3] = torch.exp(h.data) * anchor_h + # Results + _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor) + output = torch.cat( + (pred_boxes.view(bs, -1, 4) * _scale, conf.view(bs, -1, 1), pred_cls.view(bs, -1, num_classes)), -1 + ) + return output + + def new_way(input, num_classes, num_anchors, anchors, stride_h, stride_w): + raw_predictions = rearrange(input, " b (anchor prediction) h w -> prediction b anchor h w", anchor=num_anchors) + + anchors = torch.FloatTensor(anchors).to(input.device) + anchor_sizes = rearrange(anchors, "anchor dim -> dim () anchor () ()") + + _, _, _, in_h, in_w = raw_predictions.shape + grid_h = rearrange(torch.arange(in_h).float(), "h -> () () h ()").to(input.device) + grid_w = rearrange(torch.arange(in_w).float(), "w -> () () () w").to(input.device) + + predicted_bboxes = torch.zeros_like(raw_predictions) + predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_h) * stride_h # center y + predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_w) * stride_w # center x + predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes # bbox width and height + predicted_bboxes[4] = raw_predictions[4].sigmoid() # confidence + predicted_bboxes[5:] = raw_predictions[5:].sigmoid() # class predictions + # only to match results of original code, not needed + return rearrange(predicted_bboxes, "prediction b anchor h w -> b anchor h w prediction") + + stride_h = 4 + stride_w = 4 + batch_size = 5 + num_classes = 12 + anchors = [[50, 100], [100, 50], [75, 75]] + num_anchors = len(anchors) + + input = torch.randn([batch_size, num_anchors * (5 + num_classes), 1, 1]) + result1 = old_way( + input=input, + num_anchors=num_anchors, + num_classes=num_classes, + stride_h=stride_h, + stride_w=stride_w, + anchors=anchors, + ) + result2 = new_way( + input=input, + num_anchors=num_anchors, + num_classes=num_classes, + stride_h=stride_h, + stride_w=stride_w, + anchors=anchors, + ) + result1 = result1.reshape(result2.shape) + assert torch.allclose(result1, result2) diff --git a/.ports/einops/overlay/tests/test_layers.py b/.ports/einops/overlay/tests/test_layers.py new file mode 100644 index 000000000..538b3b8e1 --- /dev/null +++ b/.ports/einops/overlay/tests/test_layers.py @@ -0,0 +1,484 @@ +import pickle +from collections import namedtuple + +import numpy +import pytest + +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +rearrange = einops.rearrange +reduce = einops.reduce +EinopsError = einops.EinopsError +from . import collect_test_backends, is_backend_tested +from . import FLOAT_REDUCTIONS as REDUCTIONS + +__author__ = "Alex Rogozhnikov" + +testcase = namedtuple("testcase", ["pattern", "axes_lengths", "input_shape", "wrong_shapes"]) + +rearrangement_patterns = [ + testcase( + "b c h w -> b (c h w)", + dict(c=20), + (10, 20, 30, 40), + [(), (10,), (10, 10, 10), (10, 21, 30, 40), [1, 20, 1, 1, 1]], + ), + testcase( + "b c (h1 h2) (w1 w2) -> b (c h2 w2) h1 w1", + dict(h2=2, w2=2), + (10, 20, 30, 40), + [(), (1, 1, 1, 1), (1, 10, 3), ()], + ), + testcase( + "b ... c -> c b ...", + dict(b=10), + (10, 20, 30), + [(), (10,), (5, 10)], + ), +] + + +@pytest.mark.kernels_ci +def test_rearrange_imperative(): + for backend in collect_test_backends(symbolic=False, layers=True): + print("Test layer for ", backend.framework_name) + + for pattern, axes_lengths, input_shape, wrong_shapes in rearrangement_patterns: + x = numpy.arange(numpy.prod(input_shape), dtype="float32").reshape(input_shape) + result_numpy = rearrange(x, pattern, **axes_lengths) + layer = backend.layers().Rearrange(pattern, **axes_lengths) + for shape in wrong_shapes: + try: + layer(backend.from_numpy(numpy.zeros(shape, dtype="float32"))) + except BaseException: + pass + else: + raise AssertionError("Failure expected") + + # simple pickling / unpickling + layer2 = pickle.loads(pickle.dumps(layer)) + result1 = backend.to_numpy(layer(backend.from_numpy(x))) + result2 = backend.to_numpy(layer2(backend.from_numpy(x))) + assert numpy.allclose(result_numpy, result1) + assert numpy.allclose(result1, result2) + + just_sum = backend.layers().Reduce("...->", reduction="sum") + + variable = backend.from_numpy(x) + result = just_sum(layer(variable)) + + result.backward() + assert numpy.allclose(backend.to_numpy(variable.grad), 1) + + +@pytest.mark.kernels_ci +def test_rearrange_symbolic(): + for backend in collect_test_backends(symbolic=True, layers=True): + print("Test layer for ", backend.framework_name) + + for pattern, axes_lengths, input_shape, wrong_shapes in rearrangement_patterns: + x = numpy.arange(numpy.prod(input_shape), dtype="float32").reshape(input_shape) + result_numpy = rearrange(x, pattern, **axes_lengths) + layer = backend.layers().Rearrange(pattern, **axes_lengths) + input_shape_of_nones = [None] * len(input_shape) + shapes = [input_shape, input_shape_of_nones] + + for shape in shapes: + symbol = backend.create_symbol(shape) + eval_inputs = [(symbol, x)] + + result_symbol1 = layer(symbol) + result1 = backend.eval_symbol(result_symbol1, eval_inputs) + assert numpy.allclose(result_numpy, result1) + + layer2 = pickle.loads(pickle.dumps(layer)) + result_symbol2 = layer2(symbol) + result2 = backend.eval_symbol(result_symbol2, eval_inputs) + assert numpy.allclose(result1, result2) + + # now testing back-propagation + just_sum = backend.layers().Reduce("...->", reduction="sum") + + result_sum1 = backend.eval_symbol(just_sum(result_symbol1), eval_inputs) + result_sum2 = numpy.sum(x) + + assert numpy.allclose(result_sum1, result_sum2) + + +reduction_patterns = rearrangement_patterns + [ + testcase("b c h w -> b ()", dict(b=10), (10, 20, 30, 40), [(10,), (10, 20, 30)]), + testcase("b c (h1 h2) (w1 w2) -> b c h1 w1", dict(h1=15, h2=2, w2=2), (10, 20, 30, 40), [(10, 20, 31, 40)]), + testcase("b ... c -> b", dict(b=10), (10, 20, 30, 40), [(10,), (11, 10)]), +] + + +@pytest.mark.kernels_ci +def test_reduce_imperative(): + for backend in collect_test_backends(symbolic=False, layers=True): + print("Test layer for ", backend.framework_name) + for reduction in REDUCTIONS: + for pattern, axes_lengths, input_shape, wrong_shapes in reduction_patterns: + print(backend, reduction, pattern, axes_lengths, input_shape, wrong_shapes) + x = numpy.arange(1, 1 + numpy.prod(input_shape), dtype="float32").reshape(input_shape) + x /= x.mean() + result_numpy = reduce(x, pattern, reduction, **axes_lengths) + layer = backend.layers().Reduce(pattern, reduction, **axes_lengths) + for shape in wrong_shapes: + try: + layer(backend.from_numpy(numpy.zeros(shape, dtype="float32"))) + except BaseException: + pass + else: + raise AssertionError("Failure expected") + + # simple pickling / unpickling + layer2 = pickle.loads(pickle.dumps(layer)) + result1 = backend.to_numpy(layer(backend.from_numpy(x))) + result2 = backend.to_numpy(layer2(backend.from_numpy(x))) + assert numpy.allclose(result_numpy, result1) + assert numpy.allclose(result1, result2) + + just_sum = backend.layers().Reduce("...->", reduction="sum") + + variable = backend.from_numpy(x) + result = just_sum(layer(variable)) + + result.backward() + grad = backend.to_numpy(variable.grad) + if reduction == "sum": + assert numpy.allclose(grad, 1) + if reduction == "mean": + assert numpy.allclose(grad, grad.min()) + if reduction in ["max", "min"]: + assert numpy.all(numpy.isin(grad, [0, 1])) + assert numpy.sum(grad) > 0.5 + + +@pytest.mark.kernels_ci +def test_reduce_symbolic(): + for backend in collect_test_backends(symbolic=True, layers=True): + print("Test layer for ", backend.framework_name) + for reduction in REDUCTIONS: + for pattern, axes_lengths, input_shape, wrong_shapes in reduction_patterns: + x = numpy.arange(1, 1 + numpy.prod(input_shape), dtype="float32").reshape(input_shape) + x /= x.mean() + result_numpy = reduce(x, pattern, reduction, **axes_lengths) + layer = backend.layers().Reduce(pattern, reduction, **axes_lengths) + input_shape_of_nones = [None] * len(input_shape) + shapes = [input_shape, input_shape_of_nones] + + for shape in shapes: + symbol = backend.create_symbol(shape) + eval_inputs = [(symbol, x)] + + result_symbol1 = layer(symbol) + result1 = backend.eval_symbol(result_symbol1, eval_inputs) + assert numpy.allclose(result_numpy, result1) + + layer2 = pickle.loads(pickle.dumps(layer)) + result_symbol2 = layer2(symbol) + result2 = backend.eval_symbol(result_symbol2, eval_inputs) + assert numpy.allclose(result1, result2) + + +def create_torch_model(use_reduce=False, add_scripted_layer=False): + if not is_backend_tested("torch"): + pytest.skip() + else: + from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU + Rearrange, Reduce, EinMix = einops.layers.torch.Rearrange, einops.layers.torch.Reduce, einops.layers.torch.EinMix + import torch.jit + + return Sequential( + Conv2d(3, 6, kernel_size=(5, 5)), + Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2) if use_reduce else MaxPool2d(kernel_size=2), + Conv2d(6, 16, kernel_size=(5, 5)), + Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2), + torch.jit.script(Rearrange("b c h w -> b (c h w)")) + if add_scripted_layer + else Rearrange("b c h w -> b (c h w)"), + Linear(16 * 5 * 5, 120), + ReLU(), + Linear(120, 84), + ReLU(), + EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84), + EinMix("(b c2) -> b c3", weight_shape="c2 c3", bias_shape="c3", c2=84, c3=84), + Linear(84, 10), + ) + + +@pytest.mark.kernels_ci +def test_torch_layer(): + if not is_backend_tested("torch"): + pytest.skip() + else: + # checked that torch present + import torch + import torch.jit + + model1 = create_torch_model(use_reduce=True) + model2 = create_torch_model(use_reduce=False) + input = torch.randn([10, 3, 32, 32]) + # random models have different predictions + assert not torch.allclose(model1(input), model2(input)) + model2.load_state_dict(pickle.loads(pickle.dumps(model1.state_dict()))) + assert torch.allclose(model1(input), model2(input)) + + # tracing (freezing) + model3 = torch.jit.trace(model2, example_inputs=input) + torch.testing.assert_close(model1(input), model3(input), atol=1e-3, rtol=1e-3) + torch.testing.assert_close(model1(input + 1), model3(input + 1), atol=1e-3, rtol=1e-3) + + model4 = torch.jit.trace(model2, example_inputs=input) + torch.testing.assert_close(model1(input), model4(input), atol=1e-3, rtol=1e-3) + torch.testing.assert_close(model1(input + 1), model4(input + 1), atol=1e-3, rtol=1e-3) + + +@pytest.mark.kernels_ci +def test_torch_layers_scripting(): + if not is_backend_tested("torch"): + pytest.skip() + else: + import torch + + for script_layer in [False, True]: + model1 = create_torch_model(use_reduce=True, add_scripted_layer=script_layer) + model2 = torch.jit.script(model1) + input = torch.randn([10, 3, 32, 32]) + + torch.testing.assert_close(model1(input), model2(input), atol=1e-3, rtol=1e-3) + + +@pytest.mark.kernels_ci +def test_keras_layer(): + if not is_backend_tested("tensorflow"): + pytest.skip() + else: + import tensorflow as tf + + if tf.__version__ < "2.16.": + # current implementation of layers follows new TF interface + pytest.skip() + from tensorflow.keras.models import Sequential + from tensorflow.keras.layers import Conv2D as Conv2d, Dense as Linear, ReLU + Rearrange, Reduce, EinMix, keras_custom_objects = einops.layers.keras.Rearrange, einops.layers.keras.Reduce, einops.layers.keras.EinMix, einops.layers.keras.keras_custom_objects + + def create_keras_model(): + return Sequential( + [ + Conv2d(6, kernel_size=5, input_shape=[32, 32, 3]), + Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2), + Conv2d(16, kernel_size=5), + Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2), + Rearrange("b c h w -> b (c h w)"), + Linear(120), + ReLU(), + Linear(84), + ReLU(), + EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84), + EinMix("(b c2) -> b c3", weight_shape="c2 c3", bias_shape="c3", c2=84, c3=84), + Linear(10), + ] + ) + + model1 = create_keras_model() + model2 = create_keras_model() + + input = numpy.random.normal(size=[10, 32, 32, 3]).astype("float32") + # two randomly init models should provide different outputs + assert not numpy.allclose(model1.predict_on_batch(input), model2.predict_on_batch(input)) + + # get some temp filename + tmp_model_filename = "/tmp/einops_tf_model.h5" + # save arch + weights + print("temp_path_keras1", tmp_model_filename) + tf.keras.models.save_model(model1, tmp_model_filename) + model3 = tf.keras.models.load_model(tmp_model_filename, custom_objects=keras_custom_objects) + + numpy.testing.assert_allclose(model1.predict_on_batch(input), model3.predict_on_batch(input)) + + weight_filename = "/tmp/einops_tf_model.weights.h5" + # save arch as json + model4 = tf.keras.models.model_from_json(model1.to_json(), custom_objects=keras_custom_objects) + model1.save_weights(weight_filename) + model4.load_weights(weight_filename) + model2.load_weights(weight_filename) + # check that differently-inialized model receives same weights + numpy.testing.assert_allclose(model1.predict_on_batch(input), model2.predict_on_batch(input)) + # ulimate test + # save-load architecture, and then load weights - should return same result + numpy.testing.assert_allclose(model1.predict_on_batch(input), model4.predict_on_batch(input)) + + +@pytest.mark.kernels_ci +def test_flax_layers(): + """ + One-off simple tests for Flax layers. + Unfortunately, Flax layers have a different interface from other layers. + """ + if not is_backend_tested("jax"): + pytest.skip() + else: + import jax + import jax.numpy as jnp + + import flax + from flax import linen as nn + EinMix, Reduce, Rearrange = einops.layers.flax.EinMix, einops.layers.flax.Reduce, einops.layers.flax.Rearrange + + class NN(nn.Module): + @nn.compact + def __call__(self, x): + x = EinMix( + "b (h h2) (w w2) c -> b h w c_out", "h2 w2 c c_out", "c_out", sizes=dict(h2=2, w2=3, c=4, c_out=5) + )(x) + x = Rearrange("b h w c -> b (w h c)", sizes=dict(c=5))(x) + x = Reduce("b hwc -> b", "mean", dict(hwc=2 * 3 * 5))(x) + return x + + model = NN() + fixed_input = jnp.ones([10, 2 * 2, 3 * 3, 4]) + params = model.init(jax.random.PRNGKey(0), fixed_input) + + def eval_at_point(params): + return jnp.linalg.norm(model.apply(params, fixed_input)) + + vandg = jax.value_and_grad(eval_at_point) + value0 = eval_at_point(params) + value1, grad1 = vandg(params) + assert jnp.allclose(value0, value1) + + params2 = jax.tree_map(lambda x1, x2: x1 - x2 * 0.001, params, grad1) + + value2 = eval_at_point(params2) + assert value0 >= value2, (value0, value2) + + # check serialization + fbytes = flax.serialization.to_bytes(params) + _loaded = flax.serialization.from_bytes(params, fbytes) + + +@pytest.mark.kernels_ci +def test_einmix_decomposition(): + """ + Testing that einmix correctly decomposes into smaller transformations. + """ + _EinmixDebugger = einops.layers._einmix._EinmixDebugger + + mixin1 = _EinmixDebugger( + "a b c d e -> e d c b a", + weight_shape="d a b", + d=2, a=3, b=5, + ) # fmt: off + assert mixin1.pre_reshape_pattern is None + assert mixin1.post_reshape_pattern is None + assert mixin1.einsum_pattern == "abcde,dab->edcba" + assert mixin1.saved_weight_shape == [2, 3, 5] + assert mixin1.saved_bias_shape is None + + mixin2 = _EinmixDebugger( + "a b c d e -> e d c b a", + weight_shape="d a b", + bias_shape="a b c d e", + a=1, b=2, c=3, d=4, e=5, + ) # fmt: off + assert mixin2.pre_reshape_pattern is None + assert mixin2.post_reshape_pattern is None + assert mixin2.einsum_pattern == "abcde,dab->edcba" + assert mixin2.saved_weight_shape == [4, 1, 2] + assert mixin2.saved_bias_shape == [5, 4, 3, 2, 1] + + mixin3 = _EinmixDebugger( + "... -> ...", + weight_shape="", + bias_shape="", + ) # fmt: off + assert mixin3.pre_reshape_pattern is None + assert mixin3.post_reshape_pattern is None + assert mixin3.einsum_pattern == "...,->..." + assert mixin3.saved_weight_shape == [] + assert mixin3.saved_bias_shape == [] + + mixin4 = _EinmixDebugger( + "b a ... -> b c ...", + weight_shape="b a c", + a=1, b=2, c=3, + ) # fmt: off + assert mixin4.pre_reshape_pattern is None + assert mixin4.post_reshape_pattern is None + assert mixin4.einsum_pattern == "ba...,bac->bc..." + assert mixin4.saved_weight_shape == [2, 1, 3] + assert mixin4.saved_bias_shape is None + + mixin5 = _EinmixDebugger( + "(b a) ... -> b c (...)", + weight_shape="b a c", + a=1, b=2, c=3, + ) # fmt: off + assert mixin5.pre_reshape_pattern == "(b a) ... -> b a ..." + assert mixin5.pre_reshape_lengths == dict(a=1, b=2) + assert mixin5.post_reshape_pattern == "b c ... -> b c (...)" + assert mixin5.einsum_pattern == "ba...,bac->bc..." + assert mixin5.saved_weight_shape == [2, 1, 3] + assert mixin5.saved_bias_shape is None + + mixin6 = _EinmixDebugger( + "b ... (a c) -> b ... (a d)", + weight_shape="c d", + bias_shape="a d", + a=1, c=3, d=4, + ) # fmt: off + assert mixin6.pre_reshape_pattern == "b ... (a c) -> b ... a c" + assert mixin6.pre_reshape_lengths == dict(a=1, c=3) + assert mixin6.post_reshape_pattern == "b ... a d -> b ... (a d)" + assert mixin6.einsum_pattern == "b...ac,cd->b...ad" + assert mixin6.saved_weight_shape == [3, 4] + assert mixin6.saved_bias_shape == [1, 1, 4] # (b) a d, ellipsis does not participate + + mixin7 = _EinmixDebugger( + "a ... (b c) -> a (... d b)", + weight_shape="c d b", + bias_shape="d b", + b=2, c=3, d=4, + ) # fmt: off + assert mixin7.pre_reshape_pattern == "a ... (b c) -> a ... b c" + assert mixin7.pre_reshape_lengths == dict(b=2, c=3) + assert mixin7.post_reshape_pattern == "a ... d b -> a (... d b)" + assert mixin7.einsum_pattern == "a...bc,cdb->a...db" + assert mixin7.saved_weight_shape == [3, 4, 2] + assert mixin7.saved_bias_shape == [1, 4, 2] # (a) d b, ellipsis does not participate + + +@pytest.mark.kernels_ci +def test_einmix_restrictions(): + """ + Testing different cases + """ + _EinmixDebugger = einops.layers._einmix._EinmixDebugger + + with pytest.raises(EinopsError): + _EinmixDebugger( + "a b c d e -> e d c b a", + weight_shape="d a b", + d=2, a=3, # missing b + ) # fmt: off + + with pytest.raises(EinopsError): + _EinmixDebugger( + "a b c d e -> e d c b a", + weight_shape="w a b", + d=2, a=3, b=1 # missing d + ) # fmt: off + + with pytest.raises(EinopsError): + _EinmixDebugger( + "(...) a -> ... a", + weight_shape="a", a=1, # ellipsis on the left + ) # fmt: off + + with pytest.raises(EinopsError): + _EinmixDebugger( + "(...) a -> a ...", + weight_shape="a", a=1, # ellipsis on the right side after bias axis + bias_shape='a', + ) # fmt: off diff --git a/.ports/einops/overlay/tests/test_ops.py b/.ports/einops/overlay/tests/test_ops.py new file mode 100644 index 000000000..ff668bb50 --- /dev/null +++ b/.ports/einops/overlay/tests/test_ops.py @@ -0,0 +1,684 @@ +import itertools + +import numpy +import numpy as np +import pytest + +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +rearrange = einops.einops.rearrange +reduce = einops.einops.reduce +repeat = einops.einops.repeat +_enumerate_directions = einops.einops._enumerate_directions +from . import collect_test_backends, is_backend_tested +from . import FLOAT_REDUCTIONS as REDUCTIONS + +imp_op_backends = collect_test_backends(symbolic=False, layers=False) +sym_op_backends = collect_test_backends(symbolic=True, layers=False) + +identity_patterns = [ + "...->...", + "a b c d e-> a b c d e", + "a b c d e ...-> ... a b c d e", + "a b c d e ...-> a ... b c d e", + "... a b c d e -> ... a b c d e", + "a ... e-> a ... e", + "a ... -> a ... ", + "a ... c d e -> a (...) c d e", +] + +equivalent_rearrange_patterns = [ + ("a b c d e -> (a b) c d e", "a b ... -> (a b) ... "), + ("a b c d e -> a b (c d) e", "... c d e -> ... (c d) e"), + ("a b c d e -> a b c d e", "... -> ... "), + ("a b c d e -> (a b c d e)", "... -> (...)"), + ("a b c d e -> b (c d e) a", "a b ... -> b (...) a"), + ("a b c d e -> b (a c d) e", "a b ... e -> b (a ...) e"), +] + +equivalent_reduction_patterns = [ + ("a b c d e -> ", " ... -> "), + ("a b c d e -> (e a)", "a ... e -> (e a)"), + ("a b c d e -> d (a e)", " a b c d e ... -> d (a e) "), + ("a b c d e -> (a b)", " ... c d e -> (...) "), +] + + +@pytest.mark.kernels_ci +def test_collapsed_ellipsis_errors_out(): + x = numpy.zeros([1, 1, 1, 1, 1]) + rearrange(x, "a b c d ... -> a b c ... d") + with pytest.raises(EinopsError): + rearrange(x, "a b c d (...) -> a b c ... d") + + rearrange(x, "... -> (...)") + with pytest.raises(EinopsError): + rearrange(x, "(...) -> (...)") + + +@pytest.mark.kernels_ci +def test_ellipsis_ops_numpy(): + x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) + for pattern in identity_patterns: + assert numpy.array_equal(x, rearrange(x, pattern)), pattern + + for pattern1, pattern2 in equivalent_rearrange_patterns: + assert numpy.array_equal(rearrange(x, pattern1), rearrange(x, pattern2)) + + for reduction in ["min", "max", "sum"]: + for pattern1, pattern2 in equivalent_reduction_patterns: + assert numpy.array_equal(reduce(x, pattern1, reduction=reduction), reduce(x, pattern2, reduction=reduction)) + + # now just check coincidence with numpy + all_rearrange_patterns = [*identity_patterns] + for pattern_pairs in equivalent_rearrange_patterns: + all_rearrange_patterns.extend(pattern_pairs) + + +def check_op_against_numpy(backend, numpy_input, pattern, axes_lengths, reduction="rearrange", is_symbolic=False): + """ + Helper to test result of operation (rearrange or transpose) against numpy + if reduction == 'rearrange', rearrange op is tested, otherwise reduce + """ + + def operation(x): + if reduction == "rearrange": + return rearrange(x, pattern, **axes_lengths) + else: + return reduce(x, pattern, reduction, **axes_lengths) + + numpy_result = operation(numpy_input) + check_equal = numpy.array_equal + p_none_dimension = 0.5 + if is_symbolic: + symbol_shape = [d if numpy.random.random() >= p_none_dimension else None for d in numpy_input.shape] + symbol = backend.create_symbol(shape=symbol_shape) + result_symbol = operation(symbol) + backend_result = backend.eval_symbol(result_symbol, [(symbol, numpy_input)]) + else: + backend_result = operation(backend.from_numpy(numpy_input)) + backend_result = backend.to_numpy(backend_result) + + check_equal(numpy_result, backend_result) + + +@pytest.mark.kernels_ci +def test_ellipsis_ops_imperative(): + """Checking various patterns against numpy""" + x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) + for is_symbolic in [True, False]: + for backend in collect_test_backends(symbolic=is_symbolic, layers=False): + for pattern in identity_patterns + list(itertools.chain(*equivalent_rearrange_patterns)): + check_op_against_numpy( + backend, x, pattern, axes_lengths={}, reduction="rearrange", is_symbolic=is_symbolic + ) + + for reduction in ["min", "max", "sum"]: + for pattern in itertools.chain(*equivalent_reduction_patterns): + check_op_against_numpy( + backend, x, pattern, axes_lengths={}, reduction=reduction, is_symbolic=is_symbolic + ) + + +@pytest.mark.kernels_ci +def test_rearrange_array_api(): + import numpy as xp + AA = einops.array_api + + if xp.__version__ < "2.0.0": + pytest.skip() + + x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) + for pattern in identity_patterns + list(itertools.chain(*equivalent_rearrange_patterns)): + expected = rearrange(x, pattern) + result = AA.rearrange(xp.from_dlpack(x), pattern) + assert numpy.array_equal(AA.asnumpy(result + 0), expected) + + +@pytest.mark.kernels_ci +def test_reduce_array_api(): + import numpy as xp + AA = einops.array_api + + if xp.__version__ < "2.0.0": + pytest.skip() + + x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) + for pattern in itertools.chain(*equivalent_reduction_patterns): + for reduction in ["min", "max", "sum"]: + expected = reduce(x, pattern, reduction=reduction) + result = AA.reduce(xp.from_dlpack(x), pattern, reduction=reduction) + assert numpy.array_equal(AA.asnumpy(np.asarray(result + 0)), expected) + + +@pytest.mark.kernels_ci +def test_rearrange_consistency_numpy(): + shape = [1, 2, 3, 5, 7, 11] + x = numpy.arange(numpy.prod(shape)).reshape(shape) + for pattern in [ + "a b c d e f -> a b c d e f", + "b a c d e f -> a b d e f c", + "a b c d e f -> f e d c b a", + "a b c d e f -> (f e) d (c b a)", + "a b c d e f -> (f e d c b a)", + ]: + result = rearrange(x, pattern) + assert len(numpy.setdiff1d(x, result)) == 0 + assert result.dtype == x.dtype + + result = rearrange(x, "a b c d e f -> a (b) (c d e) f") + assert numpy.array_equal(x.flatten(), result.flatten()) + + result = rearrange(x, "a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11") + assert numpy.array_equal(x, result) + + result1 = rearrange(x, "a b c d e f -> f e d c b a") + result2 = rearrange(x, "f e d c b a -> a b c d e f") + assert numpy.array_equal(result1, result2) + + result = rearrange(rearrange(x, "a b c d e f -> (f d) c (e b) a"), "(f d) c (e b) a -> a b c d e f", b=2, d=5) + assert numpy.array_equal(x, result) + + sizes = dict(zip("abcdef", shape)) + temp = rearrange(x, "a b c d e f -> (f d) c (e b) a", **sizes) + result = rearrange(temp, "(f d) c (e b) a -> a b c d e f", **sizes) + assert numpy.array_equal(x, result) + + x2 = numpy.arange(2 * 3 * 4).reshape([2, 3, 4]) + result = rearrange(x2, "a b c -> b c a") + assert x2[1, 2, 3] == result[2, 3, 1] + assert x2[0, 1, 2] == result[1, 2, 0] + + +@pytest.mark.kernels_ci +def test_rearrange_permutations_numpy(): + # tests random permutation of axes against two independent numpy ways + for n_axes in range(1, 10): + input = numpy.arange(2**n_axes).reshape([2] * n_axes) + permutation = numpy.random.permutation(n_axes) + left_expression = " ".join("i" + str(axis) for axis in range(n_axes)) + right_expression = " ".join("i" + str(axis) for axis in permutation) + expression = left_expression + " -> " + right_expression + result = rearrange(input, expression) + + for pick in numpy.random.randint(0, 2, [10, n_axes]): + assert input[tuple(pick)] == result[tuple(pick[permutation])] + + for n_axes in range(1, 10): + input = numpy.arange(2**n_axes).reshape([2] * n_axes) + permutation = numpy.random.permutation(n_axes) + left_expression = " ".join("i" + str(axis) for axis in range(n_axes)[::-1]) + right_expression = " ".join("i" + str(axis) for axis in permutation[::-1]) + expression = left_expression + " -> " + right_expression + result = rearrange(input, expression) + assert result.shape == input.shape + expected_result = numpy.zeros_like(input) + for original_axis, result_axis in enumerate(permutation): + expected_result |= ((input >> original_axis) & 1) << result_axis + + assert numpy.array_equal(result, expected_result) + + +@pytest.mark.kernels_ci +def test_reduction_imperatives(): + for backend in imp_op_backends: + print("Reduction tests for ", backend.framework_name) + for reduction in REDUCTIONS: + # slight redundancy for simpler order - numpy version is evaluated multiple times + input = numpy.arange(2 * 3 * 4 * 5 * 6, dtype="int64").reshape([2, 3, 4, 5, 6]) + if reduction in ["mean", "prod"]: + input = input / input.astype("float64").mean() + test_cases = [ + ["a b c d e -> ", {}, getattr(input, reduction)()], + ["a ... -> ", {}, getattr(input, reduction)()], + ["(a1 a2) ... (e1 e2) -> ", dict(a1=1, e2=2), getattr(input, reduction)()], + [ + "a b c d e -> (e c) a", + {}, + getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), + ], + [ + "a ... c d e -> (e c) a", + {}, + getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), + ], + [ + "a b c d e ... -> (e c) a", + {}, + getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), + ], + ["a b c d e -> (e c a)", {}, getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1])], + ["(a a2) ... -> (a2 a) ...", dict(a2=1), input], + ] + for pattern, axes_lengths, expected_result in test_cases: + result = reduce(backend.from_numpy(input.copy()), pattern, reduction=reduction, **axes_lengths) + result = backend.to_numpy(result) + assert numpy.allclose(result, expected_result), f"Failed at {pattern}" + + +@pytest.mark.kernels_ci +def test_reduction_symbolic(): + for backend in sym_op_backends: + print("Reduction tests for ", backend.framework_name) + for reduction in REDUCTIONS: + input = numpy.arange(2 * 3 * 4 * 5 * 6, dtype="int64").reshape([2, 3, 4, 5, 6]) + input = input / input.astype("float64").mean() + # slight redundancy for simpler order - numpy version is evaluated multiple times + test_cases = [ + ["a b c d e -> ", {}, getattr(input, reduction)()], + ["a ... -> ", {}, getattr(input, reduction)()], + ["(a a2) ... (e e2) -> ", dict(a2=1, e2=1), getattr(input, reduction)()], + [ + "a b c d e -> (e c) a", + {}, + getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), + ], + [ + "a ... c d e -> (e c) a", + {}, + getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), + ], + [ + "a b c d e ... -> (e c) a", + {}, + getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), + ], + ["a b c d e -> (e c a)", {}, getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1])], + ["(a a2) ... -> (a2 a) ...", dict(a2=1), input], + ] + for pattern, axes_lengths, expected_numpy_result in test_cases: + shapes = [input.shape, [None for _ in input.shape]] + for shape in shapes: + sym = backend.create_symbol(shape) + result_sym = reduce(sym, pattern, reduction=reduction, **axes_lengths) + result = backend.eval_symbol(result_sym, [(sym, input)]) + assert numpy.allclose(result, expected_numpy_result) + + if True: + shape = [] + _axes_lengths = {**axes_lengths} + for axis, length in zip("abcde", input.shape): + # filling as much as possible with Nones + if axis in pattern: + shape.append(None) + _axes_lengths[axis] = length + else: + shape.append(length) + sym = backend.create_symbol(shape) + result_sym = reduce(sym, pattern, reduction=reduction, **_axes_lengths) + result = backend.eval_symbol(result_sym, [(sym, input)]) + assert numpy.allclose(result, expected_numpy_result) + + +@pytest.mark.kernels_ci +def test_reduction_stress_imperatives(): + for backend in imp_op_backends: + print("Stress-testing reduction for ", backend.framework_name) + for reduction in REDUCTIONS + ("rearrange",): + dtype = "int64" + coincide = numpy.array_equal + if reduction in ["mean", "prod"]: + dtype = "float64" + coincide = numpy.allclose + max_dim = 11 + if "oneflow" in backend.framework_name: + max_dim = 7 + if "paddle" in backend.framework_name: + max_dim = 9 + for n_axes in range(max_dim): + shape = numpy.random.randint(2, 4, size=n_axes) + permutation = numpy.random.permutation(n_axes) + skipped = 0 if reduction == "rearrange" else numpy.random.randint(n_axes + 1) + left = " ".join("x" + str(i) for i in range(n_axes)) + right = " ".join("x" + str(i) for i in permutation[skipped:]) + pattern = left + "->" + right + x = numpy.arange(1, 1 + numpy.prod(shape), dtype=dtype).reshape(shape) + if reduction == "prod": + x /= x.mean() # to avoid overflows + result1 = reduce(x, pattern, reduction=reduction) + result2 = x.transpose(permutation) + if skipped > 0: + result2 = getattr(result2, reduction)(axis=tuple(range(skipped))) + assert coincide(result1, result2) + check_op_against_numpy(backend, x, pattern, reduction=reduction, axes_lengths={}, is_symbolic=False) + + +@pytest.mark.kernels_ci +def test_reduction_with_callable_imperatives(): + x_numpy = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]).astype("float32") + x_numpy /= x_numpy.max() + + def logsumexp_torch(x, tuple_of_axes): + return x.logsumexp(tuple_of_axes) + + def logsumexp_tf(x, tuple_of_axes): + import tensorflow as tf + + return tf.reduce_logsumexp(x, tuple_of_axes) + + def logsumexp_keras(x, tuple_of_axes): + import tensorflow.keras.backend as k + + return k.logsumexp(x, tuple_of_axes) + + def logsumexp_numpy(x, tuple_of_axes): + # very naive logsumexp to compare to + minused = x.max(tuple_of_axes) + y = x - x.max(tuple_of_axes, keepdims=True) + y = numpy.exp(y) + y = numpy.sum(y, axis=tuple_of_axes) + return numpy.log(y) + minused + + TorchBackend = einops._backends.TorchBackend + TensorflowBackend = einops._backends.TensorflowBackend + TFKerasBackend = einops._backends.TFKerasBackend + NumpyBackend = einops._backends.NumpyBackend + + backend2callback = { + TorchBackend.framework_name: logsumexp_torch, + TensorflowBackend.framework_name: logsumexp_tf, + TFKerasBackend.framework_name: logsumexp_keras, + NumpyBackend.framework_name: logsumexp_numpy, + } + + for backend in imp_op_backends: + if backend.framework_name not in backend2callback: + continue + + backend_callback = backend2callback[backend.framework_name] + + x_backend = backend.from_numpy(x_numpy) + for pattern1, pattern2 in equivalent_reduction_patterns: + print("Test reduction with callable for ", backend.framework_name, pattern1, pattern2) + output_numpy = reduce(x_numpy, pattern1, reduction=logsumexp_numpy) + output_backend = reduce(x_backend, pattern1, reduction=backend_callback) + assert numpy.allclose( + output_numpy, + backend.to_numpy(output_backend), + ) + + +@pytest.mark.kernels_ci +def test_enumerating_directions(): + for backend in imp_op_backends: + print("testing directions for", backend.framework_name) + for shape in [[], [1], [1, 1, 1], [2, 3, 5, 7]]: + x = numpy.arange(numpy.prod(shape)).reshape(shape) + axes1 = _enumerate_directions(x) + axes2 = _enumerate_directions(backend.from_numpy(x)) + assert len(axes1) == len(axes2) == len(shape) + for ax1, ax2 in zip(axes1, axes2): + ax2 = backend.to_numpy(ax2) + assert ax1.shape == ax2.shape + assert numpy.allclose(ax1, ax2) + + +@pytest.mark.kernels_ci +def test_concatenations_and_stacking(): + for backend in imp_op_backends: + print("testing shapes for ", backend.framework_name) + for n_arrays in [1, 2, 5]: + shapes = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6] + for shape in shapes: + arrays1 = [numpy.arange(i, i + numpy.prod(shape)).reshape(shape) for i in range(n_arrays)] + arrays2 = [backend.from_numpy(array) for array in arrays1] + result0 = numpy.asarray(arrays1) + result1 = rearrange(arrays1, "...->...") + result2 = rearrange(arrays2, "...->...") + assert numpy.array_equal(result0, result1) + assert numpy.array_equal(result1, backend.to_numpy(result2)) + + result1 = rearrange(arrays1, "b ... -> ... b") + result2 = rearrange(arrays2, "b ... -> ... b") + assert numpy.array_equal(result1, backend.to_numpy(result2)) + + +@pytest.mark.kernels_ci +def test_gradients_imperatives(): + # lazy - just checking reductions + for reduction in REDUCTIONS: + if reduction in ("any", "all"): + continue # non-differentiable ops + x = numpy.arange(1, 1 + 2 * 3 * 4).reshape([2, 3, 4]).astype("float32") + results = {} + for backend in imp_op_backends: + y0 = backend.from_numpy(x) + if not hasattr(y0, "grad"): + continue + + y1 = reduce(y0, "a b c -> c a", reduction=reduction) + y2 = reduce(y1, "c a -> a c", reduction=reduction) + y3 = reduce(y2, "a (c1 c2) -> a", reduction=reduction, c1=2) + y4 = reduce(y3, "... -> ", reduction=reduction) + + y4.backward() + grad = backend.to_numpy(y0.grad) + results[backend.framework_name] = grad + + print("comparing gradients for", results.keys()) + for name1, grad1 in results.items(): + for name2, grad2 in results.items(): + assert numpy.allclose(grad1, grad2), [name1, name2, "provided different gradients"] + + +@pytest.mark.kernels_ci +def test_tiling_imperatives(): + for backend in imp_op_backends: + print("Tiling tests for ", backend.framework_name) + input = numpy.arange(2 * 3 * 5, dtype="int64").reshape([2, 1, 3, 1, 5]) + test_cases = [ + (1, 1, 1, 1, 1), + (1, 2, 1, 3, 1), + (3, 1, 1, 4, 1), + ] + for repeats in test_cases: + expected = numpy.tile(input, repeats) + converted = backend.from_numpy(input) + repeated = backend.tile(converted, repeats) + result = backend.to_numpy(repeated) + assert numpy.array_equal(result, expected) + + +@pytest.mark.kernels_ci +def test_tiling_symbolic(): + for backend in sym_op_backends: + print("Tiling tests for ", backend.framework_name) + input = numpy.arange(2 * 3 * 5, dtype="int64").reshape([2, 1, 3, 1, 5]) + test_cases = [ + (1, 1, 1, 1, 1), + (1, 2, 1, 3, 1), + (3, 1, 1, 4, 1), + ] + for repeats in test_cases: + expected = numpy.tile(input, repeats) + sym = backend.create_symbol(input.shape) + result = backend.eval_symbol(backend.tile(sym, repeats), [[sym, input]]) + assert numpy.array_equal(result, expected) + + sym = backend.create_symbol([None] * len(input.shape)) + result = backend.eval_symbol(backend.tile(sym, repeats), [[sym, input]]) + assert numpy.array_equal(result, expected) + + +repeat_test_cases = [ + # all assume that input has shape [2, 3, 5] + ("a b c -> c a b", dict()), + ("a b c -> (c copy a b)", dict(copy=2, a=2, b=3, c=5)), + ("a b c -> (a copy) b c ", dict(copy=1)), + ("a b c -> (c a) (copy1 b copy2)", dict(a=2, copy1=1, copy2=2)), + ("a ... -> a ... copy", dict(copy=4)), + ("... c -> ... (copy1 c copy2)", dict(copy1=1, copy2=2)), + ("... -> ... ", dict()), + (" ... -> copy1 ... copy2 ", dict(copy1=2, copy2=3)), + ("a b c -> copy1 a copy2 b c () ", dict(copy1=2, copy2=1)), +] + + +def check_reversion(x, repeat_pattern, **sizes): + """Checks repeat pattern by running reduction""" + left, right = repeat_pattern.split("->") + reduce_pattern = right + "->" + left + repeated = repeat(x, repeat_pattern, **sizes) + reduced_min = reduce(repeated, reduce_pattern, reduction="min", **sizes) + reduced_max = reduce(repeated, reduce_pattern, reduction="max", **sizes) + assert numpy.array_equal(x, reduced_min) + assert numpy.array_equal(x, reduced_max) + + +@pytest.mark.kernels_ci +def test_repeat_numpy(): + # check repeat vs reduce. Repeat works ok if reverse reduction with min and max work well + x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) + x1 = repeat(x, "a b c -> copy a b c ", copy=1) + assert numpy.array_equal(x[None], x1) + for pattern, axis_dimensions in repeat_test_cases: + check_reversion(x, pattern, **axis_dimensions) + + +@pytest.mark.kernels_ci +def test_repeat_imperatives(): + x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) + for backend in imp_op_backends: + print("Repeat tests for ", backend.framework_name) + + for pattern, axis_dimensions in repeat_test_cases: + expected = repeat(x, pattern, **axis_dimensions) + converted = backend.from_numpy(x) + repeated = repeat(converted, pattern, **axis_dimensions) + result = backend.to_numpy(repeated) + assert numpy.array_equal(result, expected) + + +@pytest.mark.kernels_ci +def test_repeat_symbolic(): + x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) + + for backend in sym_op_backends: + print("Repeat tests for ", backend.framework_name) + + for pattern, axis_dimensions in repeat_test_cases: + expected = repeat(x, pattern, **axis_dimensions) + + sym = backend.create_symbol(x.shape) + result = backend.eval_symbol(repeat(sym, pattern, **axis_dimensions), [[sym, x]]) + assert numpy.array_equal(result, expected) + + +@pytest.mark.kernels_ci +def test_repeat_array_api(): + import numpy as xp + AA = einops.array_api + + if xp.__version__ < "2.0.0": + pytest.skip() + + x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) + + for pattern, axis_dimensions in repeat_test_cases: + expected = repeat(x, pattern, **axis_dimensions) + + result = AA.repeat(xp.from_dlpack(x), pattern, **axis_dimensions) + assert numpy.array_equal(AA.asnumpy(result + 0), expected) + + +test_cases_repeat_anonymous = [ + # all assume that input has shape [1, 2, 4, 6] + ("a b c d -> c a d b", dict()), + ("a b c d -> (c 2 d a b)", dict(a=1, c=4, d=6)), + ("1 b c d -> (d copy 1) 3 b c ", dict(copy=3)), + ("1 ... -> 3 ... ", dict()), + ("() ... d -> 1 (copy1 d copy2) ... ", dict(copy1=2, copy2=3)), + ("1 b c d -> (1 1) (1 b) 2 c 3 d (1 1)", dict()), +] + + +@pytest.mark.kernels_ci +def test_anonymous_axes(): + x = numpy.arange(1 * 2 * 4 * 6).reshape([1, 2, 4, 6]) + for pattern, axis_dimensions in test_cases_repeat_anonymous: + check_reversion(x, pattern, **axis_dimensions) + + +@pytest.mark.kernels_ci +def test_list_inputs(): + x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) + + assert numpy.array_equal( + rearrange(list(x), "... -> (...)"), + rearrange(x, "... -> (...)"), + ) + assert numpy.array_equal( + reduce(list(x), "a ... e -> (...)", "min"), + reduce(x, "a ... e -> (...)", "min"), + ) + assert numpy.array_equal( + repeat(list(x), "... -> b (...)", b=3), + repeat(x, "... -> b (...)", b=3), + ) + + +@pytest.mark.kernels_ci +def test_torch_compile_with_dynamic_shape(): + if not is_backend_tested("torch"): + pytest.skip() + import torch + + # somewhat reasonable debug messages + torch._dynamo.config.verbose = True + + def func1(x): + # test contains ellipsis + a, b, c, *other = x.shape + x = rearrange(x, "(a a2) b c ... -> b (c a2) (a ...)", a2=2) + # test contains passing expression as axis length + x = reduce(x, "b ca2 A -> b A", "sum", ca2=c * 2) + return x + + # seems can't test static and dynamic in the same test run. + # func1_compiled_static = torch.compile(func1, dynamic=False, fullgraph=True, backend='aot_eager') + func1_compiled_dynamic = torch.compile(func1, dynamic=True, fullgraph=True, backend="aot_eager") + + x = torch.randn(size=[4, 5, 6, 3]) + assert torch.equal(func1_compiled_dynamic(x), func1(x)) + # check with input of different dimensionality, and with all shape elements changed + x = torch.randn(size=[6, 3, 4, 2, 3]) + assert torch.equal(func1_compiled_dynamic(x), func1(x)) + + +def bit_count(x): + return sum((x >> i) & 1 for i in range(20)) + + +@pytest.mark.kernels_ci +def test_reduction_imperatives_booleans(): + """Checks that any/all reduction works in all frameworks""" + x_np = numpy.asarray([(bit_count(x) % 2) == 0 for x in range(2**6)]).reshape([2] * 6) + for backend in imp_op_backends: + print("Reduction any/all tests for ", backend.framework_name) + + for axis in range(6): + expected_result_any = numpy.any(x_np, axis=axis, keepdims=True) + expected_result_all = numpy.all(x_np, axis=axis, keepdims=True) + assert not numpy.array_equal(expected_result_any, expected_result_all) + + axes = list("abcdef") + axes_in = list(axes) + axes_out = list(axes) + axes_out[axis] = "1" + pattern = (" ".join(axes_in)) + " -> " + (" ".join(axes_out)) + + res_any = reduce(backend.from_numpy(x_np), pattern, reduction="any") + res_all = reduce(backend.from_numpy(x_np), pattern, reduction="all") + + assert numpy.array_equal(expected_result_any, backend.to_numpy(res_any)) + assert numpy.array_equal(expected_result_all, backend.to_numpy(res_all)) + + # expected result: any/all + expected_result_any = numpy.any(x_np, axis=(0, 1), keepdims=True) + expected_result_all = numpy.all(x_np, axis=(0, 1), keepdims=True) + pattern = "a b ... -> 1 1 ..." + res_any = reduce(backend.from_numpy(x_np), pattern, reduction="any") + res_all = reduce(backend.from_numpy(x_np), pattern, reduction="all") + assert numpy.array_equal(expected_result_any, backend.to_numpy(res_any)) + assert numpy.array_equal(expected_result_all, backend.to_numpy(res_all)) diff --git a/.ports/einops/overlay/tests/test_other.py b/.ports/einops/overlay/tests/test_other.py new file mode 100644 index 000000000..063f3298a --- /dev/null +++ b/.ports/einops/overlay/tests/test_other.py @@ -0,0 +1,307 @@ +from doctest import testmod + +import numpy +import pytest + +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +AbstractBackend = einops._backends.AbstractBackend +rearrange = einops.einops.rearrange +parse_shape = einops.einops.parse_shape +_optimize_transformation = einops.einops._optimize_transformation +from . import collect_test_backends, is_backend_tested + +__author__ = "Alex Rogozhnikov" + + +@pytest.mark.kernels_ci +def test_doctests_examples(): + # tests docstrings, additionally + testmod(einops.layers, raise_on_error=True, extraglobs=dict(np=numpy)) + testmod(einops.einops, raise_on_error=True, extraglobs=dict(np=numpy)) + + +@pytest.mark.kernels_ci +def test_backends_installed(): + """ + This test will fail if some of backends are not installed or can't be imported + Other tests will just work and only test installed backends. + """ + from . import parse_backends_to_test + + backends_to_test = parse_backends_to_test() + errors = [] + for backend_type in AbstractBackend.__subclasses__(): + if backend_type.framework_name not in backends_to_test: + continue + try: + # instantiate + backend_type() + except Exception as e: + errors.append((backend_type.framework_name, e)) + assert len(errors) == 0, errors + + +@pytest.mark.kernels_ci +def test_optimize_transformations_numpy(): + print("Testing optimizations") + shapes = [[2] * n_dimensions for n_dimensions in range(14)] + shapes += [[3] * n_dimensions for n_dimensions in range(6)] + shapes += [[2, 3, 5, 7]] + shapes += [[2, 3, 5, 7, 11, 17]] + + for shape in shapes: + for attempt in range(5): + n_dimensions = len(shape) + x = numpy.random.randint(0, 2**12, size=shape).reshape([-1]) + init_shape = shape[:] + n_reduced = numpy.random.randint(0, n_dimensions + 1) + reduced_axes = tuple(numpy.random.permutation(n_dimensions)[:n_reduced]) + axes_reordering = numpy.random.permutation(n_dimensions - n_reduced) + final_shape = numpy.random.randint(0, 1024, size=333) # just random + + init_shape2, reduced_axes2, axes_reordering2, final_shape2 = combination2 = _optimize_transformation( + init_shape, reduced_axes, axes_reordering, final_shape + ) + + assert numpy.array_equal(final_shape, final_shape2) + result1 = x.reshape(init_shape).sum(axis=reduced_axes).transpose(axes_reordering).reshape([-1]) + result2 = x.reshape(init_shape2).sum(axis=reduced_axes2).transpose(axes_reordering2).reshape([-1]) + assert numpy.array_equal(result1, result2) + + # testing we can't optimize this formula again + combination3 = _optimize_transformation(*combination2) + for a, b in zip(combination2, combination3): + assert numpy.array_equal(a, b) + + +_IMPERATIVE_BACKENDS = collect_test_backends(symbolic=False, layers=False) + +x_np = numpy.zeros([10, 20, 30, 40]) + + +@pytest.mark.kernels_ci +def test_parse_shape_imperative(): + for backend in _IMPERATIVE_BACKENDS: + print("Shape parsing for ", backend.framework_name) + parsed1 = parse_shape(x_np, "a b c d") + parsed2 = parse_shape(backend.from_numpy(x_np), "a b c d") + assert parsed1 == parsed2 == dict(a=10, b=20, c=30, d=40) + assert parsed1 != dict(a=1, b=20, c=30, d=40) != parsed2 + + +@pytest.mark.kernels_ci +def test_underscore(): + for backend in _IMPERATIVE_BACKENDS: + parsed1 = parse_shape(x_np, "_ _ _ _") + parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ _ _") + assert parsed1 == parsed2 == dict() + + +@pytest.mark.kernels_ci +def test_underscore_one(): + for backend in _IMPERATIVE_BACKENDS: + parsed1 = parse_shape(x_np, "_ _ _ hello") + parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ _ hello") + assert parsed1 == parsed2 == dict(hello=40) + + +@pytest.mark.kernels_ci +def test_underscore_several(): + for backend in _IMPERATIVE_BACKENDS: + parsed1 = parse_shape(x_np, "_ _ a1 a1a111a") + parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ a1 a1a111a") + assert parsed1 == parsed2 == dict(a1=30, a1a111a=40) + + +@pytest.mark.kernels_ci +def test_repeating(): + with pytest.raises(einops.EinopsError): + parse_shape(x_np, "a a b b") + + for backend in _IMPERATIVE_BACKENDS: + with pytest.raises(einops.EinopsError): + parse_shape(backend.from_numpy(x_np), "a a b b") + + +@pytest.mark.kernels_ci +def test_ellipsis(): + for backend in _IMPERATIVE_BACKENDS: + for shape, pattern, expected in [ + ([10, 20], "...", dict()), + ([10], "... a", dict(a=10)), + ([10, 20], "... a", dict(a=20)), + ([10, 20, 30], "... a", dict(a=30)), + ([10, 20, 30, 40], "... a", dict(a=40)), + ([10], "a ...", dict(a=10)), + ([10, 20], "a ...", dict(a=10)), + ([10, 20, 30], "a ...", dict(a=10)), + ([10, 20, 30, 40], "a ...", dict(a=10)), + ([10, 20, 30, 40], " a ... b", dict(a=10, b=40)), + ([10, 40], " a ... b", dict(a=10, b=40)), + ]: + x = numpy.ones(shape) + parsed1 = parse_shape(x, pattern) + parsed2 = parse_shape(backend.from_numpy(x), pattern) + assert parsed1 == parsed2 == expected + + +@pytest.mark.kernels_ci +def test_parse_with_anonymous_axes(): + for backend in _IMPERATIVE_BACKENDS: + for shape, pattern, expected in [ + ([1, 2, 3, 4], "1 2 3 a", dict(a=4)), + ([10, 1, 2], "a 1 2", dict(a=10)), + ([10, 1, 2], "a () 2", dict(a=10)), + ]: + x = numpy.ones(shape) + parsed1 = parse_shape(x, pattern) + parsed2 = parse_shape(backend.from_numpy(x), pattern) + assert parsed1 == parsed2 == expected + + +@pytest.mark.kernels_ci +def test_failures(): + for backend in _IMPERATIVE_BACKENDS: + # every test should fail + for shape, pattern in [ + ([1, 2, 3, 4], "a b c"), + ([1, 2, 3, 4], "2 a b c"), + ([1, 2, 3, 4], "a b c ()"), + ([1, 2, 3, 4], "a b c d e"), + ([1, 2, 3, 4], "a b c d e ..."), + ([1, 2, 3, 4], "a b c ()"), + ]: + with pytest.raises(RuntimeError): + x = numpy.ones(shape) + parse_shape(backend.from_numpy(x), pattern) + + +_SYMBOLIC_BACKENDS = [ + *collect_test_backends(symbolic=True, layers=False), + *collect_test_backends(symbolic=True, layers=True), +] + +# tensorflow.keras needs special way to compile, +# shape vars can be used only inside layers but not as outputs +_SYMBOLIC_BACKENDS = [backend for backend in _SYMBOLIC_BACKENDS if backend.framework_name != "tensorflow.keras"] + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) +def test_parse_shape_symbolic(backend): + for shape in [ + [10, 20, 30, 40], + [10, 20, None, None], + [None, None, None, None], + ]: + print( + f"special shape parsing {backend.framework_name=} {shape=}", + ) + input_symbol = backend.create_symbol(shape) + + shape_placeholder = parse_shape(input_symbol, "a b c d") + shape = {} + for name, symbol in shape_placeholder.items(): + shape[name] = ( + symbol + if isinstance(symbol, int) + else backend.eval_symbol(symbol, [(input_symbol, numpy.zeros([10, 20, 30, 40]))]) + ) + print(shape) + result_placeholder = rearrange( + input_symbol, "a b (c1 c2) (d1 d2) -> (a b d1) c1 (c2 d2)", **parse_shape(input_symbol, "a b c1 _"), d2=2 + ) + result = backend.eval_symbol(result_placeholder, [(input_symbol, numpy.zeros([10, 20, 30, 40]))]) + print(result.shape) + assert result.shape == (10 * 20 * 20, 30, 1 * 2) + assert numpy.allclose(result, 0) + + +@pytest.mark.kernels_ci +@pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) +def test_parse_shape_symbolic_ellipsis(backend): + for static_shape, shape, pattern, expected in [ + ([10, 20], [None, None], "...", dict()), + ([10], [None], "... a", dict(a=10)), + ([10, 20], [None, None], "... a", dict(a=20)), + ([10, 20, 30], [None, None, None], "... a", dict(a=30)), + ([10, 20, 30, 40], [None, None, None, None], "... a", dict(a=40)), + ([10], [None], "a ...", dict(a=10)), + ([10, 20], [None, None], "a ...", dict(a=10)), + ([10, 20, 30], [None, None, None], "a ...", dict(a=10)), + ([10, 20, 30, 40], [None, None, None, None], "a ...", dict(a=10)), + ([10, 20, 30, 40], [None, None, None, None], " a ... b", dict(a=10, b=40)), + ([10, 40], [None, None], " a ... b ", dict(a=10, b=40)), + ]: + input_symbol = backend.create_symbol(shape) + shape_placeholder = parse_shape(input_symbol, pattern) + out_shape = {} + for name, symbol in shape_placeholder.items(): + if isinstance(symbol, int): + out_shape[name] = symbol + else: + out_shape[name] = backend.eval_symbol(symbol, [(input_symbol, numpy.zeros(static_shape))]) + assert out_shape == expected + + +@pytest.mark.kernels_ci +def test_is_float_type(): + backends = collect_test_backends(symbolic=False, layers=False) + backends += collect_test_backends(symbolic=False, layers=True) + for backend in backends: + for dtype in ["int32", "int64", "float32", "float64"]: + is_float = "float" in dtype + input = numpy.zeros([3, 4, 5], dtype=dtype) + input = backend.from_numpy(input) + assert backend.is_float_type(input) == is_float, (dtype, backend, input.dtype) + + +@pytest.mark.kernels_ci +def test_torch_compile(): + """ + Test ensures that allow_ops_in_compiled_graph allows compiling in a single graph + Additionally we ensure that after compilation cache works properly + (by changing shapes and patterns) + We additionally check that pack/unpack still can be handled + despite variable number of inputs/outputs + """ + if not is_backend_tested("torch"): + pytest.skip() + import torch + from torch import nn + repeat, reduce, pack, unpack, einsum = einops.repeat, einops.reduce, einops.pack, einops.unpack, einops.einsum + allow_ops_in_compiled_graph = einops._torch_specific.allow_ops_in_compiled_graph + + allow_ops_in_compiled_graph() + + class TorchModuleWithOperations(nn.Module): + def __init__(self) -> None: + super().__init__() + + def forward(self, x_abc, suffix=""): + a, b, c = x_abc.shape + + def suf(pattern): + parts = pattern.split() + return " ".join([p if p[-1] not in "acd" else p + suffix for p in parts]) + + # patterns look a bit strange because names a, c, d will be modified on every run + # by suf function + x_abcd = repeat(x_abc, suf("a b c -> a b c 4")) + x_abc = reduce(x_abcd, suf("a b c d -> a b c"), "min") + x_abdc, ps = pack([x_abc] * (2 + len(suffix)), suf("a b * c")) + x_array = unpack(rearrange(x_abdc, suf("a b d c -> (a b ) 1 c d")), ps, "ab one1 c *") + x1 = x_array[0] + len(x_array) + x1 = rearrange(x1, suf("(a b ) 1 c -> a b c"), b=b) + addition = einsum(x_abc, x_abcd, suf("a b c , a b c d -> d"))[0] + return x1 + addition + + original = TorchModuleWithOperations() + compiled = torch.compile(original, fullgraph=True, backend="aot_eager") + for size in [10, 20, 40]: + x = torch.rand([size, size + 1, size + 2]) + for suffix in ["", "suf1", "other_suffix"]: + result1 = compiled(x, suffix) + result2 = original(x, suffix) + assert torch.allclose(result1, result2) diff --git a/.ports/einops/overlay/tests/test_packing.py b/.ports/einops/overlay/tests/test_packing.py new file mode 100644 index 000000000..d76553475 --- /dev/null +++ b/.ports/einops/overlay/tests/test_packing.py @@ -0,0 +1,318 @@ +import dataclasses +import typing + +import numpy as np +import pytest + +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +asnumpy = einops.asnumpy +pack = einops.pack +unpack = einops.unpack +from . import collect_test_backends + + +def pack_unpack(xs, pattern): + x, ps = pack(xs, pattern) + unpacked = unpack(xs, ps, pattern) + assert len(unpacked) == len(xs) + for a, b in zip(unpacked, xs): + assert np.allclose(asnumpy(a), asnumpy(b)) + + +def unpack_and_pack(x, ps, pattern: str): + unpacked = unpack(x, ps, pattern) + packed, ps2 = pack(unpacked, pattern=pattern) + + assert np.allclose(asnumpy(packed), asnumpy(x)) + return unpacked + + +def unpack_and_pack_against_numpy(x, ps, pattern: str): + capturer_backend = CaptureException() + capturer_numpy = CaptureException() + + with capturer_backend: + unpacked = unpack(x, ps, pattern) + packed, ps2 = pack(unpacked, pattern=pattern) + + with capturer_numpy: + x_np = asnumpy(x) + unpacked_np = unpack(x_np, ps, pattern) + packed_np, ps3 = pack(unpacked_np, pattern=pattern) + + assert type(capturer_numpy.exception) == type(capturer_backend.exception) # noqa E721 + if capturer_numpy.exception is not None: + # both failed + return + else: + # neither failed, check results are identical + assert np.allclose(asnumpy(packed), asnumpy(x)) + assert np.allclose(asnumpy(packed_np), asnumpy(x)) + assert len(unpacked) == len(unpacked_np) + for a, b in zip(unpacked, unpacked_np): + assert np.allclose(asnumpy(a), b) + + +class CaptureException: + def __enter__(self): + self.exception = None + + def __exit__(self, exc_type, exc_val, exc_tb): + self.exception = exc_val + return True + + +@pytest.mark.kernels_ci +def test_numpy_trivial(H=13, W=17): + def rand(*shape): + return np.random.random(shape) + + def check(a, b): + assert a.dtype == b.dtype + assert a.shape == b.shape + assert np.all(a == b) + + r, g, b = rand(3, H, W) + embeddings = rand(H, W, 32) + + check( + np.stack([r, g, b], axis=2), + pack([r, g, b], "h w *")[0], + ) + check( + np.stack([r, g, b], axis=1), + pack([r, g, b], "h * w")[0], + ) + check( + np.stack([r, g, b], axis=0), + pack([r, g, b], "* h w")[0], + ) + + check( + np.concatenate([r, g, b], axis=1), + pack([r, g, b], "h *")[0], + ) + check( + np.concatenate([r, g, b], axis=0), + pack([r, g, b], "* w")[0], + ) + + i = np.index_exp[:, :, None] + check( + np.concatenate([r[i], g[i], b[i], embeddings], axis=2), + pack([r, g, b, embeddings], "h w *")[0], + ) + + with pytest.raises(EinopsError): + pack([r, g, b, embeddings], "h w nonexisting_axis *") + + pack([r, g, b], "some_name_for_H some_name_for_w1 *") + + with pytest.raises(EinopsError): + pack([r, g, b, embeddings], "h _w *") # no leading underscore + with pytest.raises(EinopsError): + pack([r, g, b, embeddings], "h_ w *") # no trailing underscore + with pytest.raises(EinopsError): + pack([r, g, b, embeddings], "1h_ w *") + with pytest.raises(EinopsError): + pack([r, g, b, embeddings], "1 w *") + with pytest.raises(EinopsError): + pack([r, g, b, embeddings], "h h *") + # capital and non-capital are different + pack([r, g, b, embeddings], "h H *") + + +@dataclasses.dataclass +class UnpackTestCase: + shape: typing.Tuple[int, ...] + pattern: str + + def dim(self): + return self.pattern.split().index("*") + + def selfcheck(self): + assert self.shape[self.dim()] == 5 + + +cases = [ + # NB: in all cases unpacked axis is of length 5. + # that's actively used in tests below + UnpackTestCase((5,), "*"), + UnpackTestCase((5, 7), "* seven"), + UnpackTestCase((7, 5), "seven *"), + UnpackTestCase((5, 3, 4), "* three four"), + UnpackTestCase((4, 5, 3), "four * three"), + UnpackTestCase((3, 4, 5), "three four *"), +] + + +@pytest.mark.kernels_ci +def test_pack_unpack_with_numpy(): + case: UnpackTestCase + + for case in cases: + shape = case.shape + pattern = case.pattern + + x = np.random.random(shape) + # all correct, no minus 1 + unpack_and_pack(x, [[2], [1], [2]], pattern) + # no -1, asking for wrong shapes + with pytest.raises(BaseException): + unpack_and_pack(x, [[2], [1], [2]], pattern + " non_existent_axis") + with pytest.raises(BaseException): + unpack_and_pack(x, [[2], [1], [1]], pattern) + with pytest.raises(BaseException): + unpack_and_pack(x, [[4], [1], [1]], pattern) + # all correct, with -1 + unpack_and_pack(x, [[2], [1], [-1]], pattern) + unpack_and_pack(x, [[2], [-1], [2]], pattern) + unpack_and_pack(x, [[-1], [1], [2]], pattern) + _, _, last = unpack_and_pack(x, [[2], [3], [-1]], pattern) + assert last.shape[case.dim()] == 0 + # asking for more elements than available + with pytest.raises(BaseException): + unpack(x, [[2], [4], [-1]], pattern) + # this one does not raise, because indexing x[2:1] just returns zero elements + # with pytest.raises(BaseException): + # unpack(x, [[2], [-1], [4]], pattern) + with pytest.raises(BaseException): + unpack(x, [[-1], [1], [5]], pattern) + + # all correct, -1 nested + rs = unpack_and_pack(x, [[1, 2], [1, 1], [-1, 1]], pattern) + assert all(len(r.shape) == len(x.shape) + 1 for r in rs) + rs = unpack_and_pack(x, [[1, 2], [1, -1], [1, 1]], pattern) + assert all(len(r.shape) == len(x.shape) + 1 for r in rs) + rs = unpack_and_pack(x, [[2, -1], [1, 2], [1, 1]], pattern) + assert all(len(r.shape) == len(x.shape) + 1 for r in rs) + + # asking for more elements, -1 nested + with pytest.raises(BaseException): + unpack(x, [[-1, 2], [1], [5]], pattern) + with pytest.raises(BaseException): + unpack(x, [[2, 2], [2], [5, -1]], pattern) + + # asking for non-divisible number of elements + with pytest.raises(BaseException): + unpack(x, [[2, 1], [1], [3, -1]], pattern) + with pytest.raises(BaseException): + unpack(x, [[2, 1], [3, -1], [1]], pattern) + with pytest.raises(BaseException): + unpack(x, [[3, -1], [2, 1], [1]], pattern) + + # -1 takes zero + unpack_and_pack(x, [[0], [5], [-1]], pattern) + unpack_and_pack(x, [[0], [-1], [5]], pattern) + unpack_and_pack(x, [[-1], [5], [0]], pattern) + + # -1 takes zero, -1 + unpack_and_pack(x, [[2, -1], [1, 5]], pattern) + + +@pytest.mark.kernels_ci +def test_pack_unpack_against_numpy(): + for backend in collect_test_backends(symbolic=False, layers=False): + print(f"test packing against numpy for {backend.framework_name}") + check_zero_len = True + + for case in cases: + unpack_and_pack = unpack_and_pack_against_numpy + shape = case.shape + pattern = case.pattern + + x = np.random.random(shape) + x = backend.from_numpy(x) + # all correct, no minus 1 + unpack_and_pack(x, [[2], [1], [2]], pattern) + # no -1, asking for wrong shapes + with pytest.raises(BaseException): + unpack(x, [[2], [1], [1]], pattern) + + with pytest.raises(BaseException): + unpack(x, [[4], [1], [1]], pattern) + # all correct, with -1 + unpack_and_pack(x, [[2], [1], [-1]], pattern) + unpack_and_pack(x, [[2], [-1], [2]], pattern) + unpack_and_pack(x, [[-1], [1], [2]], pattern) + + # asking for more elements than available + with pytest.raises(BaseException): + unpack(x, [[2], [4], [-1]], pattern) + # this one does not raise, because indexing x[2:1] just returns zero elements + # with pytest.raises(BaseException): + # unpack(x, [[2], [-1], [4]], pattern) + with pytest.raises(BaseException): + unpack(x, [[-1], [1], [5]], pattern) + + # all correct, -1 nested + unpack_and_pack(x, [[1, 2], [1, 1], [-1, 1]], pattern) + unpack_and_pack(x, [[1, 2], [1, -1], [1, 1]], pattern) + unpack_and_pack(x, [[2, -1], [1, 2], [1, 1]], pattern) + + # asking for more elements, -1 nested + with pytest.raises(BaseException): + unpack(x, [[-1, 2], [1], [5]], pattern) + with pytest.raises(BaseException): + unpack(x, [[2, 2], [2], [5, -1]], pattern) + + # asking for non-divisible number of elements + with pytest.raises(BaseException): + unpack(x, [[2, 1], [1], [3, -1]], pattern) + with pytest.raises(BaseException): + unpack(x, [[2, 1], [3, -1], [1]], pattern) + with pytest.raises(BaseException): + unpack(x, [[3, -1], [2, 1], [1]], pattern) + + if check_zero_len: + # -1 takes zero + unpack_and_pack(x, [[2], [3], [-1]], pattern) + unpack_and_pack(x, [[0], [5], [-1]], pattern) + unpack_and_pack(x, [[0], [-1], [5]], pattern) + unpack_and_pack(x, [[-1], [5], [0]], pattern) + + # -1 takes zero, -1 + unpack_and_pack(x, [[2, -1], [1, 5]], pattern) + + +@pytest.mark.kernels_ci +def test_pack_unpack_array_api(): + AA = einops.array_api + import numpy as xp + + if xp.__version__ < "2.0.0": + pytest.skip() + + for case in cases: + shape = case.shape + pattern = case.pattern + x_np = np.random.random(shape) + x_xp = xp.from_dlpack(x_np) + + for ps in [ + [[2], [1], [2]], + [[1], [1], [-1]], + [[1], [1], [-1, 3]], + [[2, 1], [1, 1, 1], [-1]], + ]: + x_np_split = unpack(x_np, ps, pattern) + x_xp_split = AA.unpack(x_xp, ps, pattern) + for a, b in zip(x_np_split, x_xp_split): + assert np.allclose(a, AA.asnumpy(b + 0)) + + x_agg_np, ps1 = pack(x_np_split, pattern) + x_agg_xp, ps2 = AA.pack(x_xp_split, pattern) + assert ps1 == ps2 + assert np.allclose(x_agg_np, AA.asnumpy(x_agg_xp)) + + for ps in [ + [[2, 3]], + [[1], [5]], + [[1], [5], [-1]], + [[1], [2, 3]], + [[1], [5], [-1, 2]], + ]: + with pytest.raises(BaseException): + unpack(x_np, ps, pattern) diff --git a/.ports/einops/overlay/tests/test_parsing.py b/.ports/einops/overlay/tests/test_parsing.py new file mode 100644 index 000000000..491126793 --- /dev/null +++ b/.ports/einops/overlay/tests/test_parsing.py @@ -0,0 +1,134 @@ +import pytest + +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +ParsedExpression = einops.parsing.ParsedExpression +AnonymousAxis = einops.parsing.AnonymousAxis +_ellipsis = einops.parsing._ellipsis + +__author__ = "Alex Rogozhnikov" + + +class AnonymousAxisPlaceholder: + def __init__(self, value: int): + self.value = value + assert isinstance(self.value, int) + + def __eq__(self, other): + return isinstance(other, AnonymousAxis) and self.value == other.value + + +@pytest.mark.kernels_ci +def test_anonymous_axes(): + a, b = AnonymousAxis("2"), AnonymousAxis("2") + assert a != b + c, d = AnonymousAxisPlaceholder(2), AnonymousAxisPlaceholder(3) + assert a == c and b == c + assert a != d and b != d + assert [a, 2, b] == [c, 2, c] + + +@pytest.mark.kernels_ci +def test_elementary_axis_name(): + for name in [ + "a", + "b", + "h", + "dx", + "h1", + "zz", + "i9123", + "somelongname", + "Alex", + "camelCase", + "u_n_d_e_r_score", + "unreasonablyLongAxisName", + ]: + assert ParsedExpression.check_axis_name(name) + + for name in ["", "2b", "12", "_startWithUnderscore", "endWithUnderscore_", "_", "...", _ellipsis]: + assert not ParsedExpression.check_axis_name(name) + + +@pytest.mark.kernels_ci +def test_invalid_expressions(): + # double ellipsis should raise an error + ParsedExpression("... a b c d") + with pytest.raises(EinopsError): + ParsedExpression("... a b c d ...") + with pytest.raises(EinopsError): + ParsedExpression("... a b c (d ...)") + with pytest.raises(EinopsError): + ParsedExpression("(... a) b c (d ...)") + + # double/missing/enclosed parenthesis + ParsedExpression("(a) b c (d ...)") + with pytest.raises(EinopsError): + ParsedExpression("(a)) b c (d ...)") + with pytest.raises(EinopsError): + ParsedExpression("(a b c (d ...)") + with pytest.raises(EinopsError): + ParsedExpression("(a) (()) b c (d ...)") + with pytest.raises(EinopsError): + ParsedExpression("(a) ((b c) (d ...))") + + # invalid identifiers + ParsedExpression("camelCase under_scored cApiTaLs ß ...") + with pytest.raises(EinopsError): + ParsedExpression("1a") + with pytest.raises(EinopsError): + ParsedExpression("_pre") + with pytest.raises(EinopsError): + ParsedExpression("...pre") + with pytest.raises(EinopsError): + ParsedExpression("pre...") + + +@pytest.mark.kernels_ci +def test_parse_expression(): + parsed = ParsedExpression("a1 b1 c1 d1") + assert parsed.identifiers == {"a1", "b1", "c1", "d1"} + assert parsed.composition == [["a1"], ["b1"], ["c1"], ["d1"]] + assert not parsed.has_non_unitary_anonymous_axes + assert not parsed.has_ellipsis + + parsed = ParsedExpression("() () () ()") + assert parsed.identifiers == set() + assert parsed.composition == [[], [], [], []] + assert not parsed.has_non_unitary_anonymous_axes + assert not parsed.has_ellipsis + + parsed = ParsedExpression("1 1 1 ()") + assert parsed.identifiers == set() + assert parsed.composition == [[], [], [], []] + assert not parsed.has_non_unitary_anonymous_axes + assert not parsed.has_ellipsis + + aap = AnonymousAxisPlaceholder + + parsed = ParsedExpression("5 (3 4)") + assert len(parsed.identifiers) == 3 and {i.value for i in parsed.identifiers} == {3, 4, 5} + assert parsed.composition == [[aap(5)], [aap(3), aap(4)]] + assert parsed.has_non_unitary_anonymous_axes + assert not parsed.has_ellipsis + + parsed = ParsedExpression("5 1 (1 4) 1") + assert len(parsed.identifiers) == 2 and {i.value for i in parsed.identifiers} == {4, 5} + assert parsed.composition == [[aap(5)], [], [aap(4)], []] + + parsed = ParsedExpression("name1 ... a1 12 (name2 14)") + assert len(parsed.identifiers) == 6 + assert parsed.identifiers.difference({"name1", _ellipsis, "a1", "name2"}).__len__() == 2 + assert parsed.composition == [["name1"], _ellipsis, ["a1"], [aap(12)], ["name2", aap(14)]] + assert parsed.has_non_unitary_anonymous_axes + assert parsed.has_ellipsis + assert not parsed.has_ellipsis_parenthesized + + parsed = ParsedExpression("(name1 ... a1 12) name2 14") + assert len(parsed.identifiers) == 6 + assert parsed.identifiers.difference({"name1", _ellipsis, "a1", "name2"}).__len__() == 2 + assert parsed.composition == [["name1", _ellipsis, "a1", aap(12)], ["name2"], [aap(14)]] + assert parsed.has_non_unitary_anonymous_axes + assert parsed.has_ellipsis + assert parsed.has_ellipsis_parenthesized diff --git a/.ports/einops/port.kdl b/.ports/einops/port.kdl new file mode 100644 index 000000000..2b92e57fb --- /dev/null +++ b/.ports/einops/port.kdl @@ -0,0 +1,60 @@ +recipe version=1 + +// einops +// +// einops as a Hub kernel: the package is pure Python, so the manifest is +// noarch and the port is a move plus the import rewrite that lets the package +// be imported under whatever directory name the Hub loads it from. The test +// suite is not derived - it was rewritten to pull the kernel off the Hub - so +// it comes from the overlay. +// +// kernel-port .ports/einops/port.kdl \ +// --dir --out einops + +source \ + repo="https://github.com/arogozhnikov/einops.git" \ + commit="ad2c8d67d08cbebb45e1aff2be21dc88635fc695" + +prune keep="einops/**" +delete in="einops/tests/**" + +move \ + from="einops" \ + to="torch-ext/einops" + +relativize_imports \ + in="torch-ext/einops/**" \ + package_root="torch-ext/einops" \ + changes=7 +expect \ + in="torch-ext/einops/**" \ + find="from einops" \ + count=0 + +// array_api is a submodule upstream leaves out of the package init; the kernel +// exposes it, since a Hub consumer cannot import einops.array_api by name. +replace \ + in="torch-ext/einops/__init__.py" \ + count=1 \ + find="from .packing import pack, unpack\n" \ + with="from .packing import pack, unpack\nfrom . import array_api\n" + +// The rewritten tests and the flake. +overlay from="overlay" + +expect \ + in="torch-ext/einops/**" \ + files=18 +expect \ + in="tests/**" \ + files=9 + +manifest \ + name="einops" \ + version=1 \ + edition=5 \ + license="MIT" \ + upstream="https://github.com/arogozhnikov/einops.git" \ + backends="cpu,cuda,metal,rocm,xpu" \ + repo_id="kernels-community/einops" \ + noarch=#true diff --git a/einops/.port-provenance.json b/einops/.port-provenance.json new file mode 100644 index 000000000..c99c151c3 --- /dev/null +++ b/einops/.port-provenance.json @@ -0,0 +1,12 @@ +{ + "format": 1, + "recipe": "sha256:31835e78fc155e6165bfb986a8b7bd3c58897c62bc01c6d90de7a8e1d58d16ab", + "runner": "kernel-port 0.1.0", + "sources": { + "upstream": { + "repo": "https://github.com/arogozhnikov/einops.git", + "commit": "ad2c8d67d08cbebb45e1aff2be21dc88635fc695" + } + }, + "tree": "sha256:2764199b7ae5d7a03a0780d2907b9eda7a9ae6ade00d6a879c2ef2489a1be969" +} diff --git a/einops/build.toml b/einops/build.toml index fcf94d562..4d322d3b0 100644 --- a/einops/build.toml +++ b/einops/build.toml @@ -1,10 +1,16 @@ [general] name = "einops" version = 1 -edition = 5 license = "MIT" +edition = 5 upstream = "https://github.com/arogozhnikov/einops.git" -backends = ["cpu", "cuda", "metal", "rocm", "xpu"] +backends = [ + "cpu", + "cuda", + "metal", + "rocm", + "xpu", +] [general.hub] repo-id = "kernels-community/einops" From 47ba186c48e795d989476b2c8dfd90ccbe060db2 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Wed, 26 Aug 2026 17:34:13 -0400 Subject: [PATCH 02/15] feat: move einops kernel sources under src/ Signed-off-by: David Holtz --- einops/{ => src}/build.toml | 0 einops/{ => src}/flake.lock | 0 einops/{ => src}/flake.nix | 0 einops/{ => src}/tests/__init__.py | 0 einops/{ => src}/tests/run_tests.py | 0 einops/{ => src}/tests/test_einsum.py | 0 einops/{ => src}/tests/test_examples.py | 0 einops/{ => src}/tests/test_layers.py | 0 einops/{ => src}/tests/test_ops.py | 0 einops/{ => src}/tests/test_other.py | 0 einops/{ => src}/tests/test_packing.py | 0 einops/{ => src}/tests/test_parsing.py | 0 einops/{ => src}/torch-ext/einops/__init__.py | 0 einops/{ => src}/torch-ext/einops/_backends.py | 0 einops/{ => src}/torch-ext/einops/_torch_specific.py | 0 einops/{ => src}/torch-ext/einops/array_api.py | 0 einops/{ => src}/torch-ext/einops/einops.py | 0 einops/{ => src}/torch-ext/einops/experimental/__init__.py | 0 einops/{ => src}/torch-ext/einops/experimental/indexing.py | 0 einops/{ => src}/torch-ext/einops/layers/__init__.py | 0 einops/{ => src}/torch-ext/einops/layers/_einmix.py | 0 einops/{ => src}/torch-ext/einops/layers/flax.py | 0 einops/{ => src}/torch-ext/einops/layers/keras.py | 0 einops/{ => src}/torch-ext/einops/layers/oneflow.py | 0 einops/{ => src}/torch-ext/einops/layers/paddle.py | 0 einops/{ => src}/torch-ext/einops/layers/tensorflow.py | 0 einops/{ => src}/torch-ext/einops/layers/torch.py | 0 einops/{ => src}/torch-ext/einops/packing.py | 0 einops/{ => src}/torch-ext/einops/parsing.py | 0 einops/{ => src}/torch-ext/einops/py.typed | 0 30 files changed, 0 insertions(+), 0 deletions(-) rename einops/{ => src}/build.toml (100%) rename einops/{ => src}/flake.lock (100%) rename einops/{ => src}/flake.nix (100%) rename einops/{ => src}/tests/__init__.py (100%) rename einops/{ => src}/tests/run_tests.py (100%) rename einops/{ => src}/tests/test_einsum.py (100%) rename einops/{ => src}/tests/test_examples.py (100%) rename einops/{ => src}/tests/test_layers.py (100%) rename einops/{ => src}/tests/test_ops.py (100%) rename einops/{ => src}/tests/test_other.py (100%) rename einops/{ => src}/tests/test_packing.py (100%) rename einops/{ => src}/tests/test_parsing.py (100%) rename einops/{ => src}/torch-ext/einops/__init__.py (100%) rename einops/{ => src}/torch-ext/einops/_backends.py (100%) rename einops/{ => src}/torch-ext/einops/_torch_specific.py (100%) rename einops/{ => src}/torch-ext/einops/array_api.py (100%) rename einops/{ => src}/torch-ext/einops/einops.py (100%) rename einops/{ => src}/torch-ext/einops/experimental/__init__.py (100%) rename einops/{ => src}/torch-ext/einops/experimental/indexing.py (100%) rename einops/{ => src}/torch-ext/einops/layers/__init__.py (100%) rename einops/{ => src}/torch-ext/einops/layers/_einmix.py (100%) rename einops/{ => src}/torch-ext/einops/layers/flax.py (100%) rename einops/{ => src}/torch-ext/einops/layers/keras.py (100%) rename einops/{ => src}/torch-ext/einops/layers/oneflow.py (100%) rename einops/{ => src}/torch-ext/einops/layers/paddle.py (100%) rename einops/{ => src}/torch-ext/einops/layers/tensorflow.py (100%) rename einops/{ => src}/torch-ext/einops/layers/torch.py (100%) rename einops/{ => src}/torch-ext/einops/packing.py (100%) rename einops/{ => src}/torch-ext/einops/parsing.py (100%) rename einops/{ => src}/torch-ext/einops/py.typed (100%) diff --git a/einops/build.toml b/einops/src/build.toml similarity index 100% rename from einops/build.toml rename to einops/src/build.toml diff --git a/einops/flake.lock b/einops/src/flake.lock similarity index 100% rename from einops/flake.lock rename to einops/src/flake.lock diff --git a/einops/flake.nix b/einops/src/flake.nix similarity index 100% rename from einops/flake.nix rename to einops/src/flake.nix diff --git a/einops/tests/__init__.py b/einops/src/tests/__init__.py similarity index 100% rename from einops/tests/__init__.py rename to einops/src/tests/__init__.py diff --git a/einops/tests/run_tests.py b/einops/src/tests/run_tests.py similarity index 100% rename from einops/tests/run_tests.py rename to einops/src/tests/run_tests.py diff --git a/einops/tests/test_einsum.py b/einops/src/tests/test_einsum.py similarity index 100% rename from einops/tests/test_einsum.py rename to einops/src/tests/test_einsum.py diff --git a/einops/tests/test_examples.py b/einops/src/tests/test_examples.py similarity index 100% rename from einops/tests/test_examples.py rename to einops/src/tests/test_examples.py diff --git a/einops/tests/test_layers.py b/einops/src/tests/test_layers.py similarity index 100% rename from einops/tests/test_layers.py rename to einops/src/tests/test_layers.py diff --git a/einops/tests/test_ops.py b/einops/src/tests/test_ops.py similarity index 100% rename from einops/tests/test_ops.py rename to einops/src/tests/test_ops.py diff --git a/einops/tests/test_other.py b/einops/src/tests/test_other.py similarity index 100% rename from einops/tests/test_other.py rename to einops/src/tests/test_other.py diff --git a/einops/tests/test_packing.py b/einops/src/tests/test_packing.py similarity index 100% rename from einops/tests/test_packing.py rename to einops/src/tests/test_packing.py diff --git a/einops/tests/test_parsing.py b/einops/src/tests/test_parsing.py similarity index 100% rename from einops/tests/test_parsing.py rename to einops/src/tests/test_parsing.py diff --git a/einops/torch-ext/einops/__init__.py b/einops/src/torch-ext/einops/__init__.py similarity index 100% rename from einops/torch-ext/einops/__init__.py rename to einops/src/torch-ext/einops/__init__.py diff --git a/einops/torch-ext/einops/_backends.py b/einops/src/torch-ext/einops/_backends.py similarity index 100% rename from einops/torch-ext/einops/_backends.py rename to einops/src/torch-ext/einops/_backends.py diff --git a/einops/torch-ext/einops/_torch_specific.py b/einops/src/torch-ext/einops/_torch_specific.py similarity index 100% rename from einops/torch-ext/einops/_torch_specific.py rename to einops/src/torch-ext/einops/_torch_specific.py diff --git a/einops/torch-ext/einops/array_api.py b/einops/src/torch-ext/einops/array_api.py similarity index 100% rename from einops/torch-ext/einops/array_api.py rename to einops/src/torch-ext/einops/array_api.py diff --git a/einops/torch-ext/einops/einops.py b/einops/src/torch-ext/einops/einops.py similarity index 100% rename from einops/torch-ext/einops/einops.py rename to einops/src/torch-ext/einops/einops.py diff --git a/einops/torch-ext/einops/experimental/__init__.py b/einops/src/torch-ext/einops/experimental/__init__.py similarity index 100% rename from einops/torch-ext/einops/experimental/__init__.py rename to einops/src/torch-ext/einops/experimental/__init__.py diff --git a/einops/torch-ext/einops/experimental/indexing.py b/einops/src/torch-ext/einops/experimental/indexing.py similarity index 100% rename from einops/torch-ext/einops/experimental/indexing.py rename to einops/src/torch-ext/einops/experimental/indexing.py diff --git a/einops/torch-ext/einops/layers/__init__.py b/einops/src/torch-ext/einops/layers/__init__.py similarity index 100% rename from einops/torch-ext/einops/layers/__init__.py rename to einops/src/torch-ext/einops/layers/__init__.py diff --git a/einops/torch-ext/einops/layers/_einmix.py b/einops/src/torch-ext/einops/layers/_einmix.py similarity index 100% rename from einops/torch-ext/einops/layers/_einmix.py rename to einops/src/torch-ext/einops/layers/_einmix.py diff --git a/einops/torch-ext/einops/layers/flax.py b/einops/src/torch-ext/einops/layers/flax.py similarity index 100% rename from einops/torch-ext/einops/layers/flax.py rename to einops/src/torch-ext/einops/layers/flax.py diff --git a/einops/torch-ext/einops/layers/keras.py b/einops/src/torch-ext/einops/layers/keras.py similarity index 100% rename from einops/torch-ext/einops/layers/keras.py rename to einops/src/torch-ext/einops/layers/keras.py diff --git a/einops/torch-ext/einops/layers/oneflow.py b/einops/src/torch-ext/einops/layers/oneflow.py similarity index 100% rename from einops/torch-ext/einops/layers/oneflow.py rename to einops/src/torch-ext/einops/layers/oneflow.py diff --git a/einops/torch-ext/einops/layers/paddle.py b/einops/src/torch-ext/einops/layers/paddle.py similarity index 100% rename from einops/torch-ext/einops/layers/paddle.py rename to einops/src/torch-ext/einops/layers/paddle.py diff --git a/einops/torch-ext/einops/layers/tensorflow.py b/einops/src/torch-ext/einops/layers/tensorflow.py similarity index 100% rename from einops/torch-ext/einops/layers/tensorflow.py rename to einops/src/torch-ext/einops/layers/tensorflow.py diff --git a/einops/torch-ext/einops/layers/torch.py b/einops/src/torch-ext/einops/layers/torch.py similarity index 100% rename from einops/torch-ext/einops/layers/torch.py rename to einops/src/torch-ext/einops/layers/torch.py diff --git a/einops/torch-ext/einops/packing.py b/einops/src/torch-ext/einops/packing.py similarity index 100% rename from einops/torch-ext/einops/packing.py rename to einops/src/torch-ext/einops/packing.py diff --git a/einops/torch-ext/einops/parsing.py b/einops/src/torch-ext/einops/parsing.py similarity index 100% rename from einops/torch-ext/einops/parsing.py rename to einops/src/torch-ext/einops/parsing.py diff --git a/einops/torch-ext/einops/py.typed b/einops/src/torch-ext/einops/py.typed similarity index 100% rename from einops/torch-ext/einops/py.typed rename to einops/src/torch-ext/einops/py.typed From 4c01d31d26c39ef07914ffbd7f6ade8ef750bde3 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Wed, 26 Aug 2026 17:35:39 -0400 Subject: [PATCH 03/15] feat: generate einops src/ from a port recipe Signed-off-by: David Holtz --- .gitignore | 4 - .ports/einops/overlay/tests/__init__.py | 119 --- .ports/einops/overlay/tests/run_tests.py | 85 --- .ports/einops/overlay/tests/test_einsum.py | 360 --------- .ports/einops/overlay/tests/test_examples.py | 303 -------- .ports/einops/overlay/tests/test_layers.py | 484 ------------- .ports/einops/overlay/tests/test_ops.py | 684 ------------------ .ports/einops/overlay/tests/test_other.py | 307 -------- .ports/einops/overlay/tests/test_packing.py | 318 -------- .ports/einops/overlay/tests/test_parsing.py | 134 ---- .ports/einops/port.kdl | 60 -- .../einops => einops/port}/overlay/flake.lock | 0 .../einops => einops/port}/overlay/flake.nix | 0 einops/port/port.kdl | 197 +++++ .../port-provenance.json} | 4 +- einops/src/tests/test_other.py | 4 +- 16 files changed, 201 insertions(+), 2862 deletions(-) delete mode 100644 .ports/einops/overlay/tests/__init__.py delete mode 100644 .ports/einops/overlay/tests/run_tests.py delete mode 100644 .ports/einops/overlay/tests/test_einsum.py delete mode 100644 .ports/einops/overlay/tests/test_examples.py delete mode 100644 .ports/einops/overlay/tests/test_layers.py delete mode 100644 .ports/einops/overlay/tests/test_ops.py delete mode 100644 .ports/einops/overlay/tests/test_other.py delete mode 100644 .ports/einops/overlay/tests/test_packing.py delete mode 100644 .ports/einops/overlay/tests/test_parsing.py delete mode 100644 .ports/einops/port.kdl rename {.ports/einops => einops/port}/overlay/flake.lock (100%) rename {.ports/einops => einops/port}/overlay/flake.nix (100%) create mode 100644 einops/port/port.kdl rename einops/{.port-provenance.json => src/port-provenance.json} (55%) diff --git a/.gitignore b/.gitignore index 9a41b7635..4dcaf8670 100644 --- a/.gitignore +++ b/.gitignore @@ -16,10 +16,6 @@ build .* !.github/ -# Port files and provenance -!.ports/ -!.port-provenance.json - # Python stuff __pycache__/ diff --git a/.ports/einops/overlay/tests/__init__.py b/.ports/einops/overlay/tests/__init__.py deleted file mode 100644 index b8c629217..000000000 --- a/.ports/einops/overlay/tests/__init__.py +++ /dev/null @@ -1,119 +0,0 @@ -""" -Common utils for testing. -These functions allow testing only some frameworks, not all. -""" - -import logging -import os -from functools import lru_cache -from typing import List, Tuple - -import kernels -_einops = kernels.get_kernel("kernels-community/einops", version=1) -_backends = _einops._backends -import warnings - -__author__ = "Alex Rogozhnikov" - - -# minimize noise in tests logging -logging.getLogger("tensorflow").disabled = True -logging.getLogger("matplotlib").disabled = True - -FLOAT_REDUCTIONS = ("min", "max", "sum", "mean", "prod") # not includes any/all - - -def find_names_of_all_frameworks() -> List[str]: - backend_subclasses = [] - backends = _backends.AbstractBackend.__subclasses__() - while backends: - backend = backends.pop() - backends += backend.__subclasses__() - backend_subclasses.append(backend) - return [b.framework_name for b in backend_subclasses] - - -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 ENVVAR_NAME in os.environ: - 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: - raise RuntimeError(f"Unknown framework: {backend_name}") - return parsed_backends - - # Auto-detect available frameworks. - available = [] - for module, name in [("torch", "torch"), ("jax", "jax"), ("numpy", "numpy"), ("tensorflow", "tensorflow")]: - try: - __import__(module) - available.append(name) - except ImportError: - pass - return available - - -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 collect_test_backends(symbolic=False, layers=False) -> List[_backends.AbstractBackend]: - """ - :param symbolic: symbolic or imperative frameworks? - :param layers: layers or operations? - :return: list of backends satisfying set conditions - """ - if not symbolic: - if not layers: - backend_types = [ - _backends.NumpyBackend, - _backends.JaxBackend, - _backends.TorchBackend, - _backends.TensorflowBackend, - _backends.OneFlowBackend, - _backends.PaddleBackend, - _backends.CupyBackend, - ] - else: - backend_types = [ - _backends.TorchBackend, - _backends.OneFlowBackend, - _backends.PaddleBackend, - ] - else: - if not layers: - backend_types = [ - _backends.PyTensorBackend, - ] - else: - backend_types = [ - _backends.TFKerasBackend, - ] - - backend_names_to_test = parse_backends_to_test() - result = [] - for backend_type in backend_types: - if backend_type.framework_name not in backend_names_to_test: - continue - try: - result.append(backend_type()) - except ImportError: - # problem with backend installation fails a specific test function, - # but will be skipped in all other test cases - warnings.warn("backend could not be initialized for tests: {}".format(backend_type)) - return result diff --git a/.ports/einops/overlay/tests/run_tests.py b/.ports/einops/overlay/tests/run_tests.py deleted file mode 100644 index cf7a4eb6e..000000000 --- a/.ports/einops/overlay/tests/run_tests.py +++ /dev/null @@ -1,85 +0,0 @@ -""" -Runs tests that are appropriate for framework. -""" - -import os -import sys -from subprocess import Popen -from pathlib import Path - -__author__ = "Alex Rogozhnikov" - - -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 - - -def main(): - _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]", "flax"], - "tensorflow": ["tensorflow"], - "cupy": ["cupy"], - # switch to stable paddlepaddle, because of https://github.com/PaddlePaddle/Paddle/issues/63927 - # "paddle": ["paddlepaddle==0.0.0 -f https://www.paddlepaddle.org.cn/whl/linux/cpu-mkl/develop.html"], - "paddle": ["paddlepaddle"], - "oneflow": ["oneflow==0.9.0"], - "pytensor": ["pytensor"], - } - - usage = f""" - Usage: python -m einops.tests.run_tests [--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) - return - else: - synonyms = { - "tf": "tensorflow", - "pytorch": "torch", - "paddlepaddle": "paddle", - } - frameworks = [synonyms.get(f, f) for f in frameworks] - wrong_frameworks = [f for f in frameworks if f not in framework_name2installation] - if wrong_frameworks: - print(usage) - raise RuntimeError(f"Unrecognized frameworks: {wrong_frameworks}") - - 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))) - - 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 an envvar EINOPS_TEST_BACKENDS - from einops.tests import unparse_backends - - envvar_name, envvar_value = unparse_backends(backend_names=frameworks) - return_code = run( - "python -m pytest .", - **{envvar_name: envvar_value}, - ) - assert return_code == 0 - - -if __name__ == "__main__": - main() diff --git a/.ports/einops/overlay/tests/test_einsum.py b/.ports/einops/overlay/tests/test_einsum.py deleted file mode 100644 index 2e05d939d..000000000 --- a/.ports/einops/overlay/tests/test_einsum.py +++ /dev/null @@ -1,360 +0,0 @@ -from typing import Any, Callable -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -_compactify_pattern_for_einsum = einops.einops._compactify_pattern_for_einsum -einsum = einops.einops.einsum -EinopsError = einops.EinopsError -from . import collect_test_backends -import numpy as np -import pytest -import string - - -class Arguments: - def __init__(self, *args: Any, **kargs: Any): - self.args = args - self.kwargs = kargs - - def __call__(self, function: Callable): - return function(*self.args, **self.kwargs) - - -test_layer_cases = [ - ( - Arguments("b c_in h w -> w c_out h b", "c_in c_out", bias_shape=None, c_out=13, c_in=12), - (2, 12, 3, 4), - (4, 13, 3, 2), - ), - ( - Arguments("b c_in h w -> w c_out h b", "c_in c_out", bias_shape="c_out", c_out=13, c_in=12), - (2, 12, 3, 4), - (4, 13, 3, 2), - ), - ( - Arguments("b c_in h w -> w c_in h b", "", bias_shape=None, c_in=12), - (2, 12, 3, 4), - (4, 12, 3, 2), - ), - ( - Arguments("b c_in h w -> b c_out", "c_in h w c_out", bias_shape=None, c_in=12, h=3, w=4, c_out=5), - (2, 12, 3, 4), - (2, 5), - ), - ( - Arguments("b t head c_in -> b t head c_out", "head c_in c_out", bias_shape=None, head=4, c_in=5, c_out=6), - (2, 3, 4, 5), - (2, 3, 4, 6), - ), -] - - -# Each of the form: -# (Arguments, true_einsum_pattern, in_shapes, out_shape) -test_functional_cases = [ - ( - # Basic: - "b c h w, b w -> b h", - "abcd,ad->ac", - ((2, 3, 4, 5), (2, 5)), - (2, 4), - ), - ( - # Three tensors: - "b c h w, b w, b c -> b h", - "abcd,ad,ab->ac", - ((2, 3, 40, 5), (2, 5), (2, 3)), - (2, 40), - ), - ( - # Ellipsis, and full names: - "... one two three, three four five -> ... two five", - "...abc,cde->...be", - ((32, 5, 2, 3, 4), (4, 5, 6)), - (32, 5, 3, 6), - ), - ( - # Ellipsis at the end: - "one two three ..., three four five -> two five ...", - "abc...,cde->be...", - ((2, 3, 4, 32, 5), (4, 5, 6)), - (3, 6, 32, 5), - ), - ( - # Ellipsis on multiple tensors: - "... one two three, ... three four five -> ... two five", - "...abc,...cde->...be", - ((32, 5, 2, 3, 4), (32, 5, 4, 5, 6)), - (32, 5, 3, 6), - ), - ( - # One tensor, and underscores: - "first_tensor second_tensor -> first_tensor", - "ab->a", - ((5, 4),), - (5,), - ), - ( - # Trace (repeated index) - "i i -> ", - "aa->", - ((5, 5),), - (), - ), - ( - # Too many spaces in string: - " one two , three four->two four ", - "ab,cd->bd", - ((2, 3), (4, 5)), - (3, 5), - ), - # The following tests were inspired by numpy's einsum tests - # https://github.com/numpy/numpy/blob/v1.23.0/numpy/core/tests/test_einsum.py - ( - # Trace with other indices - "i middle i -> middle", - "aba->b", - ((5, 10, 5),), - (10,), - ), - ( - # Ellipsis in the middle: - "i ... i -> ...", - "a...a->...", - ((5, 3, 2, 1, 4, 5),), - (3, 2, 1, 4), - ), - ( - # Product of first and last axes: - "i ... i -> i ...", - "a...a->a...", - ((5, 3, 2, 1, 4, 5),), - (5, 3, 2, 1, 4), - ), - ( - # Triple diagonal - "one one one -> one", - "aaa->a", - ((5, 5, 5),), - (5,), - ), - ( - # Axis swap: - "i j k -> j i k", - "abc->bac", - ((1, 2, 3),), - (2, 1, 3), - ), - ( - # Identity: - "... -> ...", - "...->...", - ((5, 4, 3, 2, 1),), - (5, 4, 3, 2, 1), - ), - ( - # Elementwise product of three tensors - "..., ..., ... -> ...", - "...,...,...->...", - ((3, 2), (3, 2), (3, 2)), - (3, 2), - ), - ( - # Basic summation: - "index ->", - "a->", - ((10,)), - (()), - ), -] - - -@pytest.mark.kernels_ci -def test_layer(): - for backend in collect_test_backends(layers=True, symbolic=False): - if backend.framework_name in ["tensorflow", "torch", "oneflow", "paddle"]: - layer_type = backend.layers().EinMix - for args, in_shape, out_shape in test_layer_cases: - layer = args(layer_type) - print("Running", layer.einsum_pattern, "for", backend.framework_name) - input = np.random.uniform(size=in_shape).astype("float32") - input_framework = backend.from_numpy(input) - output_framework = layer(input_framework) - output = backend.to_numpy(output_framework) - assert output.shape == out_shape - - -valid_backends_functional = [ - "tensorflow", - "torch", - "jax", - "numpy", - "oneflow", - "cupy", - "tensorflow.keras", - "paddle", - "pytensor", -] - - -@pytest.mark.kernels_ci -def test_functional(): - # Functional tests: - backends = filter(lambda x: x.framework_name in valid_backends_functional, collect_test_backends()) - for backend in backends: - for einops_pattern, true_pattern, in_shapes, out_shape in test_functional_cases: - print(f"Running '{einops_pattern}' for {backend.framework_name}") - - # Create pattern: - predicted_pattern = _compactify_pattern_for_einsum(einops_pattern) - assert predicted_pattern == true_pattern - - # Generate example data: - rstate = np.random.RandomState(0) - in_arrays = [rstate.uniform(size=shape).astype("float32") for shape in in_shapes] - in_arrays_framework = [backend.from_numpy(array) for array in in_arrays] - - # Loop over whether we call it manually with the backend, - # or whether we use `einops.einsum`. - for do_manual_call in [True, False]: - # Actually run einsum: - if do_manual_call: - out_array = backend.einsum(predicted_pattern, *in_arrays_framework) - else: - out_array = einsum(*in_arrays_framework, einops_pattern) - - # Check shape: - if tuple(out_array.shape) != out_shape: - raise ValueError(f"Expected output shape {out_shape} but got {out_array.shape}") - - # Check values: - true_out_array = np.einsum(true_pattern, *in_arrays) - predicted_out_array = backend.to_numpy(out_array) - np.testing.assert_array_almost_equal(predicted_out_array, true_out_array, decimal=5) - - -@pytest.mark.kernels_ci -def test_functional_symbolic(): - backends = filter( - lambda x: x.framework_name in valid_backends_functional, collect_test_backends(symbolic=True, layers=False) - ) - for backend in backends: - for einops_pattern, true_pattern, in_shapes, out_shape in test_functional_cases: - print(f"Running '{einops_pattern}' for symbolic {backend.framework_name}") - # Create pattern: - predicted_pattern = _compactify_pattern_for_einsum(einops_pattern) - assert predicted_pattern == true_pattern - - rstate = np.random.RandomState(0) - in_syms = [backend.create_symbol(in_shape) for in_shape in in_shapes] - in_data = [rstate.uniform(size=in_shape).astype("float32") for in_shape in in_shapes] - - expected_out_data = np.einsum(true_pattern, *in_data) - - for do_manual_call in [True, False]: - if do_manual_call: - predicted_out_symbol = backend.einsum(predicted_pattern, *in_syms) - else: - predicted_out_symbol = einsum(*in_syms, einops_pattern) - - predicted_out_data = backend.eval_symbol( - predicted_out_symbol, - list(zip(in_syms, in_data)), - ) - if predicted_out_data.shape != out_shape: - raise ValueError(f"Expected output shape {out_shape} but got {predicted_out_data.shape}") - np.testing.assert_array_almost_equal(predicted_out_data, expected_out_data, decimal=5) - - -@pytest.mark.kernels_ci -def test_functional_errors(): - # Specific backend does not matter, as errors are raised - # during the pattern creation. - - rstate = np.random.RandomState(0) - - def create_tensor(*shape): - return rstate.uniform(size=shape).astype("float32") - - # raise NotImplementedError("Singleton () axes are not yet supported in einsum.") - with pytest.raises(NotImplementedError, match="^Singleton"): - einsum( - create_tensor(5, 1), - "i () -> i", - ) - - # raise NotImplementedError("Shape rearrangement is not yet supported in einsum.") - with pytest.raises(NotImplementedError, match="^Shape rearrangement"): - einsum( - create_tensor(5, 1), - "a b -> (a b)", - ) - - with pytest.raises(NotImplementedError, match="^Shape rearrangement"): - einsum( - create_tensor(10, 1), - "(a b) -> a b", - ) - - # raise RuntimeError("Encountered empty axis name in einsum.") - # raise RuntimeError("Axis name in einsum must be a string.") - # ^ Not tested, these are just a failsafe in case an unexpected error occurs. - - # raise NotImplementedError("Anonymous axes are not yet supported in einsum.") - with pytest.raises(NotImplementedError, match="^Anonymous axes"): - einsum( - create_tensor(5, 1), - "i 2 -> i", - ) - - # ParsedExpression error: - with pytest.raises(EinopsError, match="^Invalid axis identifier"): - einsum( - create_tensor(5, 1), - "i 2j -> i", - ) - - # raise ValueError("Einsum pattern must contain '->'.") - with pytest.raises(ValueError, match="^Einsum pattern"): - einsum( - create_tensor(5, 3, 2), - "i j k", - ) - - # raise RuntimeError("Too many axes in einsum.") - with pytest.raises(RuntimeError, match="^Too many axes"): - einsum( - create_tensor(1), - " ".join(string.ascii_letters) + " extra ->", - ) - - # raise RuntimeError("Unknown axis on right side of einsum.") - with pytest.raises(RuntimeError, match="^Unknown axis"): - einsum( - create_tensor(5, 1), - "i j -> k", - ) - - # raise ValueError( - # "The last argument passed to `einops.einsum` must be a string," - # " representing the einsum pattern." - # ) - with pytest.raises(ValueError, match="^The last argument"): - einsum( - "i j k -> i", - create_tensor(5, 4, 3), - ) - - # raise ValueError( - # "`einops.einsum` takes at minimum two arguments: the tensors," - # " followed by the pattern." - # ) - with pytest.raises(ValueError, match="^`einops.einsum` takes"): - einsum( - "i j k -> i", - ) - with pytest.raises(ValueError, match="^`einops.einsum` takes"): - einsum( - create_tensor(5, 1), - ) - - # TODO: Include check for giving normal einsum pattern rather than einops. diff --git a/.ports/einops/overlay/tests/test_examples.py b/.ports/einops/overlay/tests/test_examples.py deleted file mode 100644 index 4a66431bc..000000000 --- a/.ports/einops/overlay/tests/test_examples.py +++ /dev/null @@ -1,303 +0,0 @@ -import numpy -import pytest - -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -rearrange = einops.rearrange -parse_shape = einops.parse_shape -reduce = einops.reduce -from . import is_backend_tested -from .test_ops import imp_op_backends - - -@pytest.mark.kernels_ci -def test_rearrange_examples(): - def test1(x): - # transpose - y = rearrange(x, "b c h w -> b h w c") - assert tuple(y.shape) == (10, 30, 40, 20) - return y - - def test2(x): - # view / reshape - y = rearrange(x, "b c h w -> b (c h w)") - assert tuple(y.shape) == (10, 20 * 30 * 40) - return y - - def test3(x): - # depth-to-space - y = rearrange(x, "b (c h1 w1) h w -> b c (h h1) (w w1)", h1=2, w1=2) - assert tuple(y.shape) == (10, 5, 30 * 2, 40 * 2) - return y - - def test4(x): - # space-to-depth - y = rearrange(x, "b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=2, w1=2) - assert tuple(y.shape) == (10, 20 * 4, 30 // 2, 40 // 2) - return y - - def test5(x): - # simple transposition - y = rearrange(x, "b1 sound b2 letter -> b1 b2 sound letter") - assert tuple(y.shape) == (10, 30, 20, 40) - return y - - def test6(x): - # parsing parameters - t = rearrange(x, "b c h w -> (b h w) c") - t = t[:, ::2] # replacement for dot-product, just changes size of second axis - assert tuple(t.shape) == (10 * 30 * 40, 10) - - y = rearrange(t, "(b h w) c2 -> b c2 h w", **parse_shape(x, "b _ h w")) - assert tuple(y.shape) == (10, 10, 30, 40) - return y - - def test7(x): - # split of embedding into groups - y1, y2 = rearrange(x, "b (c g) h w -> g b c h w", g=2) - assert tuple(y1.shape) == (10, 10, 30, 40) - assert tuple(y2.shape) == (10, 10, 30, 40) - return y1 + y2 # only one tensor is expected in output - - def test8(x): - # max-pooling - y = reduce(x, "b c (h h1) (w w1) -> b c h w", reduction="max", h1=2, w1=2) - assert tuple(y.shape) == (10, 20, 30 // 2, 40 // 2) - return y - - def test9(x): - # squeeze - unsqueeze - y = reduce(x, "b c h w -> b c () ()", reduction="max") - assert tuple(y.shape) == (10, 20, 1, 1) - y = rearrange(y, "b c () () -> c b") - assert tuple(y.shape) == (20, 10) - return y - - def test10(x): - # stack - tensors = list(x + 0) # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185 - tensors = rearrange(tensors, "b c h w -> b h w c") - assert tuple(tensors.shape) == (10, 30, 40, 20) - return tensors - - def test11(x): - # concatenate - tensors = list(x + 0) # 0 is needed https://github.com/tensorflow/tensorflow/issues/23185 - tensors = rearrange(tensors, "b c h w -> h (b w) c") - assert tuple(tensors.shape) == (30, 10 * 40, 20) - return tensors - - def shufflenet(x, convolve, c1, c2): - # shufflenet reordering example - x = convolve(x) - x = rearrange(x, "b (c1 c2) h w-> b (c2 c1) h w", c1=c1, c2=c2) - x = convolve(x) - return x - - def convolve_strided_1d(x, stride, usual_convolution): - x = rearrange(x, "b c t1 t2 -> b c (t1 t2)") # reduce dimensionality - x = rearrange(x, "b c (t stride) -> (stride b) c t", stride=stride) - x = usual_convolution(x) - x = rearrange(x, "(stride b) c t -> b c (t stride)", stride=stride) - return x - - def convolve_strided_2d(x, h_stride, w_stride, usual_convolution): - x = rearrange(x, "b c (h hs) (w ws) -> (hs ws b) c h w", hs=h_stride, ws=w_stride) - x = usual_convolution(x) - x = rearrange(x, "(hs ws b) c h w -> b c (h hs) (w ws)", hs=h_stride, ws=w_stride) - return x - - def unet_like_1d(x, usual_convolution): - # u-net like steps for increasing / reducing dimensionality - x = rearrange(x, "b c t1 t2 -> b c (t1 t2)") # reduce dimensionality - y = rearrange(x, "b c (t dt) -> b (dt c) t", dt=2) - y = usual_convolution(y) - x = x + rearrange(y, "b (dt c) t -> b c (t dt)", dt=2) - return x - - # mock for convolution (works for all backends) - def convolve_mock(x): - return x - - tests = [ - test1, - test2, - test3, - test4, - test5, - test6, - test7, - test8, - test9, - test10, - test11, - lambda x: shufflenet(x, convolve=convolve_mock, c1=4, c2=5), - lambda x: convolve_strided_1d(x, stride=2, usual_convolution=convolve_mock), - lambda x: convolve_strided_2d(x, h_stride=2, w_stride=2, usual_convolution=convolve_mock), - lambda x: unet_like_1d(x, usual_convolution=convolve_mock), - ] - - for backend in imp_op_backends: - print("testing source_examples for ", backend.framework_name) - for test in tests: - x = numpy.arange(10 * 20 * 30 * 40).reshape([10, 20, 30, 40]) - result1 = test(x) - result2 = backend.to_numpy(test(backend.from_numpy(x))) - assert numpy.array_equal(result1, result2) - - # now with strides - x = numpy.arange(10 * 2 * 20 * 3 * 30 * 1 * 40).reshape([10 * 2, 20 * 3, 30 * 1, 40 * 1]) - # known torch bug - torch doesn't support negative steps - last_step = -1 if (backend.framework_name != "torch" and backend.framework_name != "oneflow") else 1 - indexing_expression = numpy.index_exp[::2, ::3, ::1, ::last_step] - result1 = test(x[indexing_expression]) - result2 = backend.to_numpy(test(backend.from_numpy(x)[indexing_expression])) - assert numpy.array_equal(result1, result2) - - -def tensor_train_example_numpy(): - # kept here just for a collection, only tested for numpy - # https://arxiv.org/pdf/1509.06569.pdf, (5) - x = numpy.ones([3, 4, 5, 6]) - rank = 4 - if numpy.__version__ < "1.15.0": - # numpy.einsum fails here, skip test - return - # creating appropriate Gs - Gs = [numpy.ones([d, d, rank, rank]) for d in x.shape] - Gs[0] = Gs[0][:, :, :1, :] - Gs[-1] = Gs[-1][:, :, :, :1] - - # einsum way - y = x.reshape((1,) + x.shape) - for G in Gs: - # taking partial results left-to-right - # y = numpy.einsum('i j alpha beta, alpha i ... -> beta ... j', G, y) - y = numpy.einsum("i j a b, a i ... -> b ... j", G, y) - y1 = y.reshape(-1) - - # alternative way - y = x.reshape(-1) - for G in Gs: - i, j, alpha, beta = G.shape - y = rearrange(y, "(i rest alpha) -> rest (alpha i)", alpha=alpha, i=i) - y = y @ rearrange(G, "i j alpha beta -> (alpha i) (j beta)") - y = rearrange(y, "rest (beta j) -> (beta rest j)", beta=beta, j=j) - y2 = y - assert numpy.allclose(y1, y2) - - # yet another way - y = x - for G in Gs: - i, j, alpha, beta = G.shape - y = rearrange(y, "i ... (j alpha) -> ... j (alpha i)", alpha=alpha, i=i) - y = y @ rearrange(G, "i j alpha beta -> (alpha i) (j beta)") - y3 = y.reshape(-1) - assert numpy.allclose(y1, y3) - - -@pytest.mark.kernels_ci -def test_pytorch_yolo_fragment(): - if not is_backend_tested("torch"): - pytest.skip() - - import torch - - def old_way(input, num_classes, num_anchors, anchors, stride_h, stride_w): - # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L28-L44 - bs = input.size(0) - in_h = input.size(2) - in_w = input.size(3) - scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors] - - prediction = input.view(bs, num_anchors, 5 + num_classes, in_h, in_w).permute(0, 1, 3, 4, 2).contiguous() - # Get outputs - x = torch.sigmoid(prediction[..., 0]) # Center x - y = torch.sigmoid(prediction[..., 1]) # Center y - w = prediction[..., 2] # Width - h = prediction[..., 3] # Height - conf = torch.sigmoid(prediction[..., 4]) # Conf - pred_cls = torch.sigmoid(prediction[..., 5:]) # Cls pred. - - # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L70-L92 - FloatTensor = torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor - LongTensor = torch.cuda.LongTensor if x.is_cuda else torch.LongTensor - # Calculate offsets for each grid - grid_x = ( - torch.linspace(0, in_w - 1, in_w) - .repeat(in_w, 1) - .repeat(bs * num_anchors, 1, 1) - .view(x.shape) - .type(FloatTensor) - ) - grid_y = ( - torch.linspace(0, in_h - 1, in_h) - .repeat(in_h, 1) - .t() - .repeat(bs * num_anchors, 1, 1) - .view(y.shape) - .type(FloatTensor) - ) - # Calculate anchor w, h - anchor_w = FloatTensor(scaled_anchors).index_select(1, LongTensor([0])) - anchor_h = FloatTensor(scaled_anchors).index_select(1, LongTensor([1])) - anchor_w = anchor_w.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(w.shape) - anchor_h = anchor_h.repeat(bs, 1).repeat(1, 1, in_h * in_w).view(h.shape) - # Add offset and scale with anchors - pred_boxes = FloatTensor(prediction[..., :4].shape) - pred_boxes[..., 0] = x.data + grid_x - pred_boxes[..., 1] = y.data + grid_y - pred_boxes[..., 2] = torch.exp(w.data) * anchor_w - pred_boxes[..., 3] = torch.exp(h.data) * anchor_h - # Results - _scale = torch.Tensor([stride_w, stride_h] * 2).type(FloatTensor) - output = torch.cat( - (pred_boxes.view(bs, -1, 4) * _scale, conf.view(bs, -1, 1), pred_cls.view(bs, -1, num_classes)), -1 - ) - return output - - def new_way(input, num_classes, num_anchors, anchors, stride_h, stride_w): - raw_predictions = rearrange(input, " b (anchor prediction) h w -> prediction b anchor h w", anchor=num_anchors) - - anchors = torch.FloatTensor(anchors).to(input.device) - anchor_sizes = rearrange(anchors, "anchor dim -> dim () anchor () ()") - - _, _, _, in_h, in_w = raw_predictions.shape - grid_h = rearrange(torch.arange(in_h).float(), "h -> () () h ()").to(input.device) - grid_w = rearrange(torch.arange(in_w).float(), "w -> () () () w").to(input.device) - - predicted_bboxes = torch.zeros_like(raw_predictions) - predicted_bboxes[0] = (raw_predictions[0].sigmoid() + grid_h) * stride_h # center y - predicted_bboxes[1] = (raw_predictions[1].sigmoid() + grid_w) * stride_w # center x - predicted_bboxes[2:4] = (raw_predictions[2:4].exp()) * anchor_sizes # bbox width and height - predicted_bboxes[4] = raw_predictions[4].sigmoid() # confidence - predicted_bboxes[5:] = raw_predictions[5:].sigmoid() # class predictions - # only to match results of original code, not needed - return rearrange(predicted_bboxes, "prediction b anchor h w -> b anchor h w prediction") - - stride_h = 4 - stride_w = 4 - batch_size = 5 - num_classes = 12 - anchors = [[50, 100], [100, 50], [75, 75]] - num_anchors = len(anchors) - - input = torch.randn([batch_size, num_anchors * (5 + num_classes), 1, 1]) - result1 = old_way( - input=input, - num_anchors=num_anchors, - num_classes=num_classes, - stride_h=stride_h, - stride_w=stride_w, - anchors=anchors, - ) - result2 = new_way( - input=input, - num_anchors=num_anchors, - num_classes=num_classes, - stride_h=stride_h, - stride_w=stride_w, - anchors=anchors, - ) - result1 = result1.reshape(result2.shape) - assert torch.allclose(result1, result2) diff --git a/.ports/einops/overlay/tests/test_layers.py b/.ports/einops/overlay/tests/test_layers.py deleted file mode 100644 index 538b3b8e1..000000000 --- a/.ports/einops/overlay/tests/test_layers.py +++ /dev/null @@ -1,484 +0,0 @@ -import pickle -from collections import namedtuple - -import numpy -import pytest - -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -rearrange = einops.rearrange -reduce = einops.reduce -EinopsError = einops.EinopsError -from . import collect_test_backends, is_backend_tested -from . import FLOAT_REDUCTIONS as REDUCTIONS - -__author__ = "Alex Rogozhnikov" - -testcase = namedtuple("testcase", ["pattern", "axes_lengths", "input_shape", "wrong_shapes"]) - -rearrangement_patterns = [ - testcase( - "b c h w -> b (c h w)", - dict(c=20), - (10, 20, 30, 40), - [(), (10,), (10, 10, 10), (10, 21, 30, 40), [1, 20, 1, 1, 1]], - ), - testcase( - "b c (h1 h2) (w1 w2) -> b (c h2 w2) h1 w1", - dict(h2=2, w2=2), - (10, 20, 30, 40), - [(), (1, 1, 1, 1), (1, 10, 3), ()], - ), - testcase( - "b ... c -> c b ...", - dict(b=10), - (10, 20, 30), - [(), (10,), (5, 10)], - ), -] - - -@pytest.mark.kernels_ci -def test_rearrange_imperative(): - for backend in collect_test_backends(symbolic=False, layers=True): - print("Test layer for ", backend.framework_name) - - for pattern, axes_lengths, input_shape, wrong_shapes in rearrangement_patterns: - x = numpy.arange(numpy.prod(input_shape), dtype="float32").reshape(input_shape) - result_numpy = rearrange(x, pattern, **axes_lengths) - layer = backend.layers().Rearrange(pattern, **axes_lengths) - for shape in wrong_shapes: - try: - layer(backend.from_numpy(numpy.zeros(shape, dtype="float32"))) - except BaseException: - pass - else: - raise AssertionError("Failure expected") - - # simple pickling / unpickling - layer2 = pickle.loads(pickle.dumps(layer)) - result1 = backend.to_numpy(layer(backend.from_numpy(x))) - result2 = backend.to_numpy(layer2(backend.from_numpy(x))) - assert numpy.allclose(result_numpy, result1) - assert numpy.allclose(result1, result2) - - just_sum = backend.layers().Reduce("...->", reduction="sum") - - variable = backend.from_numpy(x) - result = just_sum(layer(variable)) - - result.backward() - assert numpy.allclose(backend.to_numpy(variable.grad), 1) - - -@pytest.mark.kernels_ci -def test_rearrange_symbolic(): - for backend in collect_test_backends(symbolic=True, layers=True): - print("Test layer for ", backend.framework_name) - - for pattern, axes_lengths, input_shape, wrong_shapes in rearrangement_patterns: - x = numpy.arange(numpy.prod(input_shape), dtype="float32").reshape(input_shape) - result_numpy = rearrange(x, pattern, **axes_lengths) - layer = backend.layers().Rearrange(pattern, **axes_lengths) - input_shape_of_nones = [None] * len(input_shape) - shapes = [input_shape, input_shape_of_nones] - - for shape in shapes: - symbol = backend.create_symbol(shape) - eval_inputs = [(symbol, x)] - - result_symbol1 = layer(symbol) - result1 = backend.eval_symbol(result_symbol1, eval_inputs) - assert numpy.allclose(result_numpy, result1) - - layer2 = pickle.loads(pickle.dumps(layer)) - result_symbol2 = layer2(symbol) - result2 = backend.eval_symbol(result_symbol2, eval_inputs) - assert numpy.allclose(result1, result2) - - # now testing back-propagation - just_sum = backend.layers().Reduce("...->", reduction="sum") - - result_sum1 = backend.eval_symbol(just_sum(result_symbol1), eval_inputs) - result_sum2 = numpy.sum(x) - - assert numpy.allclose(result_sum1, result_sum2) - - -reduction_patterns = rearrangement_patterns + [ - testcase("b c h w -> b ()", dict(b=10), (10, 20, 30, 40), [(10,), (10, 20, 30)]), - testcase("b c (h1 h2) (w1 w2) -> b c h1 w1", dict(h1=15, h2=2, w2=2), (10, 20, 30, 40), [(10, 20, 31, 40)]), - testcase("b ... c -> b", dict(b=10), (10, 20, 30, 40), [(10,), (11, 10)]), -] - - -@pytest.mark.kernels_ci -def test_reduce_imperative(): - for backend in collect_test_backends(symbolic=False, layers=True): - print("Test layer for ", backend.framework_name) - for reduction in REDUCTIONS: - for pattern, axes_lengths, input_shape, wrong_shapes in reduction_patterns: - print(backend, reduction, pattern, axes_lengths, input_shape, wrong_shapes) - x = numpy.arange(1, 1 + numpy.prod(input_shape), dtype="float32").reshape(input_shape) - x /= x.mean() - result_numpy = reduce(x, pattern, reduction, **axes_lengths) - layer = backend.layers().Reduce(pattern, reduction, **axes_lengths) - for shape in wrong_shapes: - try: - layer(backend.from_numpy(numpy.zeros(shape, dtype="float32"))) - except BaseException: - pass - else: - raise AssertionError("Failure expected") - - # simple pickling / unpickling - layer2 = pickle.loads(pickle.dumps(layer)) - result1 = backend.to_numpy(layer(backend.from_numpy(x))) - result2 = backend.to_numpy(layer2(backend.from_numpy(x))) - assert numpy.allclose(result_numpy, result1) - assert numpy.allclose(result1, result2) - - just_sum = backend.layers().Reduce("...->", reduction="sum") - - variable = backend.from_numpy(x) - result = just_sum(layer(variable)) - - result.backward() - grad = backend.to_numpy(variable.grad) - if reduction == "sum": - assert numpy.allclose(grad, 1) - if reduction == "mean": - assert numpy.allclose(grad, grad.min()) - if reduction in ["max", "min"]: - assert numpy.all(numpy.isin(grad, [0, 1])) - assert numpy.sum(grad) > 0.5 - - -@pytest.mark.kernels_ci -def test_reduce_symbolic(): - for backend in collect_test_backends(symbolic=True, layers=True): - print("Test layer for ", backend.framework_name) - for reduction in REDUCTIONS: - for pattern, axes_lengths, input_shape, wrong_shapes in reduction_patterns: - x = numpy.arange(1, 1 + numpy.prod(input_shape), dtype="float32").reshape(input_shape) - x /= x.mean() - result_numpy = reduce(x, pattern, reduction, **axes_lengths) - layer = backend.layers().Reduce(pattern, reduction, **axes_lengths) - input_shape_of_nones = [None] * len(input_shape) - shapes = [input_shape, input_shape_of_nones] - - for shape in shapes: - symbol = backend.create_symbol(shape) - eval_inputs = [(symbol, x)] - - result_symbol1 = layer(symbol) - result1 = backend.eval_symbol(result_symbol1, eval_inputs) - assert numpy.allclose(result_numpy, result1) - - layer2 = pickle.loads(pickle.dumps(layer)) - result_symbol2 = layer2(symbol) - result2 = backend.eval_symbol(result_symbol2, eval_inputs) - assert numpy.allclose(result1, result2) - - -def create_torch_model(use_reduce=False, add_scripted_layer=False): - if not is_backend_tested("torch"): - pytest.skip() - else: - from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU - Rearrange, Reduce, EinMix = einops.layers.torch.Rearrange, einops.layers.torch.Reduce, einops.layers.torch.EinMix - import torch.jit - - return Sequential( - Conv2d(3, 6, kernel_size=(5, 5)), - Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2) if use_reduce else MaxPool2d(kernel_size=2), - Conv2d(6, 16, kernel_size=(5, 5)), - Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2), - torch.jit.script(Rearrange("b c h w -> b (c h w)")) - if add_scripted_layer - else Rearrange("b c h w -> b (c h w)"), - Linear(16 * 5 * 5, 120), - ReLU(), - Linear(120, 84), - ReLU(), - EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84), - EinMix("(b c2) -> b c3", weight_shape="c2 c3", bias_shape="c3", c2=84, c3=84), - Linear(84, 10), - ) - - -@pytest.mark.kernels_ci -def test_torch_layer(): - if not is_backend_tested("torch"): - pytest.skip() - else: - # checked that torch present - import torch - import torch.jit - - model1 = create_torch_model(use_reduce=True) - model2 = create_torch_model(use_reduce=False) - input = torch.randn([10, 3, 32, 32]) - # random models have different predictions - assert not torch.allclose(model1(input), model2(input)) - model2.load_state_dict(pickle.loads(pickle.dumps(model1.state_dict()))) - assert torch.allclose(model1(input), model2(input)) - - # tracing (freezing) - model3 = torch.jit.trace(model2, example_inputs=input) - torch.testing.assert_close(model1(input), model3(input), atol=1e-3, rtol=1e-3) - torch.testing.assert_close(model1(input + 1), model3(input + 1), atol=1e-3, rtol=1e-3) - - model4 = torch.jit.trace(model2, example_inputs=input) - torch.testing.assert_close(model1(input), model4(input), atol=1e-3, rtol=1e-3) - torch.testing.assert_close(model1(input + 1), model4(input + 1), atol=1e-3, rtol=1e-3) - - -@pytest.mark.kernels_ci -def test_torch_layers_scripting(): - if not is_backend_tested("torch"): - pytest.skip() - else: - import torch - - for script_layer in [False, True]: - model1 = create_torch_model(use_reduce=True, add_scripted_layer=script_layer) - model2 = torch.jit.script(model1) - input = torch.randn([10, 3, 32, 32]) - - torch.testing.assert_close(model1(input), model2(input), atol=1e-3, rtol=1e-3) - - -@pytest.mark.kernels_ci -def test_keras_layer(): - if not is_backend_tested("tensorflow"): - pytest.skip() - else: - import tensorflow as tf - - if tf.__version__ < "2.16.": - # current implementation of layers follows new TF interface - pytest.skip() - from tensorflow.keras.models import Sequential - from tensorflow.keras.layers import Conv2D as Conv2d, Dense as Linear, ReLU - Rearrange, Reduce, EinMix, keras_custom_objects = einops.layers.keras.Rearrange, einops.layers.keras.Reduce, einops.layers.keras.EinMix, einops.layers.keras.keras_custom_objects - - def create_keras_model(): - return Sequential( - [ - Conv2d(6, kernel_size=5, input_shape=[32, 32, 3]), - Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2), - Conv2d(16, kernel_size=5), - Reduce("b c (h h2) (w w2) -> b c h w", "max", h2=2, w2=2), - Rearrange("b c h w -> b (c h w)"), - Linear(120), - ReLU(), - Linear(84), - ReLU(), - EinMix("b c1 -> (b c2)", weight_shape="c1 c2", bias_shape="c2", c1=84, c2=84), - EinMix("(b c2) -> b c3", weight_shape="c2 c3", bias_shape="c3", c2=84, c3=84), - Linear(10), - ] - ) - - model1 = create_keras_model() - model2 = create_keras_model() - - input = numpy.random.normal(size=[10, 32, 32, 3]).astype("float32") - # two randomly init models should provide different outputs - assert not numpy.allclose(model1.predict_on_batch(input), model2.predict_on_batch(input)) - - # get some temp filename - tmp_model_filename = "/tmp/einops_tf_model.h5" - # save arch + weights - print("temp_path_keras1", tmp_model_filename) - tf.keras.models.save_model(model1, tmp_model_filename) - model3 = tf.keras.models.load_model(tmp_model_filename, custom_objects=keras_custom_objects) - - numpy.testing.assert_allclose(model1.predict_on_batch(input), model3.predict_on_batch(input)) - - weight_filename = "/tmp/einops_tf_model.weights.h5" - # save arch as json - model4 = tf.keras.models.model_from_json(model1.to_json(), custom_objects=keras_custom_objects) - model1.save_weights(weight_filename) - model4.load_weights(weight_filename) - model2.load_weights(weight_filename) - # check that differently-inialized model receives same weights - numpy.testing.assert_allclose(model1.predict_on_batch(input), model2.predict_on_batch(input)) - # ulimate test - # save-load architecture, and then load weights - should return same result - numpy.testing.assert_allclose(model1.predict_on_batch(input), model4.predict_on_batch(input)) - - -@pytest.mark.kernels_ci -def test_flax_layers(): - """ - One-off simple tests for Flax layers. - Unfortunately, Flax layers have a different interface from other layers. - """ - if not is_backend_tested("jax"): - pytest.skip() - else: - import jax - import jax.numpy as jnp - - import flax - from flax import linen as nn - EinMix, Reduce, Rearrange = einops.layers.flax.EinMix, einops.layers.flax.Reduce, einops.layers.flax.Rearrange - - class NN(nn.Module): - @nn.compact - def __call__(self, x): - x = EinMix( - "b (h h2) (w w2) c -> b h w c_out", "h2 w2 c c_out", "c_out", sizes=dict(h2=2, w2=3, c=4, c_out=5) - )(x) - x = Rearrange("b h w c -> b (w h c)", sizes=dict(c=5))(x) - x = Reduce("b hwc -> b", "mean", dict(hwc=2 * 3 * 5))(x) - return x - - model = NN() - fixed_input = jnp.ones([10, 2 * 2, 3 * 3, 4]) - params = model.init(jax.random.PRNGKey(0), fixed_input) - - def eval_at_point(params): - return jnp.linalg.norm(model.apply(params, fixed_input)) - - vandg = jax.value_and_grad(eval_at_point) - value0 = eval_at_point(params) - value1, grad1 = vandg(params) - assert jnp.allclose(value0, value1) - - params2 = jax.tree_map(lambda x1, x2: x1 - x2 * 0.001, params, grad1) - - value2 = eval_at_point(params2) - assert value0 >= value2, (value0, value2) - - # check serialization - fbytes = flax.serialization.to_bytes(params) - _loaded = flax.serialization.from_bytes(params, fbytes) - - -@pytest.mark.kernels_ci -def test_einmix_decomposition(): - """ - Testing that einmix correctly decomposes into smaller transformations. - """ - _EinmixDebugger = einops.layers._einmix._EinmixDebugger - - mixin1 = _EinmixDebugger( - "a b c d e -> e d c b a", - weight_shape="d a b", - d=2, a=3, b=5, - ) # fmt: off - assert mixin1.pre_reshape_pattern is None - assert mixin1.post_reshape_pattern is None - assert mixin1.einsum_pattern == "abcde,dab->edcba" - assert mixin1.saved_weight_shape == [2, 3, 5] - assert mixin1.saved_bias_shape is None - - mixin2 = _EinmixDebugger( - "a b c d e -> e d c b a", - weight_shape="d a b", - bias_shape="a b c d e", - a=1, b=2, c=3, d=4, e=5, - ) # fmt: off - assert mixin2.pre_reshape_pattern is None - assert mixin2.post_reshape_pattern is None - assert mixin2.einsum_pattern == "abcde,dab->edcba" - assert mixin2.saved_weight_shape == [4, 1, 2] - assert mixin2.saved_bias_shape == [5, 4, 3, 2, 1] - - mixin3 = _EinmixDebugger( - "... -> ...", - weight_shape="", - bias_shape="", - ) # fmt: off - assert mixin3.pre_reshape_pattern is None - assert mixin3.post_reshape_pattern is None - assert mixin3.einsum_pattern == "...,->..." - assert mixin3.saved_weight_shape == [] - assert mixin3.saved_bias_shape == [] - - mixin4 = _EinmixDebugger( - "b a ... -> b c ...", - weight_shape="b a c", - a=1, b=2, c=3, - ) # fmt: off - assert mixin4.pre_reshape_pattern is None - assert mixin4.post_reshape_pattern is None - assert mixin4.einsum_pattern == "ba...,bac->bc..." - assert mixin4.saved_weight_shape == [2, 1, 3] - assert mixin4.saved_bias_shape is None - - mixin5 = _EinmixDebugger( - "(b a) ... -> b c (...)", - weight_shape="b a c", - a=1, b=2, c=3, - ) # fmt: off - assert mixin5.pre_reshape_pattern == "(b a) ... -> b a ..." - assert mixin5.pre_reshape_lengths == dict(a=1, b=2) - assert mixin5.post_reshape_pattern == "b c ... -> b c (...)" - assert mixin5.einsum_pattern == "ba...,bac->bc..." - assert mixin5.saved_weight_shape == [2, 1, 3] - assert mixin5.saved_bias_shape is None - - mixin6 = _EinmixDebugger( - "b ... (a c) -> b ... (a d)", - weight_shape="c d", - bias_shape="a d", - a=1, c=3, d=4, - ) # fmt: off - assert mixin6.pre_reshape_pattern == "b ... (a c) -> b ... a c" - assert mixin6.pre_reshape_lengths == dict(a=1, c=3) - assert mixin6.post_reshape_pattern == "b ... a d -> b ... (a d)" - assert mixin6.einsum_pattern == "b...ac,cd->b...ad" - assert mixin6.saved_weight_shape == [3, 4] - assert mixin6.saved_bias_shape == [1, 1, 4] # (b) a d, ellipsis does not participate - - mixin7 = _EinmixDebugger( - "a ... (b c) -> a (... d b)", - weight_shape="c d b", - bias_shape="d b", - b=2, c=3, d=4, - ) # fmt: off - assert mixin7.pre_reshape_pattern == "a ... (b c) -> a ... b c" - assert mixin7.pre_reshape_lengths == dict(b=2, c=3) - assert mixin7.post_reshape_pattern == "a ... d b -> a (... d b)" - assert mixin7.einsum_pattern == "a...bc,cdb->a...db" - assert mixin7.saved_weight_shape == [3, 4, 2] - assert mixin7.saved_bias_shape == [1, 4, 2] # (a) d b, ellipsis does not participate - - -@pytest.mark.kernels_ci -def test_einmix_restrictions(): - """ - Testing different cases - """ - _EinmixDebugger = einops.layers._einmix._EinmixDebugger - - with pytest.raises(EinopsError): - _EinmixDebugger( - "a b c d e -> e d c b a", - weight_shape="d a b", - d=2, a=3, # missing b - ) # fmt: off - - with pytest.raises(EinopsError): - _EinmixDebugger( - "a b c d e -> e d c b a", - weight_shape="w a b", - d=2, a=3, b=1 # missing d - ) # fmt: off - - with pytest.raises(EinopsError): - _EinmixDebugger( - "(...) a -> ... a", - weight_shape="a", a=1, # ellipsis on the left - ) # fmt: off - - with pytest.raises(EinopsError): - _EinmixDebugger( - "(...) a -> a ...", - weight_shape="a", a=1, # ellipsis on the right side after bias axis - bias_shape='a', - ) # fmt: off diff --git a/.ports/einops/overlay/tests/test_ops.py b/.ports/einops/overlay/tests/test_ops.py deleted file mode 100644 index ff668bb50..000000000 --- a/.ports/einops/overlay/tests/test_ops.py +++ /dev/null @@ -1,684 +0,0 @@ -import itertools - -import numpy -import numpy as np -import pytest - -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -rearrange = einops.einops.rearrange -reduce = einops.einops.reduce -repeat = einops.einops.repeat -_enumerate_directions = einops.einops._enumerate_directions -from . import collect_test_backends, is_backend_tested -from . import FLOAT_REDUCTIONS as REDUCTIONS - -imp_op_backends = collect_test_backends(symbolic=False, layers=False) -sym_op_backends = collect_test_backends(symbolic=True, layers=False) - -identity_patterns = [ - "...->...", - "a b c d e-> a b c d e", - "a b c d e ...-> ... a b c d e", - "a b c d e ...-> a ... b c d e", - "... a b c d e -> ... a b c d e", - "a ... e-> a ... e", - "a ... -> a ... ", - "a ... c d e -> a (...) c d e", -] - -equivalent_rearrange_patterns = [ - ("a b c d e -> (a b) c d e", "a b ... -> (a b) ... "), - ("a b c d e -> a b (c d) e", "... c d e -> ... (c d) e"), - ("a b c d e -> a b c d e", "... -> ... "), - ("a b c d e -> (a b c d e)", "... -> (...)"), - ("a b c d e -> b (c d e) a", "a b ... -> b (...) a"), - ("a b c d e -> b (a c d) e", "a b ... e -> b (a ...) e"), -] - -equivalent_reduction_patterns = [ - ("a b c d e -> ", " ... -> "), - ("a b c d e -> (e a)", "a ... e -> (e a)"), - ("a b c d e -> d (a e)", " a b c d e ... -> d (a e) "), - ("a b c d e -> (a b)", " ... c d e -> (...) "), -] - - -@pytest.mark.kernels_ci -def test_collapsed_ellipsis_errors_out(): - x = numpy.zeros([1, 1, 1, 1, 1]) - rearrange(x, "a b c d ... -> a b c ... d") - with pytest.raises(EinopsError): - rearrange(x, "a b c d (...) -> a b c ... d") - - rearrange(x, "... -> (...)") - with pytest.raises(EinopsError): - rearrange(x, "(...) -> (...)") - - -@pytest.mark.kernels_ci -def test_ellipsis_ops_numpy(): - x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) - for pattern in identity_patterns: - assert numpy.array_equal(x, rearrange(x, pattern)), pattern - - for pattern1, pattern2 in equivalent_rearrange_patterns: - assert numpy.array_equal(rearrange(x, pattern1), rearrange(x, pattern2)) - - for reduction in ["min", "max", "sum"]: - for pattern1, pattern2 in equivalent_reduction_patterns: - assert numpy.array_equal(reduce(x, pattern1, reduction=reduction), reduce(x, pattern2, reduction=reduction)) - - # now just check coincidence with numpy - all_rearrange_patterns = [*identity_patterns] - for pattern_pairs in equivalent_rearrange_patterns: - all_rearrange_patterns.extend(pattern_pairs) - - -def check_op_against_numpy(backend, numpy_input, pattern, axes_lengths, reduction="rearrange", is_symbolic=False): - """ - Helper to test result of operation (rearrange or transpose) against numpy - if reduction == 'rearrange', rearrange op is tested, otherwise reduce - """ - - def operation(x): - if reduction == "rearrange": - return rearrange(x, pattern, **axes_lengths) - else: - return reduce(x, pattern, reduction, **axes_lengths) - - numpy_result = operation(numpy_input) - check_equal = numpy.array_equal - p_none_dimension = 0.5 - if is_symbolic: - symbol_shape = [d if numpy.random.random() >= p_none_dimension else None for d in numpy_input.shape] - symbol = backend.create_symbol(shape=symbol_shape) - result_symbol = operation(symbol) - backend_result = backend.eval_symbol(result_symbol, [(symbol, numpy_input)]) - else: - backend_result = operation(backend.from_numpy(numpy_input)) - backend_result = backend.to_numpy(backend_result) - - check_equal(numpy_result, backend_result) - - -@pytest.mark.kernels_ci -def test_ellipsis_ops_imperative(): - """Checking various patterns against numpy""" - x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) - for is_symbolic in [True, False]: - for backend in collect_test_backends(symbolic=is_symbolic, layers=False): - for pattern in identity_patterns + list(itertools.chain(*equivalent_rearrange_patterns)): - check_op_against_numpy( - backend, x, pattern, axes_lengths={}, reduction="rearrange", is_symbolic=is_symbolic - ) - - for reduction in ["min", "max", "sum"]: - for pattern in itertools.chain(*equivalent_reduction_patterns): - check_op_against_numpy( - backend, x, pattern, axes_lengths={}, reduction=reduction, is_symbolic=is_symbolic - ) - - -@pytest.mark.kernels_ci -def test_rearrange_array_api(): - import numpy as xp - AA = einops.array_api - - if xp.__version__ < "2.0.0": - pytest.skip() - - x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) - for pattern in identity_patterns + list(itertools.chain(*equivalent_rearrange_patterns)): - expected = rearrange(x, pattern) - result = AA.rearrange(xp.from_dlpack(x), pattern) - assert numpy.array_equal(AA.asnumpy(result + 0), expected) - - -@pytest.mark.kernels_ci -def test_reduce_array_api(): - import numpy as xp - AA = einops.array_api - - if xp.__version__ < "2.0.0": - pytest.skip() - - x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) - for pattern in itertools.chain(*equivalent_reduction_patterns): - for reduction in ["min", "max", "sum"]: - expected = reduce(x, pattern, reduction=reduction) - result = AA.reduce(xp.from_dlpack(x), pattern, reduction=reduction) - assert numpy.array_equal(AA.asnumpy(np.asarray(result + 0)), expected) - - -@pytest.mark.kernels_ci -def test_rearrange_consistency_numpy(): - shape = [1, 2, 3, 5, 7, 11] - x = numpy.arange(numpy.prod(shape)).reshape(shape) - for pattern in [ - "a b c d e f -> a b c d e f", - "b a c d e f -> a b d e f c", - "a b c d e f -> f e d c b a", - "a b c d e f -> (f e) d (c b a)", - "a b c d e f -> (f e d c b a)", - ]: - result = rearrange(x, pattern) - assert len(numpy.setdiff1d(x, result)) == 0 - assert result.dtype == x.dtype - - result = rearrange(x, "a b c d e f -> a (b) (c d e) f") - assert numpy.array_equal(x.flatten(), result.flatten()) - - result = rearrange(x, "a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11") - assert numpy.array_equal(x, result) - - result1 = rearrange(x, "a b c d e f -> f e d c b a") - result2 = rearrange(x, "f e d c b a -> a b c d e f") - assert numpy.array_equal(result1, result2) - - result = rearrange(rearrange(x, "a b c d e f -> (f d) c (e b) a"), "(f d) c (e b) a -> a b c d e f", b=2, d=5) - assert numpy.array_equal(x, result) - - sizes = dict(zip("abcdef", shape)) - temp = rearrange(x, "a b c d e f -> (f d) c (e b) a", **sizes) - result = rearrange(temp, "(f d) c (e b) a -> a b c d e f", **sizes) - assert numpy.array_equal(x, result) - - x2 = numpy.arange(2 * 3 * 4).reshape([2, 3, 4]) - result = rearrange(x2, "a b c -> b c a") - assert x2[1, 2, 3] == result[2, 3, 1] - assert x2[0, 1, 2] == result[1, 2, 0] - - -@pytest.mark.kernels_ci -def test_rearrange_permutations_numpy(): - # tests random permutation of axes against two independent numpy ways - for n_axes in range(1, 10): - input = numpy.arange(2**n_axes).reshape([2] * n_axes) - permutation = numpy.random.permutation(n_axes) - left_expression = " ".join("i" + str(axis) for axis in range(n_axes)) - right_expression = " ".join("i" + str(axis) for axis in permutation) - expression = left_expression + " -> " + right_expression - result = rearrange(input, expression) - - for pick in numpy.random.randint(0, 2, [10, n_axes]): - assert input[tuple(pick)] == result[tuple(pick[permutation])] - - for n_axes in range(1, 10): - input = numpy.arange(2**n_axes).reshape([2] * n_axes) - permutation = numpy.random.permutation(n_axes) - left_expression = " ".join("i" + str(axis) for axis in range(n_axes)[::-1]) - right_expression = " ".join("i" + str(axis) for axis in permutation[::-1]) - expression = left_expression + " -> " + right_expression - result = rearrange(input, expression) - assert result.shape == input.shape - expected_result = numpy.zeros_like(input) - for original_axis, result_axis in enumerate(permutation): - expected_result |= ((input >> original_axis) & 1) << result_axis - - assert numpy.array_equal(result, expected_result) - - -@pytest.mark.kernels_ci -def test_reduction_imperatives(): - for backend in imp_op_backends: - print("Reduction tests for ", backend.framework_name) - for reduction in REDUCTIONS: - # slight redundancy for simpler order - numpy version is evaluated multiple times - input = numpy.arange(2 * 3 * 4 * 5 * 6, dtype="int64").reshape([2, 3, 4, 5, 6]) - if reduction in ["mean", "prod"]: - input = input / input.astype("float64").mean() - test_cases = [ - ["a b c d e -> ", {}, getattr(input, reduction)()], - ["a ... -> ", {}, getattr(input, reduction)()], - ["(a1 a2) ... (e1 e2) -> ", dict(a1=1, e2=2), getattr(input, reduction)()], - [ - "a b c d e -> (e c) a", - {}, - getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), - ], - [ - "a ... c d e -> (e c) a", - {}, - getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), - ], - [ - "a b c d e ... -> (e c) a", - {}, - getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), - ], - ["a b c d e -> (e c a)", {}, getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1])], - ["(a a2) ... -> (a2 a) ...", dict(a2=1), input], - ] - for pattern, axes_lengths, expected_result in test_cases: - result = reduce(backend.from_numpy(input.copy()), pattern, reduction=reduction, **axes_lengths) - result = backend.to_numpy(result) - assert numpy.allclose(result, expected_result), f"Failed at {pattern}" - - -@pytest.mark.kernels_ci -def test_reduction_symbolic(): - for backend in sym_op_backends: - print("Reduction tests for ", backend.framework_name) - for reduction in REDUCTIONS: - input = numpy.arange(2 * 3 * 4 * 5 * 6, dtype="int64").reshape([2, 3, 4, 5, 6]) - input = input / input.astype("float64").mean() - # slight redundancy for simpler order - numpy version is evaluated multiple times - test_cases = [ - ["a b c d e -> ", {}, getattr(input, reduction)()], - ["a ... -> ", {}, getattr(input, reduction)()], - ["(a a2) ... (e e2) -> ", dict(a2=1, e2=1), getattr(input, reduction)()], - [ - "a b c d e -> (e c) a", - {}, - getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), - ], - [ - "a ... c d e -> (e c) a", - {}, - getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), - ], - [ - "a b c d e ... -> (e c) a", - {}, - getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1, 2]), - ], - ["a b c d e -> (e c a)", {}, getattr(input, reduction)(axis=(1, 3)).transpose(2, 1, 0).reshape([-1])], - ["(a a2) ... -> (a2 a) ...", dict(a2=1), input], - ] - for pattern, axes_lengths, expected_numpy_result in test_cases: - shapes = [input.shape, [None for _ in input.shape]] - for shape in shapes: - sym = backend.create_symbol(shape) - result_sym = reduce(sym, pattern, reduction=reduction, **axes_lengths) - result = backend.eval_symbol(result_sym, [(sym, input)]) - assert numpy.allclose(result, expected_numpy_result) - - if True: - shape = [] - _axes_lengths = {**axes_lengths} - for axis, length in zip("abcde", input.shape): - # filling as much as possible with Nones - if axis in pattern: - shape.append(None) - _axes_lengths[axis] = length - else: - shape.append(length) - sym = backend.create_symbol(shape) - result_sym = reduce(sym, pattern, reduction=reduction, **_axes_lengths) - result = backend.eval_symbol(result_sym, [(sym, input)]) - assert numpy.allclose(result, expected_numpy_result) - - -@pytest.mark.kernels_ci -def test_reduction_stress_imperatives(): - for backend in imp_op_backends: - print("Stress-testing reduction for ", backend.framework_name) - for reduction in REDUCTIONS + ("rearrange",): - dtype = "int64" - coincide = numpy.array_equal - if reduction in ["mean", "prod"]: - dtype = "float64" - coincide = numpy.allclose - max_dim = 11 - if "oneflow" in backend.framework_name: - max_dim = 7 - if "paddle" in backend.framework_name: - max_dim = 9 - for n_axes in range(max_dim): - shape = numpy.random.randint(2, 4, size=n_axes) - permutation = numpy.random.permutation(n_axes) - skipped = 0 if reduction == "rearrange" else numpy.random.randint(n_axes + 1) - left = " ".join("x" + str(i) for i in range(n_axes)) - right = " ".join("x" + str(i) for i in permutation[skipped:]) - pattern = left + "->" + right - x = numpy.arange(1, 1 + numpy.prod(shape), dtype=dtype).reshape(shape) - if reduction == "prod": - x /= x.mean() # to avoid overflows - result1 = reduce(x, pattern, reduction=reduction) - result2 = x.transpose(permutation) - if skipped > 0: - result2 = getattr(result2, reduction)(axis=tuple(range(skipped))) - assert coincide(result1, result2) - check_op_against_numpy(backend, x, pattern, reduction=reduction, axes_lengths={}, is_symbolic=False) - - -@pytest.mark.kernels_ci -def test_reduction_with_callable_imperatives(): - x_numpy = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]).astype("float32") - x_numpy /= x_numpy.max() - - def logsumexp_torch(x, tuple_of_axes): - return x.logsumexp(tuple_of_axes) - - def logsumexp_tf(x, tuple_of_axes): - import tensorflow as tf - - return tf.reduce_logsumexp(x, tuple_of_axes) - - def logsumexp_keras(x, tuple_of_axes): - import tensorflow.keras.backend as k - - return k.logsumexp(x, tuple_of_axes) - - def logsumexp_numpy(x, tuple_of_axes): - # very naive logsumexp to compare to - minused = x.max(tuple_of_axes) - y = x - x.max(tuple_of_axes, keepdims=True) - y = numpy.exp(y) - y = numpy.sum(y, axis=tuple_of_axes) - return numpy.log(y) + minused - - TorchBackend = einops._backends.TorchBackend - TensorflowBackend = einops._backends.TensorflowBackend - TFKerasBackend = einops._backends.TFKerasBackend - NumpyBackend = einops._backends.NumpyBackend - - backend2callback = { - TorchBackend.framework_name: logsumexp_torch, - TensorflowBackend.framework_name: logsumexp_tf, - TFKerasBackend.framework_name: logsumexp_keras, - NumpyBackend.framework_name: logsumexp_numpy, - } - - for backend in imp_op_backends: - if backend.framework_name not in backend2callback: - continue - - backend_callback = backend2callback[backend.framework_name] - - x_backend = backend.from_numpy(x_numpy) - for pattern1, pattern2 in equivalent_reduction_patterns: - print("Test reduction with callable for ", backend.framework_name, pattern1, pattern2) - output_numpy = reduce(x_numpy, pattern1, reduction=logsumexp_numpy) - output_backend = reduce(x_backend, pattern1, reduction=backend_callback) - assert numpy.allclose( - output_numpy, - backend.to_numpy(output_backend), - ) - - -@pytest.mark.kernels_ci -def test_enumerating_directions(): - for backend in imp_op_backends: - print("testing directions for", backend.framework_name) - for shape in [[], [1], [1, 1, 1], [2, 3, 5, 7]]: - x = numpy.arange(numpy.prod(shape)).reshape(shape) - axes1 = _enumerate_directions(x) - axes2 = _enumerate_directions(backend.from_numpy(x)) - assert len(axes1) == len(axes2) == len(shape) - for ax1, ax2 in zip(axes1, axes2): - ax2 = backend.to_numpy(ax2) - assert ax1.shape == ax2.shape - assert numpy.allclose(ax1, ax2) - - -@pytest.mark.kernels_ci -def test_concatenations_and_stacking(): - for backend in imp_op_backends: - print("testing shapes for ", backend.framework_name) - for n_arrays in [1, 2, 5]: - shapes = [[], [1], [1, 1], [2, 3, 5, 7], [1] * 6] - for shape in shapes: - arrays1 = [numpy.arange(i, i + numpy.prod(shape)).reshape(shape) for i in range(n_arrays)] - arrays2 = [backend.from_numpy(array) for array in arrays1] - result0 = numpy.asarray(arrays1) - result1 = rearrange(arrays1, "...->...") - result2 = rearrange(arrays2, "...->...") - assert numpy.array_equal(result0, result1) - assert numpy.array_equal(result1, backend.to_numpy(result2)) - - result1 = rearrange(arrays1, "b ... -> ... b") - result2 = rearrange(arrays2, "b ... -> ... b") - assert numpy.array_equal(result1, backend.to_numpy(result2)) - - -@pytest.mark.kernels_ci -def test_gradients_imperatives(): - # lazy - just checking reductions - for reduction in REDUCTIONS: - if reduction in ("any", "all"): - continue # non-differentiable ops - x = numpy.arange(1, 1 + 2 * 3 * 4).reshape([2, 3, 4]).astype("float32") - results = {} - for backend in imp_op_backends: - y0 = backend.from_numpy(x) - if not hasattr(y0, "grad"): - continue - - y1 = reduce(y0, "a b c -> c a", reduction=reduction) - y2 = reduce(y1, "c a -> a c", reduction=reduction) - y3 = reduce(y2, "a (c1 c2) -> a", reduction=reduction, c1=2) - y4 = reduce(y3, "... -> ", reduction=reduction) - - y4.backward() - grad = backend.to_numpy(y0.grad) - results[backend.framework_name] = grad - - print("comparing gradients for", results.keys()) - for name1, grad1 in results.items(): - for name2, grad2 in results.items(): - assert numpy.allclose(grad1, grad2), [name1, name2, "provided different gradients"] - - -@pytest.mark.kernels_ci -def test_tiling_imperatives(): - for backend in imp_op_backends: - print("Tiling tests for ", backend.framework_name) - input = numpy.arange(2 * 3 * 5, dtype="int64").reshape([2, 1, 3, 1, 5]) - test_cases = [ - (1, 1, 1, 1, 1), - (1, 2, 1, 3, 1), - (3, 1, 1, 4, 1), - ] - for repeats in test_cases: - expected = numpy.tile(input, repeats) - converted = backend.from_numpy(input) - repeated = backend.tile(converted, repeats) - result = backend.to_numpy(repeated) - assert numpy.array_equal(result, expected) - - -@pytest.mark.kernels_ci -def test_tiling_symbolic(): - for backend in sym_op_backends: - print("Tiling tests for ", backend.framework_name) - input = numpy.arange(2 * 3 * 5, dtype="int64").reshape([2, 1, 3, 1, 5]) - test_cases = [ - (1, 1, 1, 1, 1), - (1, 2, 1, 3, 1), - (3, 1, 1, 4, 1), - ] - for repeats in test_cases: - expected = numpy.tile(input, repeats) - sym = backend.create_symbol(input.shape) - result = backend.eval_symbol(backend.tile(sym, repeats), [[sym, input]]) - assert numpy.array_equal(result, expected) - - sym = backend.create_symbol([None] * len(input.shape)) - result = backend.eval_symbol(backend.tile(sym, repeats), [[sym, input]]) - assert numpy.array_equal(result, expected) - - -repeat_test_cases = [ - # all assume that input has shape [2, 3, 5] - ("a b c -> c a b", dict()), - ("a b c -> (c copy a b)", dict(copy=2, a=2, b=3, c=5)), - ("a b c -> (a copy) b c ", dict(copy=1)), - ("a b c -> (c a) (copy1 b copy2)", dict(a=2, copy1=1, copy2=2)), - ("a ... -> a ... copy", dict(copy=4)), - ("... c -> ... (copy1 c copy2)", dict(copy1=1, copy2=2)), - ("... -> ... ", dict()), - (" ... -> copy1 ... copy2 ", dict(copy1=2, copy2=3)), - ("a b c -> copy1 a copy2 b c () ", dict(copy1=2, copy2=1)), -] - - -def check_reversion(x, repeat_pattern, **sizes): - """Checks repeat pattern by running reduction""" - left, right = repeat_pattern.split("->") - reduce_pattern = right + "->" + left - repeated = repeat(x, repeat_pattern, **sizes) - reduced_min = reduce(repeated, reduce_pattern, reduction="min", **sizes) - reduced_max = reduce(repeated, reduce_pattern, reduction="max", **sizes) - assert numpy.array_equal(x, reduced_min) - assert numpy.array_equal(x, reduced_max) - - -@pytest.mark.kernels_ci -def test_repeat_numpy(): - # check repeat vs reduce. Repeat works ok if reverse reduction with min and max work well - x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) - x1 = repeat(x, "a b c -> copy a b c ", copy=1) - assert numpy.array_equal(x[None], x1) - for pattern, axis_dimensions in repeat_test_cases: - check_reversion(x, pattern, **axis_dimensions) - - -@pytest.mark.kernels_ci -def test_repeat_imperatives(): - x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) - for backend in imp_op_backends: - print("Repeat tests for ", backend.framework_name) - - for pattern, axis_dimensions in repeat_test_cases: - expected = repeat(x, pattern, **axis_dimensions) - converted = backend.from_numpy(x) - repeated = repeat(converted, pattern, **axis_dimensions) - result = backend.to_numpy(repeated) - assert numpy.array_equal(result, expected) - - -@pytest.mark.kernels_ci -def test_repeat_symbolic(): - x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) - - for backend in sym_op_backends: - print("Repeat tests for ", backend.framework_name) - - for pattern, axis_dimensions in repeat_test_cases: - expected = repeat(x, pattern, **axis_dimensions) - - sym = backend.create_symbol(x.shape) - result = backend.eval_symbol(repeat(sym, pattern, **axis_dimensions), [[sym, x]]) - assert numpy.array_equal(result, expected) - - -@pytest.mark.kernels_ci -def test_repeat_array_api(): - import numpy as xp - AA = einops.array_api - - if xp.__version__ < "2.0.0": - pytest.skip() - - x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) - - for pattern, axis_dimensions in repeat_test_cases: - expected = repeat(x, pattern, **axis_dimensions) - - result = AA.repeat(xp.from_dlpack(x), pattern, **axis_dimensions) - assert numpy.array_equal(AA.asnumpy(result + 0), expected) - - -test_cases_repeat_anonymous = [ - # all assume that input has shape [1, 2, 4, 6] - ("a b c d -> c a d b", dict()), - ("a b c d -> (c 2 d a b)", dict(a=1, c=4, d=6)), - ("1 b c d -> (d copy 1) 3 b c ", dict(copy=3)), - ("1 ... -> 3 ... ", dict()), - ("() ... d -> 1 (copy1 d copy2) ... ", dict(copy1=2, copy2=3)), - ("1 b c d -> (1 1) (1 b) 2 c 3 d (1 1)", dict()), -] - - -@pytest.mark.kernels_ci -def test_anonymous_axes(): - x = numpy.arange(1 * 2 * 4 * 6).reshape([1, 2, 4, 6]) - for pattern, axis_dimensions in test_cases_repeat_anonymous: - check_reversion(x, pattern, **axis_dimensions) - - -@pytest.mark.kernels_ci -def test_list_inputs(): - x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) - - assert numpy.array_equal( - rearrange(list(x), "... -> (...)"), - rearrange(x, "... -> (...)"), - ) - assert numpy.array_equal( - reduce(list(x), "a ... e -> (...)", "min"), - reduce(x, "a ... e -> (...)", "min"), - ) - assert numpy.array_equal( - repeat(list(x), "... -> b (...)", b=3), - repeat(x, "... -> b (...)", b=3), - ) - - -@pytest.mark.kernels_ci -def test_torch_compile_with_dynamic_shape(): - if not is_backend_tested("torch"): - pytest.skip() - import torch - - # somewhat reasonable debug messages - torch._dynamo.config.verbose = True - - def func1(x): - # test contains ellipsis - a, b, c, *other = x.shape - x = rearrange(x, "(a a2) b c ... -> b (c a2) (a ...)", a2=2) - # test contains passing expression as axis length - x = reduce(x, "b ca2 A -> b A", "sum", ca2=c * 2) - return x - - # seems can't test static and dynamic in the same test run. - # func1_compiled_static = torch.compile(func1, dynamic=False, fullgraph=True, backend='aot_eager') - func1_compiled_dynamic = torch.compile(func1, dynamic=True, fullgraph=True, backend="aot_eager") - - x = torch.randn(size=[4, 5, 6, 3]) - assert torch.equal(func1_compiled_dynamic(x), func1(x)) - # check with input of different dimensionality, and with all shape elements changed - x = torch.randn(size=[6, 3, 4, 2, 3]) - assert torch.equal(func1_compiled_dynamic(x), func1(x)) - - -def bit_count(x): - return sum((x >> i) & 1 for i in range(20)) - - -@pytest.mark.kernels_ci -def test_reduction_imperatives_booleans(): - """Checks that any/all reduction works in all frameworks""" - x_np = numpy.asarray([(bit_count(x) % 2) == 0 for x in range(2**6)]).reshape([2] * 6) - for backend in imp_op_backends: - print("Reduction any/all tests for ", backend.framework_name) - - for axis in range(6): - expected_result_any = numpy.any(x_np, axis=axis, keepdims=True) - expected_result_all = numpy.all(x_np, axis=axis, keepdims=True) - assert not numpy.array_equal(expected_result_any, expected_result_all) - - axes = list("abcdef") - axes_in = list(axes) - axes_out = list(axes) - axes_out[axis] = "1" - pattern = (" ".join(axes_in)) + " -> " + (" ".join(axes_out)) - - res_any = reduce(backend.from_numpy(x_np), pattern, reduction="any") - res_all = reduce(backend.from_numpy(x_np), pattern, reduction="all") - - assert numpy.array_equal(expected_result_any, backend.to_numpy(res_any)) - assert numpy.array_equal(expected_result_all, backend.to_numpy(res_all)) - - # expected result: any/all - expected_result_any = numpy.any(x_np, axis=(0, 1), keepdims=True) - expected_result_all = numpy.all(x_np, axis=(0, 1), keepdims=True) - pattern = "a b ... -> 1 1 ..." - res_any = reduce(backend.from_numpy(x_np), pattern, reduction="any") - res_all = reduce(backend.from_numpy(x_np), pattern, reduction="all") - assert numpy.array_equal(expected_result_any, backend.to_numpy(res_any)) - assert numpy.array_equal(expected_result_all, backend.to_numpy(res_all)) diff --git a/.ports/einops/overlay/tests/test_other.py b/.ports/einops/overlay/tests/test_other.py deleted file mode 100644 index 063f3298a..000000000 --- a/.ports/einops/overlay/tests/test_other.py +++ /dev/null @@ -1,307 +0,0 @@ -from doctest import testmod - -import numpy -import pytest - -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -AbstractBackend = einops._backends.AbstractBackend -rearrange = einops.einops.rearrange -parse_shape = einops.einops.parse_shape -_optimize_transformation = einops.einops._optimize_transformation -from . import collect_test_backends, is_backend_tested - -__author__ = "Alex Rogozhnikov" - - -@pytest.mark.kernels_ci -def test_doctests_examples(): - # tests docstrings, additionally - testmod(einops.layers, raise_on_error=True, extraglobs=dict(np=numpy)) - testmod(einops.einops, raise_on_error=True, extraglobs=dict(np=numpy)) - - -@pytest.mark.kernels_ci -def test_backends_installed(): - """ - This test will fail if some of backends are not installed or can't be imported - Other tests will just work and only test installed backends. - """ - from . import parse_backends_to_test - - backends_to_test = parse_backends_to_test() - errors = [] - for backend_type in AbstractBackend.__subclasses__(): - if backend_type.framework_name not in backends_to_test: - continue - try: - # instantiate - backend_type() - except Exception as e: - errors.append((backend_type.framework_name, e)) - assert len(errors) == 0, errors - - -@pytest.mark.kernels_ci -def test_optimize_transformations_numpy(): - print("Testing optimizations") - shapes = [[2] * n_dimensions for n_dimensions in range(14)] - shapes += [[3] * n_dimensions for n_dimensions in range(6)] - shapes += [[2, 3, 5, 7]] - shapes += [[2, 3, 5, 7, 11, 17]] - - for shape in shapes: - for attempt in range(5): - n_dimensions = len(shape) - x = numpy.random.randint(0, 2**12, size=shape).reshape([-1]) - init_shape = shape[:] - n_reduced = numpy.random.randint(0, n_dimensions + 1) - reduced_axes = tuple(numpy.random.permutation(n_dimensions)[:n_reduced]) - axes_reordering = numpy.random.permutation(n_dimensions - n_reduced) - final_shape = numpy.random.randint(0, 1024, size=333) # just random - - init_shape2, reduced_axes2, axes_reordering2, final_shape2 = combination2 = _optimize_transformation( - init_shape, reduced_axes, axes_reordering, final_shape - ) - - assert numpy.array_equal(final_shape, final_shape2) - result1 = x.reshape(init_shape).sum(axis=reduced_axes).transpose(axes_reordering).reshape([-1]) - result2 = x.reshape(init_shape2).sum(axis=reduced_axes2).transpose(axes_reordering2).reshape([-1]) - assert numpy.array_equal(result1, result2) - - # testing we can't optimize this formula again - combination3 = _optimize_transformation(*combination2) - for a, b in zip(combination2, combination3): - assert numpy.array_equal(a, b) - - -_IMPERATIVE_BACKENDS = collect_test_backends(symbolic=False, layers=False) - -x_np = numpy.zeros([10, 20, 30, 40]) - - -@pytest.mark.kernels_ci -def test_parse_shape_imperative(): - for backend in _IMPERATIVE_BACKENDS: - print("Shape parsing for ", backend.framework_name) - parsed1 = parse_shape(x_np, "a b c d") - parsed2 = parse_shape(backend.from_numpy(x_np), "a b c d") - assert parsed1 == parsed2 == dict(a=10, b=20, c=30, d=40) - assert parsed1 != dict(a=1, b=20, c=30, d=40) != parsed2 - - -@pytest.mark.kernels_ci -def test_underscore(): - for backend in _IMPERATIVE_BACKENDS: - parsed1 = parse_shape(x_np, "_ _ _ _") - parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ _ _") - assert parsed1 == parsed2 == dict() - - -@pytest.mark.kernels_ci -def test_underscore_one(): - for backend in _IMPERATIVE_BACKENDS: - parsed1 = parse_shape(x_np, "_ _ _ hello") - parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ _ hello") - assert parsed1 == parsed2 == dict(hello=40) - - -@pytest.mark.kernels_ci -def test_underscore_several(): - for backend in _IMPERATIVE_BACKENDS: - parsed1 = parse_shape(x_np, "_ _ a1 a1a111a") - parsed2 = parse_shape(backend.from_numpy(x_np), "_ _ a1 a1a111a") - assert parsed1 == parsed2 == dict(a1=30, a1a111a=40) - - -@pytest.mark.kernels_ci -def test_repeating(): - with pytest.raises(einops.EinopsError): - parse_shape(x_np, "a a b b") - - for backend in _IMPERATIVE_BACKENDS: - with pytest.raises(einops.EinopsError): - parse_shape(backend.from_numpy(x_np), "a a b b") - - -@pytest.mark.kernels_ci -def test_ellipsis(): - for backend in _IMPERATIVE_BACKENDS: - for shape, pattern, expected in [ - ([10, 20], "...", dict()), - ([10], "... a", dict(a=10)), - ([10, 20], "... a", dict(a=20)), - ([10, 20, 30], "... a", dict(a=30)), - ([10, 20, 30, 40], "... a", dict(a=40)), - ([10], "a ...", dict(a=10)), - ([10, 20], "a ...", dict(a=10)), - ([10, 20, 30], "a ...", dict(a=10)), - ([10, 20, 30, 40], "a ...", dict(a=10)), - ([10, 20, 30, 40], " a ... b", dict(a=10, b=40)), - ([10, 40], " a ... b", dict(a=10, b=40)), - ]: - x = numpy.ones(shape) - parsed1 = parse_shape(x, pattern) - parsed2 = parse_shape(backend.from_numpy(x), pattern) - assert parsed1 == parsed2 == expected - - -@pytest.mark.kernels_ci -def test_parse_with_anonymous_axes(): - for backend in _IMPERATIVE_BACKENDS: - for shape, pattern, expected in [ - ([1, 2, 3, 4], "1 2 3 a", dict(a=4)), - ([10, 1, 2], "a 1 2", dict(a=10)), - ([10, 1, 2], "a () 2", dict(a=10)), - ]: - x = numpy.ones(shape) - parsed1 = parse_shape(x, pattern) - parsed2 = parse_shape(backend.from_numpy(x), pattern) - assert parsed1 == parsed2 == expected - - -@pytest.mark.kernels_ci -def test_failures(): - for backend in _IMPERATIVE_BACKENDS: - # every test should fail - for shape, pattern in [ - ([1, 2, 3, 4], "a b c"), - ([1, 2, 3, 4], "2 a b c"), - ([1, 2, 3, 4], "a b c ()"), - ([1, 2, 3, 4], "a b c d e"), - ([1, 2, 3, 4], "a b c d e ..."), - ([1, 2, 3, 4], "a b c ()"), - ]: - with pytest.raises(RuntimeError): - x = numpy.ones(shape) - parse_shape(backend.from_numpy(x), pattern) - - -_SYMBOLIC_BACKENDS = [ - *collect_test_backends(symbolic=True, layers=False), - *collect_test_backends(symbolic=True, layers=True), -] - -# tensorflow.keras needs special way to compile, -# shape vars can be used only inside layers but not as outputs -_SYMBOLIC_BACKENDS = [backend for backend in _SYMBOLIC_BACKENDS if backend.framework_name != "tensorflow.keras"] - - -@pytest.mark.kernels_ci -@pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) -def test_parse_shape_symbolic(backend): - for shape in [ - [10, 20, 30, 40], - [10, 20, None, None], - [None, None, None, None], - ]: - print( - f"special shape parsing {backend.framework_name=} {shape=}", - ) - input_symbol = backend.create_symbol(shape) - - shape_placeholder = parse_shape(input_symbol, "a b c d") - shape = {} - for name, symbol in shape_placeholder.items(): - shape[name] = ( - symbol - if isinstance(symbol, int) - else backend.eval_symbol(symbol, [(input_symbol, numpy.zeros([10, 20, 30, 40]))]) - ) - print(shape) - result_placeholder = rearrange( - input_symbol, "a b (c1 c2) (d1 d2) -> (a b d1) c1 (c2 d2)", **parse_shape(input_symbol, "a b c1 _"), d2=2 - ) - result = backend.eval_symbol(result_placeholder, [(input_symbol, numpy.zeros([10, 20, 30, 40]))]) - print(result.shape) - assert result.shape == (10 * 20 * 20, 30, 1 * 2) - assert numpy.allclose(result, 0) - - -@pytest.mark.kernels_ci -@pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) -def test_parse_shape_symbolic_ellipsis(backend): - for static_shape, shape, pattern, expected in [ - ([10, 20], [None, None], "...", dict()), - ([10], [None], "... a", dict(a=10)), - ([10, 20], [None, None], "... a", dict(a=20)), - ([10, 20, 30], [None, None, None], "... a", dict(a=30)), - ([10, 20, 30, 40], [None, None, None, None], "... a", dict(a=40)), - ([10], [None], "a ...", dict(a=10)), - ([10, 20], [None, None], "a ...", dict(a=10)), - ([10, 20, 30], [None, None, None], "a ...", dict(a=10)), - ([10, 20, 30, 40], [None, None, None, None], "a ...", dict(a=10)), - ([10, 20, 30, 40], [None, None, None, None], " a ... b", dict(a=10, b=40)), - ([10, 40], [None, None], " a ... b ", dict(a=10, b=40)), - ]: - input_symbol = backend.create_symbol(shape) - shape_placeholder = parse_shape(input_symbol, pattern) - out_shape = {} - for name, symbol in shape_placeholder.items(): - if isinstance(symbol, int): - out_shape[name] = symbol - else: - out_shape[name] = backend.eval_symbol(symbol, [(input_symbol, numpy.zeros(static_shape))]) - assert out_shape == expected - - -@pytest.mark.kernels_ci -def test_is_float_type(): - backends = collect_test_backends(symbolic=False, layers=False) - backends += collect_test_backends(symbolic=False, layers=True) - for backend in backends: - for dtype in ["int32", "int64", "float32", "float64"]: - is_float = "float" in dtype - input = numpy.zeros([3, 4, 5], dtype=dtype) - input = backend.from_numpy(input) - assert backend.is_float_type(input) == is_float, (dtype, backend, input.dtype) - - -@pytest.mark.kernels_ci -def test_torch_compile(): - """ - Test ensures that allow_ops_in_compiled_graph allows compiling in a single graph - Additionally we ensure that after compilation cache works properly - (by changing shapes and patterns) - We additionally check that pack/unpack still can be handled - despite variable number of inputs/outputs - """ - if not is_backend_tested("torch"): - pytest.skip() - import torch - from torch import nn - repeat, reduce, pack, unpack, einsum = einops.repeat, einops.reduce, einops.pack, einops.unpack, einops.einsum - allow_ops_in_compiled_graph = einops._torch_specific.allow_ops_in_compiled_graph - - allow_ops_in_compiled_graph() - - class TorchModuleWithOperations(nn.Module): - def __init__(self) -> None: - super().__init__() - - def forward(self, x_abc, suffix=""): - a, b, c = x_abc.shape - - def suf(pattern): - parts = pattern.split() - return " ".join([p if p[-1] not in "acd" else p + suffix for p in parts]) - - # patterns look a bit strange because names a, c, d will be modified on every run - # by suf function - x_abcd = repeat(x_abc, suf("a b c -> a b c 4")) - x_abc = reduce(x_abcd, suf("a b c d -> a b c"), "min") - x_abdc, ps = pack([x_abc] * (2 + len(suffix)), suf("a b * c")) - x_array = unpack(rearrange(x_abdc, suf("a b d c -> (a b ) 1 c d")), ps, "ab one1 c *") - x1 = x_array[0] + len(x_array) - x1 = rearrange(x1, suf("(a b ) 1 c -> a b c"), b=b) - addition = einsum(x_abc, x_abcd, suf("a b c , a b c d -> d"))[0] - return x1 + addition - - original = TorchModuleWithOperations() - compiled = torch.compile(original, fullgraph=True, backend="aot_eager") - for size in [10, 20, 40]: - x = torch.rand([size, size + 1, size + 2]) - for suffix in ["", "suf1", "other_suffix"]: - result1 = compiled(x, suffix) - result2 = original(x, suffix) - assert torch.allclose(result1, result2) diff --git a/.ports/einops/overlay/tests/test_packing.py b/.ports/einops/overlay/tests/test_packing.py deleted file mode 100644 index d76553475..000000000 --- a/.ports/einops/overlay/tests/test_packing.py +++ /dev/null @@ -1,318 +0,0 @@ -import dataclasses -import typing - -import numpy as np -import pytest - -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -asnumpy = einops.asnumpy -pack = einops.pack -unpack = einops.unpack -from . import collect_test_backends - - -def pack_unpack(xs, pattern): - x, ps = pack(xs, pattern) - unpacked = unpack(xs, ps, pattern) - assert len(unpacked) == len(xs) - for a, b in zip(unpacked, xs): - assert np.allclose(asnumpy(a), asnumpy(b)) - - -def unpack_and_pack(x, ps, pattern: str): - unpacked = unpack(x, ps, pattern) - packed, ps2 = pack(unpacked, pattern=pattern) - - assert np.allclose(asnumpy(packed), asnumpy(x)) - return unpacked - - -def unpack_and_pack_against_numpy(x, ps, pattern: str): - capturer_backend = CaptureException() - capturer_numpy = CaptureException() - - with capturer_backend: - unpacked = unpack(x, ps, pattern) - packed, ps2 = pack(unpacked, pattern=pattern) - - with capturer_numpy: - x_np = asnumpy(x) - unpacked_np = unpack(x_np, ps, pattern) - packed_np, ps3 = pack(unpacked_np, pattern=pattern) - - assert type(capturer_numpy.exception) == type(capturer_backend.exception) # noqa E721 - if capturer_numpy.exception is not None: - # both failed - return - else: - # neither failed, check results are identical - assert np.allclose(asnumpy(packed), asnumpy(x)) - assert np.allclose(asnumpy(packed_np), asnumpy(x)) - assert len(unpacked) == len(unpacked_np) - for a, b in zip(unpacked, unpacked_np): - assert np.allclose(asnumpy(a), b) - - -class CaptureException: - def __enter__(self): - self.exception = None - - def __exit__(self, exc_type, exc_val, exc_tb): - self.exception = exc_val - return True - - -@pytest.mark.kernels_ci -def test_numpy_trivial(H=13, W=17): - def rand(*shape): - return np.random.random(shape) - - def check(a, b): - assert a.dtype == b.dtype - assert a.shape == b.shape - assert np.all(a == b) - - r, g, b = rand(3, H, W) - embeddings = rand(H, W, 32) - - check( - np.stack([r, g, b], axis=2), - pack([r, g, b], "h w *")[0], - ) - check( - np.stack([r, g, b], axis=1), - pack([r, g, b], "h * w")[0], - ) - check( - np.stack([r, g, b], axis=0), - pack([r, g, b], "* h w")[0], - ) - - check( - np.concatenate([r, g, b], axis=1), - pack([r, g, b], "h *")[0], - ) - check( - np.concatenate([r, g, b], axis=0), - pack([r, g, b], "* w")[0], - ) - - i = np.index_exp[:, :, None] - check( - np.concatenate([r[i], g[i], b[i], embeddings], axis=2), - pack([r, g, b, embeddings], "h w *")[0], - ) - - with pytest.raises(EinopsError): - pack([r, g, b, embeddings], "h w nonexisting_axis *") - - pack([r, g, b], "some_name_for_H some_name_for_w1 *") - - with pytest.raises(EinopsError): - pack([r, g, b, embeddings], "h _w *") # no leading underscore - with pytest.raises(EinopsError): - pack([r, g, b, embeddings], "h_ w *") # no trailing underscore - with pytest.raises(EinopsError): - pack([r, g, b, embeddings], "1h_ w *") - with pytest.raises(EinopsError): - pack([r, g, b, embeddings], "1 w *") - with pytest.raises(EinopsError): - pack([r, g, b, embeddings], "h h *") - # capital and non-capital are different - pack([r, g, b, embeddings], "h H *") - - -@dataclasses.dataclass -class UnpackTestCase: - shape: typing.Tuple[int, ...] - pattern: str - - def dim(self): - return self.pattern.split().index("*") - - def selfcheck(self): - assert self.shape[self.dim()] == 5 - - -cases = [ - # NB: in all cases unpacked axis is of length 5. - # that's actively used in tests below - UnpackTestCase((5,), "*"), - UnpackTestCase((5, 7), "* seven"), - UnpackTestCase((7, 5), "seven *"), - UnpackTestCase((5, 3, 4), "* three four"), - UnpackTestCase((4, 5, 3), "four * three"), - UnpackTestCase((3, 4, 5), "three four *"), -] - - -@pytest.mark.kernels_ci -def test_pack_unpack_with_numpy(): - case: UnpackTestCase - - for case in cases: - shape = case.shape - pattern = case.pattern - - x = np.random.random(shape) - # all correct, no minus 1 - unpack_and_pack(x, [[2], [1], [2]], pattern) - # no -1, asking for wrong shapes - with pytest.raises(BaseException): - unpack_and_pack(x, [[2], [1], [2]], pattern + " non_existent_axis") - with pytest.raises(BaseException): - unpack_and_pack(x, [[2], [1], [1]], pattern) - with pytest.raises(BaseException): - unpack_and_pack(x, [[4], [1], [1]], pattern) - # all correct, with -1 - unpack_and_pack(x, [[2], [1], [-1]], pattern) - unpack_and_pack(x, [[2], [-1], [2]], pattern) - unpack_and_pack(x, [[-1], [1], [2]], pattern) - _, _, last = unpack_and_pack(x, [[2], [3], [-1]], pattern) - assert last.shape[case.dim()] == 0 - # asking for more elements than available - with pytest.raises(BaseException): - unpack(x, [[2], [4], [-1]], pattern) - # this one does not raise, because indexing x[2:1] just returns zero elements - # with pytest.raises(BaseException): - # unpack(x, [[2], [-1], [4]], pattern) - with pytest.raises(BaseException): - unpack(x, [[-1], [1], [5]], pattern) - - # all correct, -1 nested - rs = unpack_and_pack(x, [[1, 2], [1, 1], [-1, 1]], pattern) - assert all(len(r.shape) == len(x.shape) + 1 for r in rs) - rs = unpack_and_pack(x, [[1, 2], [1, -1], [1, 1]], pattern) - assert all(len(r.shape) == len(x.shape) + 1 for r in rs) - rs = unpack_and_pack(x, [[2, -1], [1, 2], [1, 1]], pattern) - assert all(len(r.shape) == len(x.shape) + 1 for r in rs) - - # asking for more elements, -1 nested - with pytest.raises(BaseException): - unpack(x, [[-1, 2], [1], [5]], pattern) - with pytest.raises(BaseException): - unpack(x, [[2, 2], [2], [5, -1]], pattern) - - # asking for non-divisible number of elements - with pytest.raises(BaseException): - unpack(x, [[2, 1], [1], [3, -1]], pattern) - with pytest.raises(BaseException): - unpack(x, [[2, 1], [3, -1], [1]], pattern) - with pytest.raises(BaseException): - unpack(x, [[3, -1], [2, 1], [1]], pattern) - - # -1 takes zero - unpack_and_pack(x, [[0], [5], [-1]], pattern) - unpack_and_pack(x, [[0], [-1], [5]], pattern) - unpack_and_pack(x, [[-1], [5], [0]], pattern) - - # -1 takes zero, -1 - unpack_and_pack(x, [[2, -1], [1, 5]], pattern) - - -@pytest.mark.kernels_ci -def test_pack_unpack_against_numpy(): - for backend in collect_test_backends(symbolic=False, layers=False): - print(f"test packing against numpy for {backend.framework_name}") - check_zero_len = True - - for case in cases: - unpack_and_pack = unpack_and_pack_against_numpy - shape = case.shape - pattern = case.pattern - - x = np.random.random(shape) - x = backend.from_numpy(x) - # all correct, no minus 1 - unpack_and_pack(x, [[2], [1], [2]], pattern) - # no -1, asking for wrong shapes - with pytest.raises(BaseException): - unpack(x, [[2], [1], [1]], pattern) - - with pytest.raises(BaseException): - unpack(x, [[4], [1], [1]], pattern) - # all correct, with -1 - unpack_and_pack(x, [[2], [1], [-1]], pattern) - unpack_and_pack(x, [[2], [-1], [2]], pattern) - unpack_and_pack(x, [[-1], [1], [2]], pattern) - - # asking for more elements than available - with pytest.raises(BaseException): - unpack(x, [[2], [4], [-1]], pattern) - # this one does not raise, because indexing x[2:1] just returns zero elements - # with pytest.raises(BaseException): - # unpack(x, [[2], [-1], [4]], pattern) - with pytest.raises(BaseException): - unpack(x, [[-1], [1], [5]], pattern) - - # all correct, -1 nested - unpack_and_pack(x, [[1, 2], [1, 1], [-1, 1]], pattern) - unpack_and_pack(x, [[1, 2], [1, -1], [1, 1]], pattern) - unpack_and_pack(x, [[2, -1], [1, 2], [1, 1]], pattern) - - # asking for more elements, -1 nested - with pytest.raises(BaseException): - unpack(x, [[-1, 2], [1], [5]], pattern) - with pytest.raises(BaseException): - unpack(x, [[2, 2], [2], [5, -1]], pattern) - - # asking for non-divisible number of elements - with pytest.raises(BaseException): - unpack(x, [[2, 1], [1], [3, -1]], pattern) - with pytest.raises(BaseException): - unpack(x, [[2, 1], [3, -1], [1]], pattern) - with pytest.raises(BaseException): - unpack(x, [[3, -1], [2, 1], [1]], pattern) - - if check_zero_len: - # -1 takes zero - unpack_and_pack(x, [[2], [3], [-1]], pattern) - unpack_and_pack(x, [[0], [5], [-1]], pattern) - unpack_and_pack(x, [[0], [-1], [5]], pattern) - unpack_and_pack(x, [[-1], [5], [0]], pattern) - - # -1 takes zero, -1 - unpack_and_pack(x, [[2, -1], [1, 5]], pattern) - - -@pytest.mark.kernels_ci -def test_pack_unpack_array_api(): - AA = einops.array_api - import numpy as xp - - if xp.__version__ < "2.0.0": - pytest.skip() - - for case in cases: - shape = case.shape - pattern = case.pattern - x_np = np.random.random(shape) - x_xp = xp.from_dlpack(x_np) - - for ps in [ - [[2], [1], [2]], - [[1], [1], [-1]], - [[1], [1], [-1, 3]], - [[2, 1], [1, 1, 1], [-1]], - ]: - x_np_split = unpack(x_np, ps, pattern) - x_xp_split = AA.unpack(x_xp, ps, pattern) - for a, b in zip(x_np_split, x_xp_split): - assert np.allclose(a, AA.asnumpy(b + 0)) - - x_agg_np, ps1 = pack(x_np_split, pattern) - x_agg_xp, ps2 = AA.pack(x_xp_split, pattern) - assert ps1 == ps2 - assert np.allclose(x_agg_np, AA.asnumpy(x_agg_xp)) - - for ps in [ - [[2, 3]], - [[1], [5]], - [[1], [5], [-1]], - [[1], [2, 3]], - [[1], [5], [-1, 2]], - ]: - with pytest.raises(BaseException): - unpack(x_np, ps, pattern) diff --git a/.ports/einops/overlay/tests/test_parsing.py b/.ports/einops/overlay/tests/test_parsing.py deleted file mode 100644 index 491126793..000000000 --- a/.ports/einops/overlay/tests/test_parsing.py +++ /dev/null @@ -1,134 +0,0 @@ -import pytest - -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -ParsedExpression = einops.parsing.ParsedExpression -AnonymousAxis = einops.parsing.AnonymousAxis -_ellipsis = einops.parsing._ellipsis - -__author__ = "Alex Rogozhnikov" - - -class AnonymousAxisPlaceholder: - def __init__(self, value: int): - self.value = value - assert isinstance(self.value, int) - - def __eq__(self, other): - return isinstance(other, AnonymousAxis) and self.value == other.value - - -@pytest.mark.kernels_ci -def test_anonymous_axes(): - a, b = AnonymousAxis("2"), AnonymousAxis("2") - assert a != b - c, d = AnonymousAxisPlaceholder(2), AnonymousAxisPlaceholder(3) - assert a == c and b == c - assert a != d and b != d - assert [a, 2, b] == [c, 2, c] - - -@pytest.mark.kernels_ci -def test_elementary_axis_name(): - for name in [ - "a", - "b", - "h", - "dx", - "h1", - "zz", - "i9123", - "somelongname", - "Alex", - "camelCase", - "u_n_d_e_r_score", - "unreasonablyLongAxisName", - ]: - assert ParsedExpression.check_axis_name(name) - - for name in ["", "2b", "12", "_startWithUnderscore", "endWithUnderscore_", "_", "...", _ellipsis]: - assert not ParsedExpression.check_axis_name(name) - - -@pytest.mark.kernels_ci -def test_invalid_expressions(): - # double ellipsis should raise an error - ParsedExpression("... a b c d") - with pytest.raises(EinopsError): - ParsedExpression("... a b c d ...") - with pytest.raises(EinopsError): - ParsedExpression("... a b c (d ...)") - with pytest.raises(EinopsError): - ParsedExpression("(... a) b c (d ...)") - - # double/missing/enclosed parenthesis - ParsedExpression("(a) b c (d ...)") - with pytest.raises(EinopsError): - ParsedExpression("(a)) b c (d ...)") - with pytest.raises(EinopsError): - ParsedExpression("(a b c (d ...)") - with pytest.raises(EinopsError): - ParsedExpression("(a) (()) b c (d ...)") - with pytest.raises(EinopsError): - ParsedExpression("(a) ((b c) (d ...))") - - # invalid identifiers - ParsedExpression("camelCase under_scored cApiTaLs ß ...") - with pytest.raises(EinopsError): - ParsedExpression("1a") - with pytest.raises(EinopsError): - ParsedExpression("_pre") - with pytest.raises(EinopsError): - ParsedExpression("...pre") - with pytest.raises(EinopsError): - ParsedExpression("pre...") - - -@pytest.mark.kernels_ci -def test_parse_expression(): - parsed = ParsedExpression("a1 b1 c1 d1") - assert parsed.identifiers == {"a1", "b1", "c1", "d1"} - assert parsed.composition == [["a1"], ["b1"], ["c1"], ["d1"]] - assert not parsed.has_non_unitary_anonymous_axes - assert not parsed.has_ellipsis - - parsed = ParsedExpression("() () () ()") - assert parsed.identifiers == set() - assert parsed.composition == [[], [], [], []] - assert not parsed.has_non_unitary_anonymous_axes - assert not parsed.has_ellipsis - - parsed = ParsedExpression("1 1 1 ()") - assert parsed.identifiers == set() - assert parsed.composition == [[], [], [], []] - assert not parsed.has_non_unitary_anonymous_axes - assert not parsed.has_ellipsis - - aap = AnonymousAxisPlaceholder - - parsed = ParsedExpression("5 (3 4)") - assert len(parsed.identifiers) == 3 and {i.value for i in parsed.identifiers} == {3, 4, 5} - assert parsed.composition == [[aap(5)], [aap(3), aap(4)]] - assert parsed.has_non_unitary_anonymous_axes - assert not parsed.has_ellipsis - - parsed = ParsedExpression("5 1 (1 4) 1") - assert len(parsed.identifiers) == 2 and {i.value for i in parsed.identifiers} == {4, 5} - assert parsed.composition == [[aap(5)], [], [aap(4)], []] - - parsed = ParsedExpression("name1 ... a1 12 (name2 14)") - assert len(parsed.identifiers) == 6 - assert parsed.identifiers.difference({"name1", _ellipsis, "a1", "name2"}).__len__() == 2 - assert parsed.composition == [["name1"], _ellipsis, ["a1"], [aap(12)], ["name2", aap(14)]] - assert parsed.has_non_unitary_anonymous_axes - assert parsed.has_ellipsis - assert not parsed.has_ellipsis_parenthesized - - parsed = ParsedExpression("(name1 ... a1 12) name2 14") - assert len(parsed.identifiers) == 6 - assert parsed.identifiers.difference({"name1", _ellipsis, "a1", "name2"}).__len__() == 2 - assert parsed.composition == [["name1", _ellipsis, "a1", aap(12)], ["name2"], [aap(14)]] - assert parsed.has_non_unitary_anonymous_axes - assert parsed.has_ellipsis - assert parsed.has_ellipsis_parenthesized diff --git a/.ports/einops/port.kdl b/.ports/einops/port.kdl deleted file mode 100644 index 2b92e57fb..000000000 --- a/.ports/einops/port.kdl +++ /dev/null @@ -1,60 +0,0 @@ -recipe version=1 - -// einops -// -// einops as a Hub kernel: the package is pure Python, so the manifest is -// noarch and the port is a move plus the import rewrite that lets the package -// be imported under whatever directory name the Hub loads it from. The test -// suite is not derived - it was rewritten to pull the kernel off the Hub - so -// it comes from the overlay. -// -// kernel-port .ports/einops/port.kdl \ -// --dir --out einops - -source \ - repo="https://github.com/arogozhnikov/einops.git" \ - commit="ad2c8d67d08cbebb45e1aff2be21dc88635fc695" - -prune keep="einops/**" -delete in="einops/tests/**" - -move \ - from="einops" \ - to="torch-ext/einops" - -relativize_imports \ - in="torch-ext/einops/**" \ - package_root="torch-ext/einops" \ - changes=7 -expect \ - in="torch-ext/einops/**" \ - find="from einops" \ - count=0 - -// array_api is a submodule upstream leaves out of the package init; the kernel -// exposes it, since a Hub consumer cannot import einops.array_api by name. -replace \ - in="torch-ext/einops/__init__.py" \ - count=1 \ - find="from .packing import pack, unpack\n" \ - with="from .packing import pack, unpack\nfrom . import array_api\n" - -// The rewritten tests and the flake. -overlay from="overlay" - -expect \ - in="torch-ext/einops/**" \ - files=18 -expect \ - in="tests/**" \ - files=9 - -manifest \ - name="einops" \ - version=1 \ - edition=5 \ - license="MIT" \ - upstream="https://github.com/arogozhnikov/einops.git" \ - backends="cpu,cuda,metal,rocm,xpu" \ - repo_id="kernels-community/einops" \ - noarch=#true diff --git a/.ports/einops/overlay/flake.lock b/einops/port/overlay/flake.lock similarity index 100% rename from .ports/einops/overlay/flake.lock rename to einops/port/overlay/flake.lock diff --git a/.ports/einops/overlay/flake.nix b/einops/port/overlay/flake.nix similarity index 100% rename from .ports/einops/overlay/flake.nix rename to einops/port/overlay/flake.nix diff --git a/einops/port/port.kdl b/einops/port/port.kdl new file mode 100644 index 000000000..62cf61e47 --- /dev/null +++ b/einops/port/port.kdl @@ -0,0 +1,197 @@ +recipe version=1 + +// einops +// +// einops as a Hub kernel: the package is pure Python, so the manifest is +// noarch and the port is a move plus the import rewrite that lets the package +// be imported under whatever directory name the Hub loads it from. +// +// The test suite is derived from upstream too. Upstream's tests import einops +// by name; the port rebinds those imports to the kernel loaded off the Hub and +// marks each test for kernel CI. Everything else is upstream's. +// +// The flake has no upstream equivalent, so it comes from the overlay. +// +// kernel-port einops/port/port.kdl \ +// --dir --out einops/src + +source \ + repo="https://github.com/arogozhnikov/einops.git" \ + commit="ad2c8d67d08cbebb45e1aff2be21dc88635fc695" + +prune keep="einops/**" + +move \ + from="einops/tests" \ + to="tests" +move \ + from="einops" \ + to="torch-ext/einops" + +relativize_imports \ + in="torch-ext/einops/**" \ + package_root="torch-ext/einops" \ + changes=7 +expect \ + in="torch-ext/einops/**" \ + find="from einops" \ + count=0 + +// array_api is a submodule upstream leaves out of the package init; the kernel +// exposes it, since a Hub consumer cannot import einops.array_api by name. +replace \ + in="torch-ext/einops/__init__.py" \ + count=1 \ + find="from .packing import pack, unpack\n" \ + with="from .packing import pack, unpack\nfrom . import array_api\n" + +// The test suite imports the package by name; the kernel is loaded off the +// Hub instead, so each import site is rebound to the loaded module. + +replace \ + in="tests/__init__.py" \ + count=1 \ + find="from einops import _backends\n" \ + with="import kernels\n_einops = kernels.get_kernel(\"kernels-community/einops\", version=1)\n_backends = _einops._backends\n" + +// Upstream hard-fails when EINOPS_TEST_BACKENDS is unset. In kernel CI no one +// sets it, so fall back to whichever frameworks actually import. + +replace \ + in="tests/__init__.py" \ + count=1 \ + find=" if ENVVAR_NAME not in os.environ:\n raise RuntimeError(f\"Testing frameworks were not specified, env var {ENVVAR_NAME} not set\")\n parsed_backends = os.environ[ENVVAR_NAME].split(\",\")\n _known_backends = find_names_of_all_frameworks()\n for backend_name in parsed_backends:\n if backend_name not in _known_backends:\n raise RuntimeError(f\"Unknown framework: {backend_name}\")\n\n return parsed_backends\n" \ + with=" if ENVVAR_NAME in os.environ:\n parsed_backends = os.environ[ENVVAR_NAME].split(\",\")\n _known_backends = find_names_of_all_frameworks()\n for backend_name in parsed_backends:\n if backend_name not in _known_backends:\n raise RuntimeError(f\"Unknown framework: {backend_name}\")\n return parsed_backends\n\n # Auto-detect available frameworks.\n available = []\n for module, name in [(\"torch\", \"torch\"), (\"jax\", \"jax\"), (\"numpy\", \"numpy\"), (\"tensorflow\", \"tensorflow\")]:\n try:\n __import__(module)\n available.append(name)\n except ImportError:\n pass\n return available\n" + +replace \ + in="tests/test_einsum.py" \ + count=1 \ + find="from einops.tests import collect_test_backends\nfrom einops.einops import _compactify_pattern_for_einsum, einsum, EinopsError\n" \ + with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\n_compactify_pattern_for_einsum = einops.einops._compactify_pattern_for_einsum\neinsum = einops.einops.einsum\nEinopsError = einops.EinopsError\nfrom . import collect_test_backends\n" + +replace \ + in="tests/test_examples.py" \ + count=1 \ + find="from einops import rearrange, parse_shape, reduce\nfrom einops.tests import is_backend_tested\nfrom einops.tests.test_ops import imp_op_backends\n" \ + with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nrearrange = einops.rearrange\nparse_shape = einops.parse_shape\nreduce = einops.reduce\nfrom . import is_backend_tested\nfrom .test_ops import imp_op_backends\n" + +replace \ + in="tests/test_layers.py" \ + count=1 \ + find="from einops import rearrange, reduce, EinopsError\nfrom einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS\n" \ + with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nrearrange = einops.rearrange\nreduce = einops.reduce\nEinopsError = einops.EinopsError\nfrom . import collect_test_backends, is_backend_tested\nfrom . import FLOAT_REDUCTIONS as REDUCTIONS\n" + +replace \ + in="tests/test_layers.py" \ + count=1 \ + find=" from einops.layers.torch import Rearrange, Reduce, EinMix\n" \ + with=" Rearrange, Reduce, EinMix = einops.layers.torch.Rearrange, einops.layers.torch.Reduce, einops.layers.torch.EinMix\n" + +replace \ + in="tests/test_layers.py" \ + count=1 \ + find=" from einops.layers.keras import Rearrange, Reduce, EinMix, keras_custom_objects\n" \ + with=" Rearrange, Reduce, EinMix, keras_custom_objects = einops.layers.keras.Rearrange, einops.layers.keras.Reduce, einops.layers.keras.EinMix, einops.layers.keras.keras_custom_objects\n" + +replace \ + in="tests/test_layers.py" \ + count=1 \ + find=" from einops.layers.flax import EinMix, Reduce, Rearrange\n" \ + with=" EinMix, Reduce, Rearrange = einops.layers.flax.EinMix, einops.layers.flax.Reduce, einops.layers.flax.Rearrange\n" + +replace \ + in="tests/test_layers.py" \ + count=2 \ + find=" from einops.layers._einmix import _EinmixDebugger\n" \ + with=" _EinmixDebugger = einops.layers._einmix._EinmixDebugger\n" + +// numpy 2 dropped in1d. + +replace \ + in="tests/test_layers.py" \ + count=1 \ + find="numpy.in1d(grad, [0, 1])" \ + with="numpy.isin(grad, [0, 1])" + +replace \ + in="tests/test_ops.py" \ + count=1 \ + find="from einops import EinopsError\nfrom einops.einops import rearrange, reduce, repeat, _enumerate_directions\nfrom einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS\n" \ + with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nEinopsError = einops.EinopsError\nrearrange = einops.einops.rearrange\nreduce = einops.einops.reduce\nrepeat = einops.einops.repeat\n_enumerate_directions = einops.einops._enumerate_directions\nfrom . import collect_test_backends, is_backend_tested\nfrom . import FLOAT_REDUCTIONS as REDUCTIONS\n" + +replace \ + in="tests/test_ops.py" \ + count=1 \ + find=" from einops._backends import TorchBackend, TensorflowBackend, TFKerasBackend, NumpyBackend\n" \ + with=" TorchBackend = einops._backends.TorchBackend\n TensorflowBackend = einops._backends.TensorflowBackend\n TFKerasBackend = einops._backends.TFKerasBackend\n NumpyBackend = einops._backends.NumpyBackend\n" + +replace \ + in="tests/test_other.py" \ + count=1 \ + find="import einops\nimport einops.layers\nimport einops.parsing\nfrom einops._backends import AbstractBackend\nfrom einops.einops import rearrange, parse_shape, _optimize_transformation\nfrom einops.tests import collect_test_backends, is_backend_tested\n" \ + with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nAbstractBackend = einops._backends.AbstractBackend\nrearrange = einops.einops.rearrange\nparse_shape = einops.einops.parse_shape\n_optimize_transformation = einops.einops._optimize_transformation\nfrom . import collect_test_backends, is_backend_tested\n" + +replace \ + in="tests/test_other.py" \ + count=1 \ + find=" from einops import repeat, reduce, pack, unpack, einsum\n from einops._torch_specific import allow_ops_in_compiled_graph\n" \ + with=" repeat, reduce, pack, unpack, einsum = einops.repeat, einops.reduce, einops.pack, einops.unpack, einops.einsum\n allow_ops_in_compiled_graph = einops._torch_specific.allow_ops_in_compiled_graph\n" + +replace \ + in="tests/test_packing.py" \ + count=1 \ + find="from einops import EinopsError, asnumpy, pack, unpack\nfrom einops.tests import collect_test_backends\n" \ + with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nEinopsError = einops.EinopsError\nasnumpy = einops.asnumpy\npack = einops.pack\nunpack = einops.unpack\nfrom . import collect_test_backends\n" + +replace \ + in="tests/test_parsing.py" \ + count=1 \ + find="from einops import EinopsError\nfrom einops.parsing import ParsedExpression, AnonymousAxis, _ellipsis\n" \ + with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nEinopsError = einops.EinopsError\nParsedExpression = einops.parsing.ParsedExpression\nAnonymousAxis = einops.parsing.AnonymousAxis\n_ellipsis = einops.parsing._ellipsis\n" + +// array_api is reached as an attribute of the loaded kernel (see the +// __init__.py edit above), not as a submodule import. + +replace \ + in="tests/**" \ + count=4 \ + find=" from einops import array_api as AA\n" \ + with=" AA = einops.array_api\n" + +// Kernel CI selects tests by marker; upstream has none. + +replace \ + in="tests/**" \ + count=63 \ + find="\ndef test_" \ + with="\n@pytest.mark.kernels_ci\ndef test_" + +// The flake: hand-written, no upstream equivalent. +overlay from="overlay" + +expect \ + in="torch-ext/einops/**" \ + files=18 +expect \ + in="tests/**" \ + files=9 +// run_tests.py is upstream's standalone pip-install-and-run driver. Kernel CI +// does not use it, so it is shipped verbatim and excluded from this check. +expect \ + in="tests/__init__.py" \ + find="from einops" \ + count=0 +expect \ + in="tests/test_*.py" \ + find="from einops" \ + count=0 + +manifest \ + name="einops" \ + version=1 \ + edition=5 \ + license="MIT" \ + upstream="https://github.com/arogozhnikov/einops.git" \ + backends="cpu,cuda,metal,rocm,xpu" \ + repo_id="kernels-community/einops" \ + noarch=#true diff --git a/einops/.port-provenance.json b/einops/src/port-provenance.json similarity index 55% rename from einops/.port-provenance.json rename to einops/src/port-provenance.json index c99c151c3..f0a2f76fc 100644 --- a/einops/.port-provenance.json +++ b/einops/src/port-provenance.json @@ -1,6 +1,6 @@ { "format": 1, - "recipe": "sha256:31835e78fc155e6165bfb986a8b7bd3c58897c62bc01c6d90de7a8e1d58d16ab", + "recipe": "sha256:7b9b97f928ded68c67302ce872b875a604abd35d2be6543730a20a38da97bf47", "runner": "kernel-port 0.1.0", "sources": { "upstream": { @@ -8,5 +8,5 @@ "commit": "ad2c8d67d08cbebb45e1aff2be21dc88635fc695" } }, - "tree": "sha256:2764199b7ae5d7a03a0780d2907b9eda7a9ae6ade00d6a879c2ef2489a1be969" + "tree": "sha256:f7dab98a4d28b88488879056c333147c7981814aa777031aefac7a2b8c32c1b9" } diff --git a/einops/src/tests/test_other.py b/einops/src/tests/test_other.py index 063f3298a..71fcb9ce3 100644 --- a/einops/src/tests/test_other.py +++ b/einops/src/tests/test_other.py @@ -187,8 +187,8 @@ def test_failures(): _SYMBOLIC_BACKENDS = [backend for backend in _SYMBOLIC_BACKENDS if backend.framework_name != "tensorflow.keras"] -@pytest.mark.kernels_ci @pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) +@pytest.mark.kernels_ci def test_parse_shape_symbolic(backend): for shape in [ [10, 20, 30, 40], @@ -218,8 +218,8 @@ def test_parse_shape_symbolic(backend): assert numpy.allclose(result, 0) -@pytest.mark.kernels_ci @pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) +@pytest.mark.kernels_ci def test_parse_shape_symbolic_ellipsis(backend): for static_shape, shape, pattern, expected in [ ([10, 20], [None, None], "...", dict()), From 324a1e2b26cd7499938be206e8dacd16e224fdb8 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Wed, 26 Aug 2026 17:50:48 -0400 Subject: [PATCH 04/15] feat: use multi-line strings in the einops port recipe Signed-off-by: David Holtz --- einops/port/port.kdl | 297 +++++++++++++++++++++----------- einops/src/port-provenance.json | 2 +- 2 files changed, 198 insertions(+), 101 deletions(-) diff --git a/einops/port/port.kdl b/einops/port/port.kdl index 62cf61e47..faa5903dd 100644 --- a/einops/port/port.kdl +++ b/einops/port/port.kdl @@ -12,6 +12,10 @@ recipe version=1 // // The flake has no upstream equivalent, so it comes from the overlay. // +// Multi-line payloads below use KDL raw strings. KDL dedents them against the +// closing delimiter, so every "\"\"\"# sits at column 0 and the payload keeps +// the exact indentation it must match in the source. +// // kernel-port einops/port/port.kdl \ // --dir --out einops/src @@ -39,124 +43,217 @@ expect \ // array_api is a submodule upstream leaves out of the package init; the kernel // exposes it, since a Hub consumer cannot import einops.array_api by name. -replace \ - in="torch-ext/einops/__init__.py" \ - count=1 \ - find="from .packing import pack, unpack\n" \ - with="from .packing import pack, unpack\nfrom . import array_api\n" +replace in="torch-ext/einops/__init__.py" count=1 \ + find="from .packing import pack, unpack" \ + with=#""" +from .packing import pack, unpack +from . import array_api +"""# // The test suite imports the package by name; the kernel is loaded off the // Hub instead, so each import site is rebound to the loaded module. -replace \ - in="tests/__init__.py" \ - count=1 \ - find="from einops import _backends\n" \ - with="import kernels\n_einops = kernels.get_kernel(\"kernels-community/einops\", version=1)\n_backends = _einops._backends\n" +replace in="tests/__init__.py" count=1 \ + find="from einops import _backends" \ + with=#""" +import kernels +_einops = kernels.get_kernel("kernels-community/einops", version=1) +_backends = _einops._backends +"""# // Upstream hard-fails when EINOPS_TEST_BACKENDS is unset. In kernel CI no one // sets it, so fall back to whichever frameworks actually import. -replace \ - in="tests/__init__.py" \ - count=1 \ - find=" if ENVVAR_NAME not in os.environ:\n raise RuntimeError(f\"Testing frameworks were not specified, env var {ENVVAR_NAME} not set\")\n parsed_backends = os.environ[ENVVAR_NAME].split(\",\")\n _known_backends = find_names_of_all_frameworks()\n for backend_name in parsed_backends:\n if backend_name not in _known_backends:\n raise RuntimeError(f\"Unknown framework: {backend_name}\")\n\n return parsed_backends\n" \ - with=" if ENVVAR_NAME in os.environ:\n parsed_backends = os.environ[ENVVAR_NAME].split(\",\")\n _known_backends = find_names_of_all_frameworks()\n for backend_name in parsed_backends:\n if backend_name not in _known_backends:\n raise RuntimeError(f\"Unknown framework: {backend_name}\")\n return parsed_backends\n\n # Auto-detect available frameworks.\n available = []\n for module, name in [(\"torch\", \"torch\"), (\"jax\", \"jax\"), (\"numpy\", \"numpy\"), (\"tensorflow\", \"tensorflow\")]:\n try:\n __import__(module)\n available.append(name)\n except ImportError:\n pass\n return available\n" - -replace \ - in="tests/test_einsum.py" \ - count=1 \ - find="from einops.tests import collect_test_backends\nfrom einops.einops import _compactify_pattern_for_einsum, einsum, EinopsError\n" \ - with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\n_compactify_pattern_for_einsum = einops.einops._compactify_pattern_for_einsum\neinsum = einops.einops.einsum\nEinopsError = einops.EinopsError\nfrom . import collect_test_backends\n" - -replace \ - in="tests/test_examples.py" \ - count=1 \ - find="from einops import rearrange, parse_shape, reduce\nfrom einops.tests import is_backend_tested\nfrom einops.tests.test_ops import imp_op_backends\n" \ - with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nrearrange = einops.rearrange\nparse_shape = einops.parse_shape\nreduce = einops.reduce\nfrom . import is_backend_tested\nfrom .test_ops import imp_op_backends\n" - -replace \ - in="tests/test_layers.py" \ - count=1 \ - find="from einops import rearrange, reduce, EinopsError\nfrom einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS\n" \ - with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nrearrange = einops.rearrange\nreduce = einops.reduce\nEinopsError = einops.EinopsError\nfrom . import collect_test_backends, is_backend_tested\nfrom . import FLOAT_REDUCTIONS as REDUCTIONS\n" - -replace \ - in="tests/test_layers.py" \ - count=1 \ - find=" from einops.layers.torch import Rearrange, Reduce, EinMix\n" \ - with=" Rearrange, Reduce, EinMix = einops.layers.torch.Rearrange, einops.layers.torch.Reduce, einops.layers.torch.EinMix\n" - -replace \ - in="tests/test_layers.py" \ - count=1 \ - find=" from einops.layers.keras import Rearrange, Reduce, EinMix, keras_custom_objects\n" \ - with=" Rearrange, Reduce, EinMix, keras_custom_objects = einops.layers.keras.Rearrange, einops.layers.keras.Reduce, einops.layers.keras.EinMix, einops.layers.keras.keras_custom_objects\n" - -replace \ - in="tests/test_layers.py" \ - count=1 \ - find=" from einops.layers.flax import EinMix, Reduce, Rearrange\n" \ - with=" EinMix, Reduce, Rearrange = einops.layers.flax.EinMix, einops.layers.flax.Reduce, einops.layers.flax.Rearrange\n" - -replace \ - in="tests/test_layers.py" \ - count=2 \ - find=" from einops.layers._einmix import _EinmixDebugger\n" \ - with=" _EinmixDebugger = einops.layers._einmix._EinmixDebugger\n" +replace in="tests/__init__.py" count=1 \ + find=#""" + 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: + raise RuntimeError(f"Unknown framework: {backend_name}") + + return parsed_backends +"""# \ + with=#""" + if ENVVAR_NAME in os.environ: + 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: + raise RuntimeError(f"Unknown framework: {backend_name}") + return parsed_backends + + # Auto-detect available frameworks. + available = [] + for module, name in [("torch", "torch"), ("jax", "jax"), ("numpy", "numpy"), ("tensorflow", "tensorflow")]: + try: + __import__(module) + available.append(name) + except ImportError: + pass + return available +"""# + +replace in="tests/test_einsum.py" count=1 \ + find=#""" +from einops.tests import collect_test_backends +from einops.einops import _compactify_pattern_for_einsum, einsum, EinopsError +"""# \ + with=#""" +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +_compactify_pattern_for_einsum = einops.einops._compactify_pattern_for_einsum +einsum = einops.einops.einsum +EinopsError = einops.EinopsError +from . import collect_test_backends +"""# + +replace in="tests/test_examples.py" count=1 \ + find=#""" +from einops import rearrange, parse_shape, reduce +from einops.tests import is_backend_tested +from einops.tests.test_ops import imp_op_backends +"""# \ + with=#""" +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +rearrange = einops.rearrange +parse_shape = einops.parse_shape +reduce = einops.reduce +from . import is_backend_tested +from .test_ops import imp_op_backends +"""# + +replace in="tests/test_layers.py" count=1 \ + find=#""" +from einops import rearrange, reduce, EinopsError +from einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS +"""# \ + with=#""" +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +rearrange = einops.rearrange +reduce = einops.reduce +EinopsError = einops.EinopsError +from . import collect_test_backends, is_backend_tested +from . import FLOAT_REDUCTIONS as REDUCTIONS +"""# + +replace in="tests/test_layers.py" count=1 \ + find=" from einops.layers.torch import Rearrange, Reduce, EinMix" \ + with=" Rearrange, Reduce, EinMix = einops.layers.torch.Rearrange, einops.layers.torch.Reduce, einops.layers.torch.EinMix" + +replace in="tests/test_layers.py" count=1 \ + find=" from einops.layers.keras import Rearrange, Reduce, EinMix, keras_custom_objects" \ + with=" Rearrange, Reduce, EinMix, keras_custom_objects = einops.layers.keras.Rearrange, einops.layers.keras.Reduce, einops.layers.keras.EinMix, einops.layers.keras.keras_custom_objects" + +replace in="tests/test_layers.py" count=1 \ + find=" from einops.layers.flax import EinMix, Reduce, Rearrange" \ + with=" EinMix, Reduce, Rearrange = einops.layers.flax.EinMix, einops.layers.flax.Reduce, einops.layers.flax.Rearrange" + +replace in="tests/test_layers.py" count=2 \ + find=" from einops.layers._einmix import _EinmixDebugger" \ + with=" _EinmixDebugger = einops.layers._einmix._EinmixDebugger" // numpy 2 dropped in1d. -replace \ - in="tests/test_layers.py" \ - count=1 \ +replace in="tests/test_layers.py" count=1 \ find="numpy.in1d(grad, [0, 1])" \ with="numpy.isin(grad, [0, 1])" -replace \ - in="tests/test_ops.py" \ - count=1 \ - find="from einops import EinopsError\nfrom einops.einops import rearrange, reduce, repeat, _enumerate_directions\nfrom einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS\n" \ - with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nEinopsError = einops.EinopsError\nrearrange = einops.einops.rearrange\nreduce = einops.einops.reduce\nrepeat = einops.einops.repeat\n_enumerate_directions = einops.einops._enumerate_directions\nfrom . import collect_test_backends, is_backend_tested\nfrom . import FLOAT_REDUCTIONS as REDUCTIONS\n" - -replace \ - in="tests/test_ops.py" \ - count=1 \ - find=" from einops._backends import TorchBackend, TensorflowBackend, TFKerasBackend, NumpyBackend\n" \ - with=" TorchBackend = einops._backends.TorchBackend\n TensorflowBackend = einops._backends.TensorflowBackend\n TFKerasBackend = einops._backends.TFKerasBackend\n NumpyBackend = einops._backends.NumpyBackend\n" - -replace \ - in="tests/test_other.py" \ - count=1 \ - find="import einops\nimport einops.layers\nimport einops.parsing\nfrom einops._backends import AbstractBackend\nfrom einops.einops import rearrange, parse_shape, _optimize_transformation\nfrom einops.tests import collect_test_backends, is_backend_tested\n" \ - with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nAbstractBackend = einops._backends.AbstractBackend\nrearrange = einops.einops.rearrange\nparse_shape = einops.einops.parse_shape\n_optimize_transformation = einops.einops._optimize_transformation\nfrom . import collect_test_backends, is_backend_tested\n" - -replace \ - in="tests/test_other.py" \ - count=1 \ - find=" from einops import repeat, reduce, pack, unpack, einsum\n from einops._torch_specific import allow_ops_in_compiled_graph\n" \ - with=" repeat, reduce, pack, unpack, einsum = einops.repeat, einops.reduce, einops.pack, einops.unpack, einops.einsum\n allow_ops_in_compiled_graph = einops._torch_specific.allow_ops_in_compiled_graph\n" - -replace \ - in="tests/test_packing.py" \ - count=1 \ - find="from einops import EinopsError, asnumpy, pack, unpack\nfrom einops.tests import collect_test_backends\n" \ - with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nEinopsError = einops.EinopsError\nasnumpy = einops.asnumpy\npack = einops.pack\nunpack = einops.unpack\nfrom . import collect_test_backends\n" - -replace \ - in="tests/test_parsing.py" \ - count=1 \ - find="from einops import EinopsError\nfrom einops.parsing import ParsedExpression, AnonymousAxis, _ellipsis\n" \ - with="import kernels\neinops = kernels.get_kernel(\"kernels-community/einops\", version=1)\nEinopsError = einops.EinopsError\nParsedExpression = einops.parsing.ParsedExpression\nAnonymousAxis = einops.parsing.AnonymousAxis\n_ellipsis = einops.parsing._ellipsis\n" +replace in="tests/test_ops.py" count=1 \ + find=#""" +from einops import EinopsError +from einops.einops import rearrange, reduce, repeat, _enumerate_directions +from einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS +"""# \ + with=#""" +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +rearrange = einops.einops.rearrange +reduce = einops.einops.reduce +repeat = einops.einops.repeat +_enumerate_directions = einops.einops._enumerate_directions +from . import collect_test_backends, is_backend_tested +from . import FLOAT_REDUCTIONS as REDUCTIONS +"""# + +replace in="tests/test_ops.py" count=1 \ + find=" from einops._backends import TorchBackend, TensorflowBackend, TFKerasBackend, NumpyBackend" \ + with=#""" + TorchBackend = einops._backends.TorchBackend + TensorflowBackend = einops._backends.TensorflowBackend + TFKerasBackend = einops._backends.TFKerasBackend + NumpyBackend = einops._backends.NumpyBackend +"""# + +replace in="tests/test_other.py" count=1 \ + find=#""" +import einops +import einops.layers +import einops.parsing +from einops._backends import AbstractBackend +from einops.einops import rearrange, parse_shape, _optimize_transformation +from einops.tests import collect_test_backends, is_backend_tested +"""# \ + with=#""" +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +AbstractBackend = einops._backends.AbstractBackend +rearrange = einops.einops.rearrange +parse_shape = einops.einops.parse_shape +_optimize_transformation = einops.einops._optimize_transformation +from . import collect_test_backends, is_backend_tested +"""# + +replace in="tests/test_other.py" count=1 \ + find=#""" + from einops import repeat, reduce, pack, unpack, einsum + from einops._torch_specific import allow_ops_in_compiled_graph +"""# \ + with=#""" + repeat, reduce, pack, unpack, einsum = einops.repeat, einops.reduce, einops.pack, einops.unpack, einops.einsum + allow_ops_in_compiled_graph = einops._torch_specific.allow_ops_in_compiled_graph +"""# + +replace in="tests/test_packing.py" count=1 \ + find=#""" +from einops import EinopsError, asnumpy, pack, unpack +from einops.tests import collect_test_backends +"""# \ + with=#""" +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +asnumpy = einops.asnumpy +pack = einops.pack +unpack = einops.unpack +from . import collect_test_backends +"""# + +replace in="tests/test_parsing.py" count=1 \ + find=#""" +from einops import EinopsError +from einops.parsing import ParsedExpression, AnonymousAxis, _ellipsis +"""# \ + with=#""" +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +ParsedExpression = einops.parsing.ParsedExpression +AnonymousAxis = einops.parsing.AnonymousAxis +_ellipsis = einops.parsing._ellipsis +"""# // array_api is reached as an attribute of the loaded kernel (see the // __init__.py edit above), not as a submodule import. -replace \ - in="tests/**" \ - count=4 \ - find=" from einops import array_api as AA\n" \ - with=" AA = einops.array_api\n" +replace in="tests/**" count=4 \ + find=" from einops import array_api as AA" \ + with=" AA = einops.array_api" // Kernel CI selects tests by marker; upstream has none. diff --git a/einops/src/port-provenance.json b/einops/src/port-provenance.json index f0a2f76fc..f9c32e78d 100644 --- a/einops/src/port-provenance.json +++ b/einops/src/port-provenance.json @@ -1,6 +1,6 @@ { "format": 1, - "recipe": "sha256:7b9b97f928ded68c67302ce872b875a604abd35d2be6543730a20a38da97bf47", + "recipe": "sha256:6da89bfee0093cbe1e8f231bfb2784d32e24a62355feae5449b4dbfe5a8de84a", "runner": "kernel-port 0.1.0", "sources": { "upstream": { From dbad83e6d2be28596b2c801176e762972dfb7891 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Wed, 26 Aug 2026 17:51:09 -0400 Subject: [PATCH 05/15] feat: move the einops flake out of the port overlay Signed-off-by: David Holtz --- einops/{src => }/flake.lock | 0 einops/{port/overlay => }/flake.nix | 2 +- einops/port/overlay/flake.lock | 117 ---------------------------- einops/port/port.kdl | 7 +- einops/src/flake.nix | 17 ---- einops/src/port-provenance.json | 4 +- 6 files changed, 6 insertions(+), 141 deletions(-) rename einops/{src => }/flake.lock (100%) rename einops/{port/overlay => }/flake.nix (92%) delete mode 100644 einops/port/overlay/flake.lock delete mode 100644 einops/src/flake.nix diff --git a/einops/src/flake.lock b/einops/flake.lock similarity index 100% rename from einops/src/flake.lock rename to einops/flake.lock diff --git a/einops/port/overlay/flake.nix b/einops/flake.nix similarity index 92% rename from einops/port/overlay/flake.nix rename to einops/flake.nix index 0b2a92f1e..b7b6ea3ec 100644 --- a/einops/port/overlay/flake.nix +++ b/einops/flake.nix @@ -12,6 +12,6 @@ }: kernel-builder.lib.genKernelFlakeOutputs { inherit self; - path = ./.; + path = ./src; }; } diff --git a/einops/port/overlay/flake.lock b/einops/port/overlay/flake.lock deleted file mode 100644 index 1446e6473..000000000 --- a/einops/port/overlay/flake.lock +++ /dev/null @@ -1,117 +0,0 @@ -{ - "nodes": { - "flake-compat": { - "locked": { - "lastModified": 1767039857, - "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=", - "owner": "edolstra", - "repo": "flake-compat", - "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab", - "type": "github" - }, - "original": { - "owner": "edolstra", - "repo": "flake-compat", - "type": "github" - } - }, - "flake-utils": { - "inputs": { - "systems": "systems" - }, - "locked": { - "lastModified": 1731533236, - "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=", - "owner": "numtide", - "repo": "flake-utils", - "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b", - "type": "github" - }, - "original": { - "owner": "numtide", - "repo": "flake-utils", - "type": "github" - } - }, - "kernel-builder": { - "inputs": { - "flake-compat": "flake-compat", - "flake-utils": "flake-utils", - "nixpkgs": "nixpkgs", - "rust-overlay": "rust-overlay" - }, - "locked": { - "lastModified": 1785427612, - "narHash": "sha256-hUFEwQKDjVHw9v6xgC6Lx402TYfUA8zXEi8hDE/F4y0=", - "owner": "huggingface", - "repo": "kernels", - "rev": "d0610aa58db33b142c86b59598a2a1c730f52996", - "type": "github" - }, - "original": { - "owner": "huggingface", - "repo": "kernels", - "type": "github" - } - }, - "nixpkgs": { - "locked": { - "lastModified": 1783284758, - "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=", - "owner": "NixOS", - "repo": "nixpkgs", - "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", - "type": "github" - }, - "original": { - "owner": "NixOS", - "repo": "nixpkgs", - "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d", - "type": "github" - } - }, - "root": { - "inputs": { - "kernel-builder": "kernel-builder" - } - }, - "rust-overlay": { - "inputs": { - "nixpkgs": [ - "kernel-builder", - "nixpkgs" - ] - }, - "locked": { - "lastModified": 1783320166, - "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=", - "owner": "oxalica", - "repo": "rust-overlay", - "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673", - "type": "github" - }, - "original": { - "owner": "oxalica", - "repo": "rust-overlay", - "type": "github" - } - }, - "systems": { - "locked": { - "lastModified": 1681028828, - "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=", - "owner": "nix-systems", - "repo": "default", - "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e", - "type": "github" - }, - "original": { - "owner": "nix-systems", - "repo": "default", - "type": "github" - } - } - }, - "root": "root", - "version": 7 -} diff --git a/einops/port/port.kdl b/einops/port/port.kdl index faa5903dd..b6e269361 100644 --- a/einops/port/port.kdl +++ b/einops/port/port.kdl @@ -10,7 +10,9 @@ recipe version=1 // by name; the port rebinds those imports to the kernel loaded off the Hub and // marks each test for kernel CI. Everything else is upstream's. // -// The flake has no upstream equivalent, so it comes from the overlay. +// The flake is hand-written and lives at einops/, outside --out, so it +// survives regeneration. Its `path = ./src` points kernel-builder at the +// generated tree. // // Multi-line payloads below use KDL raw strings. KDL dedents them against the // closing delimiter, so every "\"\"\"# sits at column 0 and the payload keeps @@ -263,9 +265,6 @@ replace \ find="\ndef test_" \ with="\n@pytest.mark.kernels_ci\ndef test_" -// The flake: hand-written, no upstream equivalent. -overlay from="overlay" - expect \ in="torch-ext/einops/**" \ files=18 diff --git a/einops/src/flake.nix b/einops/src/flake.nix deleted file mode 100644 index 0b2a92f1e..000000000 --- a/einops/src/flake.nix +++ /dev/null @@ -1,17 +0,0 @@ -{ - description = "Flake for eigen kernel"; - - inputs = { - kernel-builder.url = "github:huggingface/kernels"; - }; - - outputs = - { - self, - kernel-builder, - }: - kernel-builder.lib.genKernelFlakeOutputs { - inherit self; - path = ./.; - }; -} diff --git a/einops/src/port-provenance.json b/einops/src/port-provenance.json index f9c32e78d..2524463f8 100644 --- a/einops/src/port-provenance.json +++ b/einops/src/port-provenance.json @@ -1,6 +1,6 @@ { "format": 1, - "recipe": "sha256:6da89bfee0093cbe1e8f231bfb2784d32e24a62355feae5449b4dbfe5a8de84a", + "recipe": "sha256:95d664ac2005149af525538d155557fb0066342fbc87c99058838ea65a991f10", "runner": "kernel-port 0.1.0", "sources": { "upstream": { @@ -8,5 +8,5 @@ "commit": "ad2c8d67d08cbebb45e1aff2be21dc88635fc695" } }, - "tree": "sha256:f7dab98a4d28b88488879056c333147c7981814aa777031aefac7a2b8c32c1b9" + "tree": "sha256:c617d65fd27e3cd928acb157a66cd3d765ccc8c7987f9b990782adaf1d2ef6dc" } From 7193964c8ce1c01d5b1d1769c231696c51d77852 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Wed, 26 Aug 2026 19:47:04 -0400 Subject: [PATCH 06/15] feat: verify ported kernels regenerate from their recipe in CI Signed-off-by: David Holtz --- .gitattributes | 8 ++ .github/workflows/check-ports.yaml | 41 ++++++++ .gitignore | 1 + scripts/check_ports.py | 157 +++++++++++++++++++++++++++++ 4 files changed, 207 insertions(+) create mode 100644 .gitattributes create mode 100644 .github/workflows/check-ports.yaml create mode 100644 scripts/check_ports.py diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 000000000..33aee19ec --- /dev/null +++ b/.gitattributes @@ -0,0 +1,8 @@ +# Ported kernels are generated: /port/port.kdl plus /port/overlay +# are the inputs, and /src is kernel-port's output, committed so that +# what ships stays reviewable and auditable. +# +# Collapse the generated tree in diffs so review lands on the inputs. The output +# is not taken on trust: .github/workflows/check-ports.yaml regenerates every +# port and fails if /src does not match byte for byte. +*/src/** linguist-generated=true diff --git a/.github/workflows/check-ports.yaml b/.github/workflows/check-ports.yaml new file mode 100644 index 000000000..bdaf84b7e --- /dev/null +++ b/.github/workflows/check-ports.yaml @@ -0,0 +1,41 @@ +name: Check Ports + +# A ported kernel's /src is generated from /port/port.kdl and +# committed. Regenerate it here and fail if it drifted, so that a stale overlay +# silently reverting an in-tree fix, a hand-edited build.toml field the manifest +# op drops, or an overlay file edited without its generated twin are all loud. +on: + push: + branches: [main] + pull_request: + branches: [main] + types: [opened, synchronize, reopened] + paths: + - "*/port/**" + - "*/src/**" + - "scripts/check_ports.py" + - ".github/workflows/check-ports.yaml" + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }} + cancel-in-progress: true + +permissions: + contents: read + +jobs: + check-ports: + name: Check ports reproduce + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + + - uses: DeterminateSystems/nix-installer-action@ef8a148080ab6020fd15196c2084a2eea5ff2d25 # v22 + + # Each recipe pins its own upstream commit; the runner verifies the + # origin URL, HEAD and cleanliness of the checkout it is handed. + - name: Regenerate every port and compare + env: + KERNEL_PORT: "nix run github:huggingface/kernels#kernel-port --" + run: python3 scripts/check_ports.py diff --git a/.gitignore b/.gitignore index 4dcaf8670..884823ecb 100644 --- a/.gitignore +++ b/.gitignore @@ -15,6 +15,7 @@ build # Editor files .* !.github/ +!.gitattributes # Python stuff __pycache__/ diff --git a/scripts/check_ports.py b/scripts/check_ports.py new file mode 100644 index 000000000..a60a7cadf --- /dev/null +++ b/scripts/check_ports.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +"""Verify that every ported kernel's committed tree matches its recipe. + +A ported kernel is a directory holding `port/port.kdl`. The recipe generates +`/src`, which is committed so that what ships stays reviewable. This +script re-runs the port against a fresh upstream checkout at the pinned commit +and fails if the result differs from what is committed. + +Without this, three failure modes are silent: + + - an overlay file goes stale, so regenerating reverts a fix made in-tree + - a build.toml field is edited by hand, and the manifest op drops it + - an overlay file and its generated twin are edited independently and drift + +Usage: + python3 scripts/check_ports.py [kernel ...] + +Defaults to every port found. Set KERNEL_PORT to override the runner command +(e.g. "nix run github:huggingface/kernels#kernel-port --"). +""" + +import filecmp +import os +import re +import shlex +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +# `source repo="..." commit="..."` is the recipe's first op; the runner already +# enforces a 40-char sha and a clean checkout, so a loose match is enough here. +SOURCE_RE = re.compile( + r"^source\b(?P(?:[^\n\\]|\\\s*\n)*)", re.MULTILINE +) +ARG_RE = re.compile(r'(\w+)="([^"]*)"') + +DEFAULT_RUNNER = "kernel-port" + + +def runner_cmd() -> list: + return shlex.split(os.environ.get("KERNEL_PORT", DEFAULT_RUNNER)) + + +def find_ports(root: Path) -> list: + return sorted(p.parent.parent.name for p in root.glob("*/port/port.kdl")) + + +def parse_source(recipe: Path) -> tuple: + text = recipe.read_text() + m = SOURCE_RE.search(text) + if not m: + raise SystemExit(f"{recipe}: no `source` op found") + args = dict(ARG_RE.findall(m.group("body").replace("\\\n", " "))) + try: + return args["repo"], args["commit"] + except KeyError as exc: + raise SystemExit(f"{recipe}: source is missing {exc}") from exc + + +def checkout(repo: str, commit: str, dest: Path) -> None: + # The `source` op verifies origin URL, HEAD, and cleanliness, so clone from + # the exact URL the recipe names and land on the exact commit. + subprocess.run( + ["git", "clone", "--quiet", "--filter=blob:none", "--no-checkout", repo, str(dest)], + check=True, + ) + subprocess.run(["git", "-C", str(dest), "checkout", "--quiet", commit], check=True) + + +def tree_files(root: Path) -> set: + return { + p.relative_to(root).as_posix() + for p in root.rglob("*") + if p.is_file() and "__pycache__" not in p.parts + } + + +def compare(generated: Path, committed: Path) -> list: + gen, com = tree_files(generated), tree_files(committed) + problems = [] + for rel in sorted(com - gen): + problems.append(f"committed but not generated: {rel}") + for rel in sorted(gen - com): + problems.append(f"generated but not committed: {rel}") + for rel in sorted(gen & com): + if not filecmp.cmp(generated / rel, committed / rel, shallow=False): + problems.append(f"differs: {rel}") + return problems + + +def check(kernel: str, root: Path, workdir: Path) -> list: + recipe = root / kernel / "port" / "port.kdl" + repo, commit = parse_source(recipe) + print(f" {kernel}: {repo} @ {commit[:12]}") + + # One checkout per (repo, commit): two kernels may share an upstream. + src = workdir / f"upstream-{abs(hash((repo, commit)))}" + if not src.exists(): + checkout(repo, commit, src) + + out = workdir / f"out-{kernel}" + shutil.rmtree(out, ignore_errors=True) + proc = subprocess.run( + [*runner_cmd(), str(recipe), "--dir", str(src), "--out", str(out)], + capture_output=True, + text=True, + ) + if proc.returncode != 0: + tail = (proc.stderr or proc.stdout).strip().splitlines()[-8:] + return [f"kernel-port failed:\n " + "\n ".join(tail)] + + return compare(out, root / kernel / "src") + + +def main() -> int: + root = Path.cwd() + kernels = sys.argv[1:] or find_ports(root) + if not kernels: + print("No ported kernels found (no */port/port.kdl).") + return 0 + + print(f"Checking {len(kernels)} port(s) reproduce their committed tree:\n") + failed = {} + with tempfile.TemporaryDirectory() as tmp: + workdir = Path(tmp) + for kernel in kernels: + if not (root / kernel / "port" / "port.kdl").is_file(): + failed[kernel] = [f"no port/port.kdl under {kernel}/"] + continue + problems = check(kernel, root, workdir) + if problems: + failed[kernel] = problems + else: + print(" ok - regenerates byte for byte") + + if failed: + print("\nERROR: committed tree does not match the recipe.\n", file=sys.stderr) + for kernel, problems in failed.items(): + print(f" {kernel}:", file=sys.stderr) + for p in problems: + print(f" - {p}", file=sys.stderr) + print( + "\nRegenerate with:\n" + " kernel-port /port/port.kdl --dir --out /src\n" + "If an overlay file went stale, refresh it from /src first.", + file=sys.stderr, + ) + return 1 + + print("\nAll ports reproduce their committed tree.") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) From 48c52d009dad51796832839a14c65fcc32c52bf4 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Wed, 26 Aug 2026 19:55:57 -0400 Subject: [PATCH 07/15] fix: point the port check at the kernel-port-tool branch Signed-off-by: David Holtz --- .github/workflows/check-ports.yaml | 4 +++- .gitignore | 2 +- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/check-ports.yaml b/.github/workflows/check-ports.yaml index bdaf84b7e..1d9b29138 100644 --- a/.github/workflows/check-ports.yaml +++ b/.github/workflows/check-ports.yaml @@ -37,5 +37,7 @@ jobs: # origin URL, HEAD and cleanliness of the checkout it is handed. - name: Regenerate every port and compare env: - KERNEL_PORT: "nix run github:huggingface/kernels#kernel-port --" + # kernel-port still lives on the kernel-port-tool branch; point this at + # github:huggingface/kernels#kernel-port once that branch merges. + KERNEL_PORT: "nix run github:huggingface/kernels/kernel-port-tool#kernel-port --" run: python3 scripts/check_ports.py diff --git a/.gitignore b/.gitignore index 884823ecb..c9216b6cf 100644 --- a/.gitignore +++ b/.gitignore @@ -15,7 +15,7 @@ build # Editor files .* !.github/ -!.gitattributes +!/.gitattributes # Python stuff __pycache__/ From 4a6e1bfae1a62a52abe0c012eb1248dfb32dcc41 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Thu, 27 Aug 2026 15:34:31 -0400 Subject: [PATCH 08/15] feat: use ensure_import and kernelize_imports to make recipe much more concise --- einops/port/port.kdl | 307 +++++------------------------- einops/src/port-provenance.json | 4 +- einops/src/tests/__init__.py | 40 ++-- einops/src/tests/test_einsum.py | 22 ++- einops/src/tests/test_examples.py | 20 +- einops/src/tests/test_layers.py | 41 ++-- einops/src/tests/test_ops.py | 59 ++---- einops/src/tests/test_other.py | 43 ++--- einops/src/tests/test_packing.py | 25 +-- einops/src/tests/test_parsing.py | 24 ++- 10 files changed, 181 insertions(+), 404 deletions(-) diff --git a/einops/port/port.kdl b/einops/port/port.kdl index b6e269361..84aeb429b 100644 --- a/einops/port/port.kdl +++ b/einops/port/port.kdl @@ -1,287 +1,74 @@ recipe version=1 -// einops -// -// einops as a Hub kernel: the package is pure Python, so the manifest is -// noarch and the port is a move plus the import rewrite that lets the package -// be imported under whatever directory name the Hub loads it from. -// -// The test suite is derived from upstream too. Upstream's tests import einops -// by name; the port rebinds those imports to the kernel loaded off the Hub and -// marks each test for kernel CI. Everything else is upstream's. -// -// The flake is hand-written and lives at einops/, outside --out, so it -// survives regeneration. Its `path = ./src` points kernel-builder at the -// generated tree. -// -// Multi-line payloads below use KDL raw strings. KDL dedents them against the -// closing delimiter, so every "\"\"\"# sits at column 0 and the payload keeps -// the exact indentation it must match in the source. -// -// kernel-port einops/port/port.kdl \ -// --dir --out einops/src - +// Pin the exact clean upstream tree this recipe transforms. source \ repo="https://github.com/arogozhnikov/einops.git" \ commit="ad2c8d67d08cbebb45e1aff2be21dc88635fc695" +// Retain only the Python package and its upstream tests. prune keep="einops/**" +// Lift upstream tests out of the package into the kernel test root. move \ from="einops/tests" \ to="tests" + +// Place the distributable package in the torch extension tree. move \ from="einops" \ to="torch-ext/einops" +// Make package-internal imports work under the Hub-generated module name. relativize_imports \ in="torch-ext/einops/**" \ - package_root="torch-ext/einops" \ - changes=7 -expect \ - in="torch-ext/einops/**" \ - find="from einops" \ - count=0 - -// array_api is a submodule upstream leaves out of the package init; the kernel -// exposes it, since a Hub consumer cannot import einops.array_api by name. -replace in="torch-ext/einops/__init__.py" count=1 \ - find="from .packing import pack, unpack" \ - with=#""" -from .packing import pack, unpack -from . import array_api -"""# - -// The test suite imports the package by name; the kernel is loaded off the -// Hub instead, so each import site is rebound to the loaded module. - -replace in="tests/__init__.py" count=1 \ - find="from einops import _backends" \ - with=#""" -import kernels -_einops = kernels.get_kernel("kernels-community/einops", version=1) -_backends = _einops._backends -"""# - -// Upstream hard-fails when EINOPS_TEST_BACKENDS is unset. In kernel CI no one -// sets it, so fall back to whichever frameworks actually import. - -replace in="tests/__init__.py" count=1 \ - find=#""" - 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: - raise RuntimeError(f"Unknown framework: {backend_name}") - - return parsed_backends -"""# \ - with=#""" - if ENVVAR_NAME in os.environ: - 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: - raise RuntimeError(f"Unknown framework: {backend_name}") - return parsed_backends - - # Auto-detect available frameworks. - available = [] - for module, name in [("torch", "torch"), ("jax", "jax"), ("numpy", "numpy"), ("tensorflow", "tensorflow")]: - try: - __import__(module) - available.append(name) - except ImportError: - pass - return available -"""# - -replace in="tests/test_einsum.py" count=1 \ - find=#""" -from einops.tests import collect_test_backends -from einops.einops import _compactify_pattern_for_einsum, einsum, EinopsError -"""# \ - with=#""" -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -_compactify_pattern_for_einsum = einops.einops._compactify_pattern_for_einsum -einsum = einops.einops.einsum -EinopsError = einops.EinopsError -from . import collect_test_backends -"""# - -replace in="tests/test_examples.py" count=1 \ - find=#""" -from einops import rearrange, parse_shape, reduce -from einops.tests import is_backend_tested -from einops.tests.test_ops import imp_op_backends -"""# \ - with=#""" -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -rearrange = einops.rearrange -parse_shape = einops.parse_shape -reduce = einops.reduce -from . import is_backend_tested -from .test_ops import imp_op_backends -"""# - -replace in="tests/test_layers.py" count=1 \ - find=#""" -from einops import rearrange, reduce, EinopsError -from einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS -"""# \ - with=#""" -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -rearrange = einops.rearrange -reduce = einops.reduce -EinopsError = einops.EinopsError -from . import collect_test_backends, is_backend_tested -from . import FLOAT_REDUCTIONS as REDUCTIONS -"""# - -replace in="tests/test_layers.py" count=1 \ - find=" from einops.layers.torch import Rearrange, Reduce, EinMix" \ - with=" Rearrange, Reduce, EinMix = einops.layers.torch.Rearrange, einops.layers.torch.Reduce, einops.layers.torch.EinMix" - -replace in="tests/test_layers.py" count=1 \ - find=" from einops.layers.keras import Rearrange, Reduce, EinMix, keras_custom_objects" \ - with=" Rearrange, Reduce, EinMix, keras_custom_objects = einops.layers.keras.Rearrange, einops.layers.keras.Reduce, einops.layers.keras.EinMix, einops.layers.keras.keras_custom_objects" - -replace in="tests/test_layers.py" count=1 \ - find=" from einops.layers.flax import EinMix, Reduce, Rearrange" \ - with=" EinMix, Reduce, Rearrange = einops.layers.flax.EinMix, einops.layers.flax.Reduce, einops.layers.flax.Rearrange" - -replace in="tests/test_layers.py" count=2 \ - find=" from einops.layers._einmix import _EinmixDebugger" \ - with=" _EinmixDebugger = einops.layers._einmix._EinmixDebugger" - -// numpy 2 dropped in1d. + package_root="torch-ext/einops" + +// Expose array_api because Hub consumers cannot import it by package name. +ensure_import \ + in="torch-ext/einops/__init__.py" \ + from="." \ + name="array_api" + +// Point imports of the former einops.tests package at the lifted tests. +remap_module \ + in="tests/{__init__,test_*}.py" \ + from="einops.tests" \ + to="tests" -replace in="tests/test_layers.py" count=1 \ +// Make imports between lifted test modules package-relative. +relativize_imports \ + in="tests/{__init__,test_*}.py" \ + package_root="tests" + +// Resolve product imports through the locally built Hub kernel. +kernelize_imports \ + in="tests/{__init__,test_*}.py" \ + package="einops" \ + kernel="kernels-community/einops" \ + version=1 + +// Default to the two Kernel CI backends when upstream's env var is absent. +replace in="tests/__init__.py" \ + find=#" raise RuntimeError(f"Testing frameworks were not specified, env var {ENVVAR_NAME} not set")"# \ + with=#" return ["torch", "numpy"]"# + +// Use the supported replacement for NumPy 2's removed in1d API. +replace in="tests/test_layers.py" \ find="numpy.in1d(grad, [0, 1])" \ with="numpy.isin(grad, [0, 1])" -replace in="tests/test_ops.py" count=1 \ - find=#""" -from einops import EinopsError -from einops.einops import rearrange, reduce, repeat, _enumerate_directions -from einops.tests import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS -"""# \ - with=#""" -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -rearrange = einops.einops.rearrange -reduce = einops.einops.reduce -repeat = einops.einops.repeat -_enumerate_directions = einops.einops._enumerate_directions -from . import collect_test_backends, is_backend_tested -from . import FLOAT_REDUCTIONS as REDUCTIONS -"""# - -replace in="tests/test_ops.py" count=1 \ - find=" from einops._backends import TorchBackend, TensorflowBackend, TFKerasBackend, NumpyBackend" \ - with=#""" - TorchBackend = einops._backends.TorchBackend - TensorflowBackend = einops._backends.TensorflowBackend - TFKerasBackend = einops._backends.TFKerasBackend - NumpyBackend = einops._backends.NumpyBackend -"""# - -replace in="tests/test_other.py" count=1 \ - find=#""" -import einops -import einops.layers -import einops.parsing -from einops._backends import AbstractBackend -from einops.einops import rearrange, parse_shape, _optimize_transformation -from einops.tests import collect_test_backends, is_backend_tested -"""# \ - with=#""" -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -AbstractBackend = einops._backends.AbstractBackend -rearrange = einops.einops.rearrange -parse_shape = einops.einops.parse_shape -_optimize_transformation = einops.einops._optimize_transformation -from . import collect_test_backends, is_backend_tested -"""# - -replace in="tests/test_other.py" count=1 \ - find=#""" - from einops import repeat, reduce, pack, unpack, einsum - from einops._torch_specific import allow_ops_in_compiled_graph -"""# \ - with=#""" - repeat, reduce, pack, unpack, einsum = einops.repeat, einops.reduce, einops.pack, einops.unpack, einops.einsum - allow_ops_in_compiled_graph = einops._torch_specific.allow_ops_in_compiled_graph -"""# - -replace in="tests/test_packing.py" count=1 \ - find=#""" -from einops import EinopsError, asnumpy, pack, unpack -from einops.tests import collect_test_backends -"""# \ - with=#""" -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -asnumpy = einops.asnumpy -pack = einops.pack -unpack = einops.unpack -from . import collect_test_backends -"""# - -replace in="tests/test_parsing.py" count=1 \ - find=#""" -from einops import EinopsError -from einops.parsing import ParsedExpression, AnonymousAxis, _ellipsis -"""# \ - with=#""" -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -ParsedExpression = einops.parsing.ParsedExpression -AnonymousAxis = einops.parsing.AnonymousAxis -_ellipsis = einops.parsing._ellipsis -"""# - -// array_api is reached as an attribute of the loaded kernel (see the -// __init__.py edit above), not as a submodule import. - -replace in="tests/**" count=4 \ - find=" from einops import array_api as AA" \ - with=" AA = einops.array_api" - -// Kernel CI selects tests by marker; upstream has none. - +// Select every upstream test module for the kernel CI suite. replace \ - in="tests/**" \ - count=63 \ - find="\ndef test_" \ - with="\n@pytest.mark.kernels_ci\ndef test_" - -expect \ - in="torch-ext/einops/**" \ - files=18 -expect \ - in="tests/**" \ - files=9 -// run_tests.py is upstream's standalone pip-install-and-run driver. Kernel CI -// does not use it, so it is shipped verbatim and excluded from this check. -expect \ - in="tests/__init__.py" \ - find="from einops" \ - count=0 -expect \ in="tests/test_*.py" \ - find="from einops" \ - count=0 + find="import pytest" \ + with="import pytest\n\npytestmark = pytest.mark.kernels_ci" + +// Keep pytest from collecting the imported doctest helper as a marked test. +replace in="tests/test_other.py" \ + find="from doctest import testmod" \ + with="from doctest import testmod\ntestmod.__test__ = False" +// Emit the noarch kernel manifest and supported runtime backends. manifest \ name="einops" \ version=1 \ diff --git a/einops/src/port-provenance.json b/einops/src/port-provenance.json index 2524463f8..d0ee2129b 100644 --- a/einops/src/port-provenance.json +++ b/einops/src/port-provenance.json @@ -1,6 +1,6 @@ { "format": 1, - "recipe": "sha256:95d664ac2005149af525538d155557fb0066342fbc87c99058838ea65a991f10", + "recipe": "sha256:553ea0a8342085862371639c2f12947c592617ffb7af4343b02a3048986a026c", "runner": "kernel-port 0.1.0", "sources": { "upstream": { @@ -8,5 +8,5 @@ "commit": "ad2c8d67d08cbebb45e1aff2be21dc88635fc695" } }, - "tree": "sha256:c617d65fd27e3cd928acb157a66cd3d765ccc8c7987f9b990782adaf1d2ef6dc" + "tree": "sha256:ed3299bb86f88fa6fc9fcf72551a64830b327b879f9a4da791d29981a4b840f4" } diff --git a/einops/src/tests/__init__.py b/einops/src/tests/__init__.py index b8c629217..977745afb 100644 --- a/einops/src/tests/__init__.py +++ b/einops/src/tests/__init__.py @@ -2,15 +2,23 @@ Common utils for testing. These functions allow testing only some frameworks, not all. """ +__kernel_port_einops_root = None +def __kernel_port_einops(module=""): + global __kernel_port_einops_root + if __kernel_port_einops_root is None: + __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) + root = __kernel_port_einops_root + if not module: + return root + return __import__("importlib").import_module(root.__name__ + "." + module) + import logging import os from functools import lru_cache from typing import List, Tuple -import kernels -_einops = kernels.get_kernel("kernels-community/einops", version=1) -_backends = _einops._backends +_backends = getattr(__kernel_port_einops(), "_backends") import warnings __author__ = "Alex Rogozhnikov" @@ -46,23 +54,15 @@ def unparse_backends(backend_names: List[str]) -> Tuple[str, str]: @lru_cache(maxsize=1) def parse_backends_to_test() -> List[str]: - if ENVVAR_NAME in os.environ: - 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: - raise RuntimeError(f"Unknown framework: {backend_name}") - return parsed_backends - - # Auto-detect available frameworks. - available = [] - for module, name in [("torch", "torch"), ("jax", "jax"), ("numpy", "numpy"), ("tensorflow", "tensorflow")]: - try: - __import__(module) - available.append(name) - except ImportError: - pass - return available + if ENVVAR_NAME not in os.environ: + return ["torch", "numpy"] + 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: + raise RuntimeError(f"Unknown framework: {backend_name}") + + return parsed_backends def is_backend_tested(backend: str) -> bool: diff --git a/einops/src/tests/test_einsum.py b/einops/src/tests/test_einsum.py index 2e05d939d..97a1257ef 100644 --- a/einops/src/tests/test_einsum.py +++ b/einops/src/tests/test_einsum.py @@ -1,12 +1,20 @@ +__kernel_port_einops_root = None +def __kernel_port_einops(module=""): + global __kernel_port_einops_root + if __kernel_port_einops_root is None: + __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) + root = __kernel_port_einops_root + if not module: + return root + return __import__("importlib").import_module(root.__name__ + "." + module) + from typing import Any, Callable -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -_compactify_pattern_for_einsum = einops.einops._compactify_pattern_for_einsum -einsum = einops.einops.einsum -EinopsError = einops.EinopsError from . import collect_test_backends +_compactify_pattern_for_einsum, einsum, EinopsError = getattr(__kernel_port_einops("einops"), "_compactify_pattern_for_einsum"), getattr(__kernel_port_einops("einops"), "einsum"), getattr(__kernel_port_einops("einops"), "EinopsError") import numpy as np import pytest + +pytestmark = pytest.mark.kernels_ci import string @@ -168,7 +176,6 @@ def __call__(self, function: Callable): ] -@pytest.mark.kernels_ci def test_layer(): for backend in collect_test_backends(layers=True, symbolic=False): if backend.framework_name in ["tensorflow", "torch", "oneflow", "paddle"]: @@ -196,7 +203,6 @@ def test_layer(): ] -@pytest.mark.kernels_ci def test_functional(): # Functional tests: backends = filter(lambda x: x.framework_name in valid_backends_functional, collect_test_backends()) @@ -232,7 +238,6 @@ def test_functional(): np.testing.assert_array_almost_equal(predicted_out_array, true_out_array, decimal=5) -@pytest.mark.kernels_ci def test_functional_symbolic(): backends = filter( lambda x: x.framework_name in valid_backends_functional, collect_test_backends(symbolic=True, layers=False) @@ -265,7 +270,6 @@ def test_functional_symbolic(): np.testing.assert_array_almost_equal(predicted_out_data, expected_out_data, decimal=5) -@pytest.mark.kernels_ci def test_functional_errors(): # Specific backend does not matter, as errors are raised # during the pattern creation. diff --git a/einops/src/tests/test_examples.py b/einops/src/tests/test_examples.py index 4a66431bc..8cda27f01 100644 --- a/einops/src/tests/test_examples.py +++ b/einops/src/tests/test_examples.py @@ -1,16 +1,23 @@ +__kernel_port_einops_root = None +def __kernel_port_einops(module=""): + global __kernel_port_einops_root + if __kernel_port_einops_root is None: + __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) + root = __kernel_port_einops_root + if not module: + return root + return __import__("importlib").import_module(root.__name__ + "." + module) + import numpy import pytest -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -rearrange = einops.rearrange -parse_shape = einops.parse_shape -reduce = einops.reduce +pytestmark = pytest.mark.kernels_ci + +rearrange, parse_shape, reduce = getattr(__kernel_port_einops(), "rearrange"), getattr(__kernel_port_einops(), "parse_shape"), getattr(__kernel_port_einops(), "reduce") from . import is_backend_tested from .test_ops import imp_op_backends -@pytest.mark.kernels_ci def test_rearrange_examples(): def test1(x): # transpose @@ -196,7 +203,6 @@ def tensor_train_example_numpy(): assert numpy.allclose(y1, y3) -@pytest.mark.kernels_ci def test_pytorch_yolo_fragment(): if not is_backend_tested("torch"): pytest.skip() diff --git a/einops/src/tests/test_layers.py b/einops/src/tests/test_layers.py index 538b3b8e1..46afdd069 100644 --- a/einops/src/tests/test_layers.py +++ b/einops/src/tests/test_layers.py @@ -1,16 +1,23 @@ +__kernel_port_einops_root = None +def __kernel_port_einops(module=""): + global __kernel_port_einops_root + if __kernel_port_einops_root is None: + __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) + root = __kernel_port_einops_root + if not module: + return root + return __import__("importlib").import_module(root.__name__ + "." + module) + import pickle from collections import namedtuple import numpy import pytest -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -rearrange = einops.rearrange -reduce = einops.reduce -EinopsError = einops.EinopsError -from . import collect_test_backends, is_backend_tested -from . import FLOAT_REDUCTIONS as REDUCTIONS +pytestmark = pytest.mark.kernels_ci + +rearrange, reduce, EinopsError = getattr(__kernel_port_einops(), "rearrange"), getattr(__kernel_port_einops(), "reduce"), getattr(__kernel_port_einops(), "EinopsError") +from . import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS __author__ = "Alex Rogozhnikov" @@ -38,7 +45,6 @@ ] -@pytest.mark.kernels_ci def test_rearrange_imperative(): for backend in collect_test_backends(symbolic=False, layers=True): print("Test layer for ", backend.framework_name) @@ -71,7 +77,6 @@ def test_rearrange_imperative(): assert numpy.allclose(backend.to_numpy(variable.grad), 1) -@pytest.mark.kernels_ci def test_rearrange_symbolic(): for backend in collect_test_backends(symbolic=True, layers=True): print("Test layer for ", backend.framework_name) @@ -112,7 +117,6 @@ def test_rearrange_symbolic(): ] -@pytest.mark.kernels_ci def test_reduce_imperative(): for backend in collect_test_backends(symbolic=False, layers=True): print("Test layer for ", backend.framework_name) @@ -154,7 +158,6 @@ def test_reduce_imperative(): assert numpy.sum(grad) > 0.5 -@pytest.mark.kernels_ci def test_reduce_symbolic(): for backend in collect_test_backends(symbolic=True, layers=True): print("Test layer for ", backend.framework_name) @@ -186,7 +189,7 @@ def create_torch_model(use_reduce=False, add_scripted_layer=False): pytest.skip() else: from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU - Rearrange, Reduce, EinMix = einops.layers.torch.Rearrange, einops.layers.torch.Reduce, einops.layers.torch.EinMix + Rearrange, Reduce, EinMix = getattr(__kernel_port_einops("layers.torch"), "Rearrange"), getattr(__kernel_port_einops("layers.torch"), "Reduce"), getattr(__kernel_port_einops("layers.torch"), "EinMix") import torch.jit return Sequential( @@ -207,7 +210,6 @@ def create_torch_model(use_reduce=False, add_scripted_layer=False): ) -@pytest.mark.kernels_ci def test_torch_layer(): if not is_backend_tested("torch"): pytest.skip() @@ -234,7 +236,6 @@ def test_torch_layer(): torch.testing.assert_close(model1(input + 1), model4(input + 1), atol=1e-3, rtol=1e-3) -@pytest.mark.kernels_ci def test_torch_layers_scripting(): if not is_backend_tested("torch"): pytest.skip() @@ -249,7 +250,6 @@ def test_torch_layers_scripting(): torch.testing.assert_close(model1(input), model2(input), atol=1e-3, rtol=1e-3) -@pytest.mark.kernels_ci def test_keras_layer(): if not is_backend_tested("tensorflow"): pytest.skip() @@ -261,7 +261,7 @@ def test_keras_layer(): pytest.skip() from tensorflow.keras.models import Sequential from tensorflow.keras.layers import Conv2D as Conv2d, Dense as Linear, ReLU - Rearrange, Reduce, EinMix, keras_custom_objects = einops.layers.keras.Rearrange, einops.layers.keras.Reduce, einops.layers.keras.EinMix, einops.layers.keras.keras_custom_objects + Rearrange, Reduce, EinMix, keras_custom_objects = getattr(__kernel_port_einops("layers.keras"), "Rearrange"), getattr(__kernel_port_einops("layers.keras"), "Reduce"), getattr(__kernel_port_einops("layers.keras"), "EinMix"), getattr(__kernel_port_einops("layers.keras"), "keras_custom_objects") def create_keras_model(): return Sequential( @@ -310,7 +310,6 @@ def create_keras_model(): numpy.testing.assert_allclose(model1.predict_on_batch(input), model4.predict_on_batch(input)) -@pytest.mark.kernels_ci def test_flax_layers(): """ One-off simple tests for Flax layers. @@ -324,7 +323,7 @@ def test_flax_layers(): import flax from flax import linen as nn - EinMix, Reduce, Rearrange = einops.layers.flax.EinMix, einops.layers.flax.Reduce, einops.layers.flax.Rearrange + EinMix, Reduce, Rearrange = getattr(__kernel_port_einops("layers.flax"), "EinMix"), getattr(__kernel_port_einops("layers.flax"), "Reduce"), getattr(__kernel_port_einops("layers.flax"), "Rearrange") class NN(nn.Module): @nn.compact @@ -358,12 +357,11 @@ def eval_at_point(params): _loaded = flax.serialization.from_bytes(params, fbytes) -@pytest.mark.kernels_ci def test_einmix_decomposition(): """ Testing that einmix correctly decomposes into smaller transformations. """ - _EinmixDebugger = einops.layers._einmix._EinmixDebugger + _EinmixDebugger = getattr(__kernel_port_einops("layers._einmix"), "_EinmixDebugger") mixin1 = _EinmixDebugger( "a b c d e -> e d c b a", @@ -449,12 +447,11 @@ def test_einmix_decomposition(): assert mixin7.saved_bias_shape == [1, 4, 2] # (a) d b, ellipsis does not participate -@pytest.mark.kernels_ci def test_einmix_restrictions(): """ Testing different cases """ - _EinmixDebugger = einops.layers._einmix._EinmixDebugger + _EinmixDebugger = getattr(__kernel_port_einops("layers._einmix"), "_EinmixDebugger") with pytest.raises(EinopsError): _EinmixDebugger( diff --git a/einops/src/tests/test_ops.py b/einops/src/tests/test_ops.py index ff668bb50..ef0196242 100644 --- a/einops/src/tests/test_ops.py +++ b/einops/src/tests/test_ops.py @@ -1,18 +1,24 @@ +__kernel_port_einops_root = None +def __kernel_port_einops(module=""): + global __kernel_port_einops_root + if __kernel_port_einops_root is None: + __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) + root = __kernel_port_einops_root + if not module: + return root + return __import__("importlib").import_module(root.__name__ + "." + module) + import itertools import numpy import numpy as np import pytest -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -rearrange = einops.einops.rearrange -reduce = einops.einops.reduce -repeat = einops.einops.repeat -_enumerate_directions = einops.einops._enumerate_directions -from . import collect_test_backends, is_backend_tested -from . import FLOAT_REDUCTIONS as REDUCTIONS +pytestmark = pytest.mark.kernels_ci + +EinopsError = getattr(__kernel_port_einops(), "EinopsError") +rearrange, reduce, repeat, _enumerate_directions = getattr(__kernel_port_einops("einops"), "rearrange"), getattr(__kernel_port_einops("einops"), "reduce"), getattr(__kernel_port_einops("einops"), "repeat"), getattr(__kernel_port_einops("einops"), "_enumerate_directions") +from . 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) @@ -45,7 +51,6 @@ ] -@pytest.mark.kernels_ci def test_collapsed_ellipsis_errors_out(): x = numpy.zeros([1, 1, 1, 1, 1]) rearrange(x, "a b c d ... -> a b c ... d") @@ -57,7 +62,6 @@ def test_collapsed_ellipsis_errors_out(): rearrange(x, "(...) -> (...)") -@pytest.mark.kernels_ci def test_ellipsis_ops_numpy(): x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) for pattern in identity_patterns: @@ -103,7 +107,6 @@ def operation(x): check_equal(numpy_result, backend_result) -@pytest.mark.kernels_ci def test_ellipsis_ops_imperative(): """Checking various patterns against numpy""" x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) @@ -121,10 +124,9 @@ def test_ellipsis_ops_imperative(): ) -@pytest.mark.kernels_ci def test_rearrange_array_api(): import numpy as xp - AA = einops.array_api + AA = getattr(__kernel_port_einops(), "array_api") if xp.__version__ < "2.0.0": pytest.skip() @@ -136,10 +138,9 @@ def test_rearrange_array_api(): assert numpy.array_equal(AA.asnumpy(result + 0), expected) -@pytest.mark.kernels_ci def test_reduce_array_api(): import numpy as xp - AA = einops.array_api + AA = getattr(__kernel_port_einops(), "array_api") if xp.__version__ < "2.0.0": pytest.skip() @@ -152,7 +153,6 @@ def test_reduce_array_api(): assert numpy.array_equal(AA.asnumpy(np.asarray(result + 0)), expected) -@pytest.mark.kernels_ci def test_rearrange_consistency_numpy(): shape = [1, 2, 3, 5, 7, 11] x = numpy.arange(numpy.prod(shape)).reshape(shape) @@ -191,7 +191,6 @@ def test_rearrange_consistency_numpy(): assert x2[0, 1, 2] == result[1, 2, 0] -@pytest.mark.kernels_ci def test_rearrange_permutations_numpy(): # tests random permutation of axes against two independent numpy ways for n_axes in range(1, 10): @@ -220,7 +219,6 @@ def test_rearrange_permutations_numpy(): assert numpy.array_equal(result, expected_result) -@pytest.mark.kernels_ci def test_reduction_imperatives(): for backend in imp_op_backends: print("Reduction tests for ", backend.framework_name) @@ -257,7 +255,6 @@ def test_reduction_imperatives(): assert numpy.allclose(result, expected_result), f"Failed at {pattern}" -@pytest.mark.kernels_ci def test_reduction_symbolic(): for backend in sym_op_backends: print("Reduction tests for ", backend.framework_name) @@ -311,7 +308,6 @@ def test_reduction_symbolic(): assert numpy.allclose(result, expected_numpy_result) -@pytest.mark.kernels_ci def test_reduction_stress_imperatives(): for backend in imp_op_backends: print("Stress-testing reduction for ", backend.framework_name) @@ -344,7 +340,6 @@ def test_reduction_stress_imperatives(): check_op_against_numpy(backend, x, pattern, reduction=reduction, axes_lengths={}, is_symbolic=False) -@pytest.mark.kernels_ci def test_reduction_with_callable_imperatives(): x_numpy = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]).astype("float32") x_numpy /= x_numpy.max() @@ -370,10 +365,7 @@ def logsumexp_numpy(x, tuple_of_axes): y = numpy.sum(y, axis=tuple_of_axes) return numpy.log(y) + minused - TorchBackend = einops._backends.TorchBackend - TensorflowBackend = einops._backends.TensorflowBackend - TFKerasBackend = einops._backends.TFKerasBackend - NumpyBackend = einops._backends.NumpyBackend + TorchBackend, TensorflowBackend, TFKerasBackend, NumpyBackend = getattr(__kernel_port_einops("_backends"), "TorchBackend"), getattr(__kernel_port_einops("_backends"), "TensorflowBackend"), getattr(__kernel_port_einops("_backends"), "TFKerasBackend"), getattr(__kernel_port_einops("_backends"), "NumpyBackend") backend2callback = { TorchBackend.framework_name: logsumexp_torch, @@ -399,7 +391,6 @@ def logsumexp_numpy(x, tuple_of_axes): ) -@pytest.mark.kernels_ci def test_enumerating_directions(): for backend in imp_op_backends: print("testing directions for", backend.framework_name) @@ -414,7 +405,6 @@ def test_enumerating_directions(): assert numpy.allclose(ax1, ax2) -@pytest.mark.kernels_ci def test_concatenations_and_stacking(): for backend in imp_op_backends: print("testing shapes for ", backend.framework_name) @@ -434,7 +424,6 @@ def test_concatenations_and_stacking(): assert numpy.array_equal(result1, backend.to_numpy(result2)) -@pytest.mark.kernels_ci def test_gradients_imperatives(): # lazy - just checking reductions for reduction in REDUCTIONS: @@ -462,7 +451,6 @@ def test_gradients_imperatives(): assert numpy.allclose(grad1, grad2), [name1, name2, "provided different gradients"] -@pytest.mark.kernels_ci def test_tiling_imperatives(): for backend in imp_op_backends: print("Tiling tests for ", backend.framework_name) @@ -480,7 +468,6 @@ def test_tiling_imperatives(): assert numpy.array_equal(result, expected) -@pytest.mark.kernels_ci def test_tiling_symbolic(): for backend in sym_op_backends: print("Tiling tests for ", backend.framework_name) @@ -526,7 +513,6 @@ def check_reversion(x, repeat_pattern, **sizes): assert numpy.array_equal(x, reduced_max) -@pytest.mark.kernels_ci def test_repeat_numpy(): # check repeat vs reduce. Repeat works ok if reverse reduction with min and max work well x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) @@ -536,7 +522,6 @@ def test_repeat_numpy(): check_reversion(x, pattern, **axis_dimensions) -@pytest.mark.kernels_ci def test_repeat_imperatives(): x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) for backend in imp_op_backends: @@ -550,7 +535,6 @@ def test_repeat_imperatives(): assert numpy.array_equal(result, expected) -@pytest.mark.kernels_ci def test_repeat_symbolic(): x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) @@ -565,10 +549,9 @@ def test_repeat_symbolic(): assert numpy.array_equal(result, expected) -@pytest.mark.kernels_ci def test_repeat_array_api(): import numpy as xp - AA = einops.array_api + AA = getattr(__kernel_port_einops(), "array_api") if xp.__version__ < "2.0.0": pytest.skip() @@ -593,14 +576,12 @@ def test_repeat_array_api(): ] -@pytest.mark.kernels_ci def test_anonymous_axes(): x = numpy.arange(1 * 2 * 4 * 6).reshape([1, 2, 4, 6]) for pattern, axis_dimensions in test_cases_repeat_anonymous: check_reversion(x, pattern, **axis_dimensions) -@pytest.mark.kernels_ci def test_list_inputs(): x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) @@ -618,7 +599,6 @@ def test_list_inputs(): ) -@pytest.mark.kernels_ci def test_torch_compile_with_dynamic_shape(): if not is_backend_tested("torch"): pytest.skip() @@ -650,7 +630,6 @@ def bit_count(x): return sum((x >> i) & 1 for i in range(20)) -@pytest.mark.kernels_ci def test_reduction_imperatives_booleans(): """Checks that any/all reduction works in all frameworks""" x_np = numpy.asarray([(bit_count(x) % 2) == 0 for x in range(2**6)]).reshape([2] * 6) diff --git a/einops/src/tests/test_other.py b/einops/src/tests/test_other.py index 71fcb9ce3..ff4346851 100644 --- a/einops/src/tests/test_other.py +++ b/einops/src/tests/test_other.py @@ -1,27 +1,37 @@ +__kernel_port_einops_root = None +def __kernel_port_einops(module=""): + global __kernel_port_einops_root + if __kernel_port_einops_root is None: + __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) + root = __kernel_port_einops_root + if not module: + return root + return __import__("importlib").import_module(root.__name__ + "." + module) + from doctest import testmod +testmod.__test__ = False import numpy import pytest -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -AbstractBackend = einops._backends.AbstractBackend -rearrange = einops.einops.rearrange -parse_shape = einops.einops.parse_shape -_optimize_transformation = einops.einops._optimize_transformation +pytestmark = pytest.mark.kernels_ci + +einops = __kernel_port_einops() +einops = __kernel_port_einops(); __kernel_port_einops("layers") +einops = __kernel_port_einops(); __kernel_port_einops("parsing") +AbstractBackend = getattr(__kernel_port_einops("_backends"), "AbstractBackend") +rearrange, parse_shape, _optimize_transformation = getattr(__kernel_port_einops("einops"), "rearrange"), getattr(__kernel_port_einops("einops"), "parse_shape"), getattr(__kernel_port_einops("einops"), "_optimize_transformation") from . import collect_test_backends, is_backend_tested __author__ = "Alex Rogozhnikov" -@pytest.mark.kernels_ci def test_doctests_examples(): # tests docstrings, additionally testmod(einops.layers, raise_on_error=True, extraglobs=dict(np=numpy)) testmod(einops.einops, raise_on_error=True, extraglobs=dict(np=numpy)) -@pytest.mark.kernels_ci def test_backends_installed(): """ This test will fail if some of backends are not installed or can't be imported @@ -42,7 +52,6 @@ def test_backends_installed(): assert len(errors) == 0, errors -@pytest.mark.kernels_ci def test_optimize_transformations_numpy(): print("Testing optimizations") shapes = [[2] * n_dimensions for n_dimensions in range(14)] @@ -80,7 +89,6 @@ def test_optimize_transformations_numpy(): x_np = numpy.zeros([10, 20, 30, 40]) -@pytest.mark.kernels_ci def test_parse_shape_imperative(): for backend in _IMPERATIVE_BACKENDS: print("Shape parsing for ", backend.framework_name) @@ -90,7 +98,6 @@ def test_parse_shape_imperative(): assert parsed1 != dict(a=1, b=20, c=30, d=40) != parsed2 -@pytest.mark.kernels_ci def test_underscore(): for backend in _IMPERATIVE_BACKENDS: parsed1 = parse_shape(x_np, "_ _ _ _") @@ -98,7 +105,6 @@ def test_underscore(): assert parsed1 == parsed2 == dict() -@pytest.mark.kernels_ci def test_underscore_one(): for backend in _IMPERATIVE_BACKENDS: parsed1 = parse_shape(x_np, "_ _ _ hello") @@ -106,7 +112,6 @@ def test_underscore_one(): assert parsed1 == parsed2 == dict(hello=40) -@pytest.mark.kernels_ci def test_underscore_several(): for backend in _IMPERATIVE_BACKENDS: parsed1 = parse_shape(x_np, "_ _ a1 a1a111a") @@ -114,7 +119,6 @@ def test_underscore_several(): assert parsed1 == parsed2 == dict(a1=30, a1a111a=40) -@pytest.mark.kernels_ci def test_repeating(): with pytest.raises(einops.EinopsError): parse_shape(x_np, "a a b b") @@ -124,7 +128,6 @@ def test_repeating(): parse_shape(backend.from_numpy(x_np), "a a b b") -@pytest.mark.kernels_ci def test_ellipsis(): for backend in _IMPERATIVE_BACKENDS: for shape, pattern, expected in [ @@ -146,7 +149,6 @@ def test_ellipsis(): assert parsed1 == parsed2 == expected -@pytest.mark.kernels_ci def test_parse_with_anonymous_axes(): for backend in _IMPERATIVE_BACKENDS: for shape, pattern, expected in [ @@ -160,7 +162,6 @@ def test_parse_with_anonymous_axes(): assert parsed1 == parsed2 == expected -@pytest.mark.kernels_ci def test_failures(): for backend in _IMPERATIVE_BACKENDS: # every test should fail @@ -188,7 +189,6 @@ def test_failures(): @pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) -@pytest.mark.kernels_ci def test_parse_shape_symbolic(backend): for shape in [ [10, 20, 30, 40], @@ -219,7 +219,6 @@ def test_parse_shape_symbolic(backend): @pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) -@pytest.mark.kernels_ci def test_parse_shape_symbolic_ellipsis(backend): for static_shape, shape, pattern, expected in [ ([10, 20], [None, None], "...", dict()), @@ -245,7 +244,6 @@ def test_parse_shape_symbolic_ellipsis(backend): assert out_shape == expected -@pytest.mark.kernels_ci def test_is_float_type(): backends = collect_test_backends(symbolic=False, layers=False) backends += collect_test_backends(symbolic=False, layers=True) @@ -257,7 +255,6 @@ def test_is_float_type(): assert backend.is_float_type(input) == is_float, (dtype, backend, input.dtype) -@pytest.mark.kernels_ci def test_torch_compile(): """ Test ensures that allow_ops_in_compiled_graph allows compiling in a single graph @@ -270,8 +267,8 @@ def test_torch_compile(): pytest.skip() import torch from torch import nn - repeat, reduce, pack, unpack, einsum = einops.repeat, einops.reduce, einops.pack, einops.unpack, einops.einsum - allow_ops_in_compiled_graph = einops._torch_specific.allow_ops_in_compiled_graph + repeat, reduce, pack, unpack, einsum = getattr(__kernel_port_einops(), "repeat"), getattr(__kernel_port_einops(), "reduce"), getattr(__kernel_port_einops(), "pack"), getattr(__kernel_port_einops(), "unpack"), getattr(__kernel_port_einops(), "einsum") + allow_ops_in_compiled_graph = getattr(__kernel_port_einops("_torch_specific"), "allow_ops_in_compiled_graph") allow_ops_in_compiled_graph() diff --git a/einops/src/tests/test_packing.py b/einops/src/tests/test_packing.py index d76553475..09869c3fb 100644 --- a/einops/src/tests/test_packing.py +++ b/einops/src/tests/test_packing.py @@ -1,15 +1,22 @@ +__kernel_port_einops_root = None +def __kernel_port_einops(module=""): + global __kernel_port_einops_root + if __kernel_port_einops_root is None: + __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) + root = __kernel_port_einops_root + if not module: + return root + return __import__("importlib").import_module(root.__name__ + "." + module) + import dataclasses import typing import numpy as np import pytest -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -asnumpy = einops.asnumpy -pack = einops.pack -unpack = einops.unpack +pytestmark = pytest.mark.kernels_ci + +EinopsError, asnumpy, pack, unpack = getattr(__kernel_port_einops(), "EinopsError"), getattr(__kernel_port_einops(), "asnumpy"), getattr(__kernel_port_einops(), "pack"), getattr(__kernel_port_einops(), "unpack") from . import collect_test_backends @@ -64,7 +71,6 @@ def __exit__(self, exc_type, exc_val, exc_tb): return True -@pytest.mark.kernels_ci def test_numpy_trivial(H=13, W=17): def rand(*shape): return np.random.random(shape) @@ -148,7 +154,6 @@ def selfcheck(self): ] -@pytest.mark.kernels_ci def test_pack_unpack_with_numpy(): case: UnpackTestCase @@ -212,7 +217,6 @@ def test_pack_unpack_with_numpy(): unpack_and_pack(x, [[2, -1], [1, 5]], pattern) -@pytest.mark.kernels_ci def test_pack_unpack_against_numpy(): for backend in collect_test_backends(symbolic=False, layers=False): print(f"test packing against numpy for {backend.framework_name}") @@ -277,9 +281,8 @@ def test_pack_unpack_against_numpy(): unpack_and_pack(x, [[2, -1], [1, 5]], pattern) -@pytest.mark.kernels_ci def test_pack_unpack_array_api(): - AA = einops.array_api + AA = getattr(__kernel_port_einops(), "array_api") import numpy as xp if xp.__version__ < "2.0.0": diff --git a/einops/src/tests/test_parsing.py b/einops/src/tests/test_parsing.py index 491126793..6687adee8 100644 --- a/einops/src/tests/test_parsing.py +++ b/einops/src/tests/test_parsing.py @@ -1,11 +1,19 @@ +__kernel_port_einops_root = None +def __kernel_port_einops(module=""): + global __kernel_port_einops_root + if __kernel_port_einops_root is None: + __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) + root = __kernel_port_einops_root + if not module: + return root + return __import__("importlib").import_module(root.__name__ + "." + module) + import pytest -import kernels -einops = kernels.get_kernel("kernels-community/einops", version=1) -EinopsError = einops.EinopsError -ParsedExpression = einops.parsing.ParsedExpression -AnonymousAxis = einops.parsing.AnonymousAxis -_ellipsis = einops.parsing._ellipsis +pytestmark = pytest.mark.kernels_ci + +EinopsError = getattr(__kernel_port_einops(), "EinopsError") +ParsedExpression, AnonymousAxis, _ellipsis = getattr(__kernel_port_einops("parsing"), "ParsedExpression"), getattr(__kernel_port_einops("parsing"), "AnonymousAxis"), getattr(__kernel_port_einops("parsing"), "_ellipsis") __author__ = "Alex Rogozhnikov" @@ -19,7 +27,6 @@ def __eq__(self, other): return isinstance(other, AnonymousAxis) and self.value == other.value -@pytest.mark.kernels_ci def test_anonymous_axes(): a, b = AnonymousAxis("2"), AnonymousAxis("2") assert a != b @@ -29,7 +36,6 @@ def test_anonymous_axes(): assert [a, 2, b] == [c, 2, c] -@pytest.mark.kernels_ci def test_elementary_axis_name(): for name in [ "a", @@ -51,7 +57,6 @@ def test_elementary_axis_name(): assert not ParsedExpression.check_axis_name(name) -@pytest.mark.kernels_ci def test_invalid_expressions(): # double ellipsis should raise an error ParsedExpression("... a b c d") @@ -85,7 +90,6 @@ def test_invalid_expressions(): ParsedExpression("pre...") -@pytest.mark.kernels_ci def test_parse_expression(): parsed = ParsedExpression("a1 b1 c1 d1") assert parsed.identifiers == {"a1", "b1", "c1", "d1"} From cb337de8a6b604eadf1737cf020b5cad5b131627 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Wed, 2 Sep 2026 15:19:33 -0400 Subject: [PATCH 09/15] fix: improve port ci checks and fix typos --- .github/workflows/check-ports.yaml | 28 +++++++--- einops/flake.nix | 2 +- einops/port/port.kdl | 3 ++ einops/src/port-provenance.json | 4 +- einops/src/tests/run_tests.py | 85 ------------------------------ scripts/check_ports.py | 44 +++++++++++++--- 6 files changed, 65 insertions(+), 101 deletions(-) delete mode 100644 einops/src/tests/run_tests.py diff --git a/.github/workflows/check-ports.yaml b/.github/workflows/check-ports.yaml index 1d9b29138..464cf26d7 100644 --- a/.github/workflows/check-ports.yaml +++ b/.github/workflows/check-ports.yaml @@ -1,9 +1,14 @@ name: Check Ports # A ported kernel's /src is generated from /port/port.kdl and -# committed. Regenerate it here and fail if it drifted, so that a stale overlay -# silently reverting an in-tree fix, a hand-edited build.toml field the manifest -# op drops, or an overlay file edited without its generated twin are all loud. +# committed. Regenerate it here and fail if the two disagree: a stale overlay +# reverting an in-tree fix, a hand-edited build.toml field the manifest op +# drops, an overlay file edited without its generated twin, or a recipe edited +# without rerunning the port (src/port-provenance.json records the recipe hash, +# so that case shows up as a diff in the provenance file). +# +# Pull requests only check the ports they touch; pushes to main and manual +# runs check every port. on: push: branches: [main] @@ -30,6 +35,8 @@ jobs: runs-on: ubuntu-latest steps: - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + fetch-depth: 0 - uses: DeterminateSystems/nix-installer-action@ef8a148080ab6020fd15196c2084a2eea5ff2d25 # v22 @@ -37,7 +44,14 @@ jobs: # origin URL, HEAD and cleanliness of the checkout it is handed. - name: Regenerate every port and compare env: - # kernel-port still lives on the kernel-port-tool branch; point this at - # github:huggingface/kernels#kernel-port once that branch merges. - KERNEL_PORT: "nix run github:huggingface/kernels/kernel-port-tool#kernel-port --" - run: python3 scripts/check_ports.py + # TODO: kernel-port still lives on the kernel-port-tool branch, pinned + # to a rev so a push there cannot break every open PR. Point this at + # github:huggingface/kernels#kernel-port once the branch merges. + KERNEL_PORT: "nix run github:huggingface/kernels/c29a6751b5c699de555c98060fb26b4ef0e6b520#kernel-port --" + BASE_REF: ${{ github.event_name == 'pull_request' && format('origin/{0}', github.base_ref) || '' }} + run: | + if [ -n "$BASE_REF" ]; then + python3 scripts/check_ports.py --changed-since "$BASE_REF" + else + python3 scripts/check_ports.py + fi diff --git a/einops/flake.nix b/einops/flake.nix index b7b6ea3ec..70f28b83d 100644 --- a/einops/flake.nix +++ b/einops/flake.nix @@ -1,5 +1,5 @@ { - description = "Flake for eigen kernel"; + description = "Flake for einops kernel"; inputs = { kernel-builder.url = "github:huggingface/kernels"; diff --git a/einops/port/port.kdl b/einops/port/port.kdl index 84aeb429b..ce9aa3e63 100644 --- a/einops/port/port.kdl +++ b/einops/port/port.kdl @@ -8,6 +8,9 @@ source \ // Retain only the Python package and its upstream tests. prune keep="einops/**" +// Drop upstream's pip-install test runner; it is not a test. +delete in="einops/tests/run_tests.py" + // Lift upstream tests out of the package into the kernel test root. move \ from="einops/tests" \ diff --git a/einops/src/port-provenance.json b/einops/src/port-provenance.json index d0ee2129b..c7259c5e2 100644 --- a/einops/src/port-provenance.json +++ b/einops/src/port-provenance.json @@ -1,6 +1,6 @@ { "format": 1, - "recipe": "sha256:553ea0a8342085862371639c2f12947c592617ffb7af4343b02a3048986a026c", + "recipe": "sha256:09ea109865a0ef1de6779af7db28fc5baa2cd01d0b52ee5795f07c063881ee18", "runner": "kernel-port 0.1.0", "sources": { "upstream": { @@ -8,5 +8,5 @@ "commit": "ad2c8d67d08cbebb45e1aff2be21dc88635fc695" } }, - "tree": "sha256:ed3299bb86f88fa6fc9fcf72551a64830b327b879f9a4da791d29981a4b840f4" + "tree": "sha256:bd455bcb9afd28fb5c0c2b3c6c9013739fc71554017a2035e5bd782398c4dd62" } diff --git a/einops/src/tests/run_tests.py b/einops/src/tests/run_tests.py deleted file mode 100644 index cf7a4eb6e..000000000 --- a/einops/src/tests/run_tests.py +++ /dev/null @@ -1,85 +0,0 @@ -""" -Runs tests that are appropriate for framework. -""" - -import os -import sys -from subprocess import Popen -from pathlib import Path - -__author__ = "Alex Rogozhnikov" - - -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 - - -def main(): - _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]", "flax"], - "tensorflow": ["tensorflow"], - "cupy": ["cupy"], - # switch to stable paddlepaddle, because of https://github.com/PaddlePaddle/Paddle/issues/63927 - # "paddle": ["paddlepaddle==0.0.0 -f https://www.paddlepaddle.org.cn/whl/linux/cpu-mkl/develop.html"], - "paddle": ["paddlepaddle"], - "oneflow": ["oneflow==0.9.0"], - "pytensor": ["pytensor"], - } - - usage = f""" - Usage: python -m einops.tests.run_tests [--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) - return - else: - synonyms = { - "tf": "tensorflow", - "pytorch": "torch", - "paddlepaddle": "paddle", - } - frameworks = [synonyms.get(f, f) for f in frameworks] - wrong_frameworks = [f for f in frameworks if f not in framework_name2installation] - if wrong_frameworks: - print(usage) - raise RuntimeError(f"Unrecognized frameworks: {wrong_frameworks}") - - 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))) - - 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 an envvar EINOPS_TEST_BACKENDS - from einops.tests import unparse_backends - - envvar_name, envvar_value = unparse_backends(backend_names=frameworks) - return_code = run( - "python -m pytest .", - **{envvar_name: envvar_value}, - ) - assert return_code == 0 - - -if __name__ == "__main__": - main() diff --git a/scripts/check_ports.py b/scripts/check_ports.py index a60a7cadf..8adf44078 100644 --- a/scripts/check_ports.py +++ b/scripts/check_ports.py @@ -6,17 +6,22 @@ script re-runs the port against a fresh upstream checkout at the pinned commit and fails if the result differs from what is committed. -Without this, three failure modes are silent: +Without this, four failure modes are silent: - an overlay file goes stale, so regenerating reverts a fix made in-tree - a build.toml field is edited by hand, and the manifest op drops it - an overlay file and its generated twin are edited independently and drift + - the recipe is edited without rerunning the port; src/port-provenance.json + records the recipe hash, so this shows up as a diff in that file Usage: python3 scripts/check_ports.py [kernel ...] + python3 scripts/check_ports.py --changed-since -Defaults to every port found. Set KERNEL_PORT to override the runner command -(e.g. "nix run github:huggingface/kernels#kernel-port --"). +Defaults to every port found. --changed-since limits the check to ports whose +port/ or src/ tree differs from , which is what CI uses for pull requests. +Set KERNEL_PORT to override the runner command (e.g. "nix run +github:huggingface/kernels#kernel-port --"). """ import filecmp @@ -47,6 +52,22 @@ def find_ports(root: Path) -> list: return sorted(p.parent.parent.name for p in root.glob("*/port/port.kdl")) +def changed_ports(root: Path, ref: str) -> list: + """Ports whose port/ or src/ tree differs from `ref`.""" + out = subprocess.run( + ["git", "diff", "--name-only", f"{ref}...HEAD"], + capture_output=True, + text=True, + check=True, + ).stdout + touched = set() + for line in out.splitlines(): + parts = line.split("/") + if len(parts) >= 3 and parts[1] in ("port", "src"): + touched.add(parts[0]) + return sorted(k for k in find_ports(root) if k in touched) + + def parse_source(recipe: Path) -> tuple: text = recipe.read_text() m = SOURCE_RE.search(text) @@ -96,7 +117,7 @@ def check(kernel: str, root: Path, workdir: Path) -> list: print(f" {kernel}: {repo} @ {commit[:12]}") # One checkout per (repo, commit): two kernels may share an upstream. - src = workdir / f"upstream-{abs(hash((repo, commit)))}" + src = workdir / f"upstream-{commit[:12]}" if not src.exists(): checkout(repo, commit, src) @@ -116,7 +137,16 @@ def check(kernel: str, root: Path, workdir: Path) -> list: def main() -> int: root = Path.cwd() - kernels = sys.argv[1:] or find_ports(root) + args = sys.argv[1:] + if args[:1] == ["--changed-since"]: + if len(args) != 2: + raise SystemExit("usage: check_ports.py --changed-since ") + kernels = changed_ports(root, args[1]) + if not kernels: + print(f"No ported kernels changed since {args[1]}.") + return 0 + else: + kernels = args or find_ports(root) if not kernels: print("No ported kernels found (no */port/port.kdl).") return 0 @@ -144,7 +174,9 @@ def main() -> int: print( "\nRegenerate with:\n" " kernel-port /port/port.kdl --dir --out /src\n" - "If an overlay file went stale, refresh it from /src first.", + "If an overlay file went stale, refresh it from /src first.\n" + "If only port-provenance.json differs, the recipe changed and the\n" + "port was not rerun.", file=sys.stderr, ) return 1 From 11d283ce242c7db9b378590dd9f5fb3e46721b41 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 29 Sep 2026 20:42:45 +0000 Subject: [PATCH 10/15] feat: pin kernels version for port reproduction --- .github/workflows/check-ports.yaml | 9 +++++---- scripts/check_ports.py | 13 +++++++++++-- 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/.github/workflows/check-ports.yaml b/.github/workflows/check-ports.yaml index 464cf26d7..8b835981d 100644 --- a/.github/workflows/check-ports.yaml +++ b/.github/workflows/check-ports.yaml @@ -44,10 +44,11 @@ jobs: # origin URL, HEAD and cleanliness of the checkout it is handed. - name: Regenerate every port and compare env: - # TODO: kernel-port still lives on the kernel-port-tool branch, pinned - # to a rev so a push there cannot break every open PR. Point this at - # github:huggingface/kernels#kernel-port once the branch merges. - KERNEL_PORT: "nix run github:huggingface/kernels/c29a6751b5c699de555c98060fb26b4ef0e6b520#kernel-port --" + # kernel-port is experimental upstream (no CLI or op stability + # guarantee), so pin a huggingface/kernels main rev rather than + # tracking the branch: a push there cannot break every open PR, and + # bumping the pin is an explicit PR that re-checks every port. + KERNEL_PORT: "nix run github:huggingface/kernels/ca9095966486b76168115a34e2027c6c98805919#kernel-port --" BASE_REF: ${{ github.event_name == 'pull_request' && format('origin/{0}', github.base_ref) || '' }} run: | if [ -n "$BASE_REF" ]; then diff --git a/scripts/check_ports.py b/scripts/check_ports.py index 8adf44078..a7ab25a2d 100644 --- a/scripts/check_ports.py +++ b/scripts/check_ports.py @@ -19,7 +19,8 @@ python3 scripts/check_ports.py --changed-since Defaults to every port found. --changed-since limits the check to ports whose -port/ or src/ tree differs from , which is what CI uses for pull requests. +port/ or src/ tree differs from (or every port, if the checker itself +changed), which is what CI uses for pull requests. Set KERNEL_PORT to override the runner command (e.g. "nix run github:huggingface/kernels#kernel-port --"). """ @@ -43,6 +44,8 @@ DEFAULT_RUNNER = "kernel-port" +CHECKER_FILES = {"scripts/check_ports.py", ".github/workflows/check-ports.yaml"} + def runner_cmd() -> list: return shlex.split(os.environ.get("KERNEL_PORT", DEFAULT_RUNNER)) @@ -53,7 +56,11 @@ def find_ports(root: Path) -> list: def changed_ports(root: Path, ref: str) -> list: - """Ports whose port/ or src/ tree differs from `ref`.""" + """Ports whose port/ or src/ tree differs from `ref`. + + A change to the checker itself (this script, or the workflow that pins the + runner) can break any port, so it selects all of them. + """ out = subprocess.run( ["git", "diff", "--name-only", f"{ref}...HEAD"], capture_output=True, @@ -62,6 +69,8 @@ def changed_ports(root: Path, ref: str) -> list: ).stdout touched = set() for line in out.splitlines(): + if line in CHECKER_FILES: + return find_ports(root) parts = line.split("/") if len(parts) >= 3 and parts[1] in ("port", "src"): touched.add(parts[0]) From ac0b2264618ddd8fa994fe5b6bf4b83c1e9747e1 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 29 Sep 2026 20:52:24 +0000 Subject: [PATCH 11/15] fix: find build.toml under /src in release dispatch and ABI checks Ported kernels keep build.toml in /src. The manual release dispatch and run_kernels_checks.py only looked at /build.toml, so they skipped ported kernels such as einops. --- .github/workflows/build-release-dispatch.yaml | 6 +++++- scripts/run_kernels_checks.py | 5 ++++- 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/.github/workflows/build-release-dispatch.yaml b/.github/workflows/build-release-dispatch.yaml index c67b72a4a..cc09b4faa 100644 --- a/.github/workflows/build-release-dispatch.yaml +++ b/.github/workflows/build-release-dispatch.yaml @@ -40,7 +40,11 @@ jobs: KERNEL: ${{ inputs.kernel_name }} run: | if [ "${{ github.event_name }}" = "workflow_dispatch" ]; then - if [ -d "$KERNEL" ] && [ -f "$KERNEL/flake.nix" ] && [ -f "$KERNEL/build.toml" ]; then + # Ported kernels keep build.toml in /src; the flake stays at + # and points nix at it. + CONFIG="$KERNEL" + [ -f "$KERNEL/src/build.toml" ] && CONFIG="$KERNEL/src" + if [ -d "$KERNEL" ] && [ -f "$KERNEL/flake.nix" ] && [ -f "$CONFIG/build.toml" ]; then echo "kernel=$KERNEL" >> $GITHUB_OUTPUT echo "skip=false" >> $GITHUB_OUTPUT else diff --git a/scripts/run_kernels_checks.py b/scripts/run_kernels_checks.py index 39a75b2e1..ca854d43c 100644 --- a/scripts/run_kernels_checks.py +++ b/scripts/run_kernels_checks.py @@ -57,7 +57,10 @@ def discover_kernel_dirs(root: Path, excludes: list[str]) -> list[str]: name = entry.name if name in filtered: continue - elif entry.is_dir() and not (entry / "build.toml").exists(): + # Ported kernels keep build.toml in /src. + elif entry.is_dir() and not ( + (entry / "build.toml").exists() or (entry / "src" / "build.toml").exists() + ): logging.debug(f"Skipping {name} because it doesn't contain a `build.toml`.") continue if entry.is_dir(): From 61add17de47727c954876170028a277e680ddf4b Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 29 Sep 2026 21:01:07 +0000 Subject: [PATCH 12/15] feat: update checks for ported kernels --- .github/scripts/check_kernel_test_versions.py | 6 ++++-- .github/scripts/kernel_maintainers.py | 2 ++ .github/scripts/pr_kernel_build_guidance.py | 8 ++++++-- .../scripts/test_check_kernel_test_versions.py | 15 +++++++++++++++ .github/scripts/test_pr_kernel_build_guidance.py | 16 ++++++++++++++++ 5 files changed, 43 insertions(+), 4 deletions(-) diff --git a/.github/scripts/check_kernel_test_versions.py b/.github/scripts/check_kernel_test_versions.py index 433e5a7f1..f271084cc 100644 --- a/.github/scripts/check_kernel_test_versions.py +++ b/.github/scripts/check_kernel_test_versions.py @@ -40,7 +40,9 @@ def discover_kernels(repository_root: Path) -> list[Kernel]: if not manifest.is_file(): # Ported kernels keep generated sources and build.toml under src/. manifest = root / "src" / "build.toml" - if not manifest.is_file() or not (root / "tests").is_dir(): + # A ported kernel generates its tests next to build.toml in src/. + tests_root = manifest.parent if (manifest.parent / "tests").is_dir() else root + if not manifest.is_file() or not (tests_root / "tests").is_dir(): continue with manifest.open("rb") as handle: @@ -52,7 +54,7 @@ def discover_kernels(repository_root: Path) -> list[Kernel]: raise ValueError(f"{manifest}: [general.hub].repo-id must be a string") if type(version) is not int: raise ValueError(f"{manifest}: [general].version must be an integer") - kernels.append(Kernel(root=root, repo_id=repo_id, version=version)) + kernels.append(Kernel(root=tests_root, repo_id=repo_id, version=version)) return kernels diff --git a/.github/scripts/kernel_maintainers.py b/.github/scripts/kernel_maintainers.py index 5bd220db1..ee5ca0383 100755 --- a/.github/scripts/kernel_maintainers.py +++ b/.github/scripts/kernel_maintainers.py @@ -22,7 +22,9 @@ def load_registry() -> dict: def check_registry() -> None: registry = load_registry() + # Ported kernels keep build.toml in /src. kernels = {path.parent.name for path in Path(".").glob("*/build.toml")} + kernels |= {path.parent.parent.name for path in Path(".").glob("*/src/build.toml")} registered = set(registry) problems = [ diff --git a/.github/scripts/pr_kernel_build_guidance.py b/.github/scripts/pr_kernel_build_guidance.py index 1f100203d..4be60b4dd 100644 --- a/.github/scripts/pr_kernel_build_guidance.py +++ b/.github/scripts/pr_kernel_build_guidance.py @@ -43,7 +43,10 @@ def paginate(url: str, token: str) -> list[dict]: def kernel_directories(root: Path) -> set[str]: - return {manifest.parent.name for manifest in root.glob("*/build.toml")} + # Ported kernels keep build.toml in /src. + return {manifest.parent.name for manifest in root.glob("*/build.toml")} | { + manifest.parent.parent.name for manifest in root.glob("*/src/build.toml") + } def is_documentation(path: PurePosixPath) -> bool: @@ -55,7 +58,8 @@ def touched_source_kernels(files: list[dict], existing_kernels: set[str]) -> lis new_kernels = { path.parts[0] for path in paths - if len(path.parts) == 2 and path.name == "build.toml" + if path.name == "build.toml" + and (len(path.parts) == 2 or (len(path.parts) == 3 and path.parts[1] == "src")) } kernels = existing_kernels | new_kernels touched = set(new_kernels) diff --git a/.github/scripts/test_check_kernel_test_versions.py b/.github/scripts/test_check_kernel_test_versions.py index d60b8caaf..27eea2b86 100644 --- a/.github/scripts/test_check_kernel_test_versions.py +++ b/.github/scripts/test_check_kernel_test_versions.py @@ -46,6 +46,21 @@ def test_matching_import_with_nested_manifest(tmp_path): assert checker.check_repository(tmp_path) == ([], 1, 1) +def test_ported_kernel_tests_under_src(tmp_path): + kernel = tmp_path / "example" + (kernel / "src" / "tests").mkdir(parents=True) + (kernel / "src" / "build.toml").write_text( + '[general]\nversion = 2\n\n[general.hub]\nrepo-id = "kernels-community/example"\n' + ) + (kernel / "src" / "tests" / "test_example.py").write_text( + "from kernels import get_kernel\n" + 'example = get_kernel("kernels-community/example", version=1)\n' + ) + + problems, _, _ = checker.check_repository(tmp_path) + assert len(problems) == 1 + + def test_matching_call_is_formatting_agnostic(tmp_path): kernel = write_kernel(tmp_path) (kernel / "tests" / "test_example.py").write_text( diff --git a/.github/scripts/test_pr_kernel_build_guidance.py b/.github/scripts/test_pr_kernel_build_guidance.py index 8ad5fed83..b7218814c 100644 --- a/.github/scripts/test_pr_kernel_build_guidance.py +++ b/.github/scripts/test_pr_kernel_build_guidance.py @@ -39,6 +39,22 @@ def test_unknown_extensions_and_test_changes_are_build_relevant(tmp_path): ) == ["relu"] +def test_detects_ported_kernel_source(tmp_path): + (tmp_path / "einops" / "src").mkdir(parents=True) + (tmp_path / "einops" / "src" / "build.toml").touch() + files = [{"filename": "einops/port/port.kdl", "status": "modified"}] + + assert guidance.touched_source_kernels( + files, guidance.kernel_directories(tmp_path) + ) == ["einops"] + + +def test_detects_new_ported_kernel_from_added_manifest(): + files = [{"filename": "new-kernel/src/build.toml", "status": "added"}] + + assert guidance.touched_source_kernels(files, set()) == ["new-kernel"] + + def test_detects_new_kernel_from_added_manifest(): files = [ {"filename": "new-kernel/build.toml", "status": "added"}, From 7b7fc4d59f3bd39ee9888856e81b3915a753572a Mon Sep 17 00:00:00 2001 From: David Holtz Date: Mon, 5 Oct 2026 15:23:04 +0000 Subject: [PATCH 13/15] feat: better test kernel import --- .github/workflows/check-ports.yaml | 17 +++--- einops/port/overlay/tests/conftest.py | 13 ++++ einops/port/port.kdl | 34 +++++------ einops/src/build.toml | 2 +- einops/src/port-provenance.json | 4 +- einops/src/tests/__init__.py | 16 ++--- einops/src/tests/conftest.py | 13 ++++ einops/src/tests/run_tests.py | 85 +++++++++++++++++++++++++++ einops/src/tests/test_einsum.py | 22 +++---- einops/src/tests/test_examples.py | 20 +++---- einops/src/tests/test_layers.py | 38 ++++++------ einops/src/tests/test_ops.py | 53 +++++++++++------ einops/src/tests/test_other.py | 43 +++++++------- einops/src/tests/test_packing.py | 25 ++++---- einops/src/tests/test_parsing.py | 24 ++++---- 15 files changed, 258 insertions(+), 151 deletions(-) create mode 100644 einops/port/overlay/tests/conftest.py create mode 100644 einops/src/tests/conftest.py create mode 100644 einops/src/tests/run_tests.py diff --git a/.github/workflows/check-ports.yaml b/.github/workflows/check-ports.yaml index 8b835981d..dab3574fd 100644 --- a/.github/workflows/check-ports.yaml +++ b/.github/workflows/check-ports.yaml @@ -7,11 +7,15 @@ name: Check Ports # without rerunning the port (src/port-provenance.json records the recipe hash, # so that case shows up as a diff in the provenance file). # -# Pull requests only check the ports they touch; pushes to main and manual -# runs check every port. +# Pull requests only check the ports they touch; pushes to main, the daily +# run and manual runs check every port. on: push: branches: [main] + # kernel-port tracks huggingface/kernels main, so also run without a change + # here: a kernel-port regression should surface even when no port PR is open. + schedule: + - cron: "0 6 * * *" pull_request: branches: [main] types: [opened, synchronize, reopened] @@ -44,11 +48,10 @@ jobs: # origin URL, HEAD and cleanliness of the checkout it is handed. - name: Regenerate every port and compare env: - # kernel-port is experimental upstream (no CLI or op stability - # guarantee), so pin a huggingface/kernels main rev rather than - # tracking the branch: a push there cannot break every open PR, and - # bumping the pin is an explicit PR that re-checks every port. - KERNEL_PORT: "nix run github:huggingface/kernels/ca9095966486b76168115a34e2027c6c98805919#kernel-port --" + # Track huggingface/kernels main rather than pinning a rev, so that a + # kernel-port change that breaks a port is caught immediately. + # --refresh skips Nix's flake cache so every run sees the latest main. + KERNEL_PORT: "nix run --refresh github:huggingface/kernels#kernel-port --" BASE_REF: ${{ github.event_name == 'pull_request' && format('origin/{0}', github.base_ref) || '' }} run: | if [ -n "$BASE_REF" ]; then diff --git a/einops/port/overlay/tests/conftest.py b/einops/port/overlay/tests/conftest.py new file mode 100644 index 000000000..33118d584 --- /dev/null +++ b/einops/port/overlay/tests/conftest.py @@ -0,0 +1,13 @@ +import os + +# Upstream's test suite refuses to run without EINOPS_TEST_BACKENDS. Default it +# to the frameworks that are installed, so `pytest` works out of the box. +if "EINOPS_TEST_BACKENDS" not in os.environ: + available = [] + for name in ["torch", "jax", "numpy", "tensorflow"]: + try: + __import__(name) + available.append(name) + except ImportError: + pass + os.environ["EINOPS_TEST_BACKENDS"] = ",".join(available) diff --git a/einops/port/port.kdl b/einops/port/port.kdl index ce9aa3e63..88c980fc6 100644 --- a/einops/port/port.kdl +++ b/einops/port/port.kdl @@ -8,9 +8,6 @@ source \ // Retain only the Python package and its upstream tests. prune keep="einops/**" -// Drop upstream's pip-install test runner; it is not a test. -delete in="einops/tests/run_tests.py" - // Lift upstream tests out of the package into the kernel test root. move \ from="einops/tests" \ @@ -43,33 +40,32 @@ relativize_imports \ in="tests/{__init__,test_*}.py" \ package_root="tests" -// Resolve product imports through the locally built Hub kernel. +// Resolve product imports through the locally built Hub kernel. The package +// init binds it privately so `from . import *` cannot pick it up. kernelize_imports \ - in="tests/{__init__,test_*}.py" \ + in="tests/__init__.py" \ + package="einops" \ + binding="_einops" \ + kernel="kernels-community/einops" \ + version=1 + +kernelize_imports \ + in="tests/test_*.py" \ package="einops" \ kernel="kernels-community/einops" \ version=1 -// Default to the two Kernel CI backends when upstream's env var is absent. -replace in="tests/__init__.py" \ - find=#" raise RuntimeError(f"Testing frameworks were not specified, env var {ENVVAR_NAME} not set")"# \ - with=#" return ["torch", "numpy"]"# +// Pick the test backends from the installed frameworks; upstream requires +// EINOPS_TEST_BACKENDS to be set, which Kernel CI cannot do. +overlay from="overlay" // Use the supported replacement for NumPy 2's removed in1d API. replace in="tests/test_layers.py" \ find="numpy.in1d(grad, [0, 1])" \ with="numpy.isin(grad, [0, 1])" -// Select every upstream test module for the kernel CI suite. -replace \ - in="tests/test_*.py" \ - find="import pytest" \ - with="import pytest\n\npytestmark = pytest.mark.kernels_ci" - -// Keep pytest from collecting the imported doctest helper as a marked test. -replace in="tests/test_other.py" \ - find="from doctest import testmod" \ - with="from doctest import testmod\ntestmod.__test__ = False" +// Select every upstream test for the kernel CI suite. +mark_tests in="tests/test_*.py" marker="kernels_ci" changes=63 // Emit the noarch kernel manifest and supported runtime backends. manifest \ diff --git a/einops/src/build.toml b/einops/src/build.toml index 4d322d3b0..5fd31b09b 100644 --- a/einops/src/build.toml +++ b/einops/src/build.toml @@ -1,8 +1,8 @@ [general] name = "einops" version = 1 -license = "MIT" edition = 5 +license = "MIT" upstream = "https://github.com/arogozhnikov/einops.git" backends = [ "cpu", diff --git a/einops/src/port-provenance.json b/einops/src/port-provenance.json index c7259c5e2..f07fdb1c5 100644 --- a/einops/src/port-provenance.json +++ b/einops/src/port-provenance.json @@ -1,6 +1,6 @@ { "format": 1, - "recipe": "sha256:09ea109865a0ef1de6779af7db28fc5baa2cd01d0b52ee5795f07c063881ee18", + "recipe": "sha256:a1a1accf96eb178d5a2ec8470b70741a3dad80fc60a30e8df37f4ce7999789ac", "runner": "kernel-port 0.1.0", "sources": { "upstream": { @@ -8,5 +8,5 @@ "commit": "ad2c8d67d08cbebb45e1aff2be21dc88635fc695" } }, - "tree": "sha256:bd455bcb9afd28fb5c0c2b3c6c9013739fc71554017a2035e5bd782398c4dd62" + "tree": "sha256:efc73241cf8f4c4b74ea12b12906604665449c1fb40b696196397bf8390e55dd" } diff --git a/einops/src/tests/__init__.py b/einops/src/tests/__init__.py index 977745afb..8972f0707 100644 --- a/einops/src/tests/__init__.py +++ b/einops/src/tests/__init__.py @@ -2,23 +2,15 @@ Common utils for testing. These functions allow testing only some frameworks, not all. """ -__kernel_port_einops_root = None -def __kernel_port_einops(module=""): - global __kernel_port_einops_root - if __kernel_port_einops_root is None: - __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) - root = __kernel_port_einops_root - if not module: - return root - return __import__("importlib").import_module(root.__name__ + "." + module) - import logging import os from functools import lru_cache from typing import List, Tuple -_backends = getattr(__kernel_port_einops(), "_backends") +import kernels +_einops = kernels.get_kernel("kernels-community/einops", version=1) +_backends = _einops._backends import warnings __author__ = "Alex Rogozhnikov" @@ -55,7 +47,7 @@ def unparse_backends(backend_names: List[str]) -> Tuple[str, str]: @lru_cache(maxsize=1) def parse_backends_to_test() -> List[str]: if ENVVAR_NAME not in os.environ: - return ["torch", "numpy"] + 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: diff --git a/einops/src/tests/conftest.py b/einops/src/tests/conftest.py new file mode 100644 index 000000000..33118d584 --- /dev/null +++ b/einops/src/tests/conftest.py @@ -0,0 +1,13 @@ +import os + +# Upstream's test suite refuses to run without EINOPS_TEST_BACKENDS. Default it +# to the frameworks that are installed, so `pytest` works out of the box. +if "EINOPS_TEST_BACKENDS" not in os.environ: + available = [] + for name in ["torch", "jax", "numpy", "tensorflow"]: + try: + __import__(name) + available.append(name) + except ImportError: + pass + os.environ["EINOPS_TEST_BACKENDS"] = ",".join(available) diff --git a/einops/src/tests/run_tests.py b/einops/src/tests/run_tests.py new file mode 100644 index 000000000..cf7a4eb6e --- /dev/null +++ b/einops/src/tests/run_tests.py @@ -0,0 +1,85 @@ +""" +Runs tests that are appropriate for framework. +""" + +import os +import sys +from subprocess import Popen +from pathlib import Path + +__author__ = "Alex Rogozhnikov" + + +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 + + +def main(): + _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]", "flax"], + "tensorflow": ["tensorflow"], + "cupy": ["cupy"], + # switch to stable paddlepaddle, because of https://github.com/PaddlePaddle/Paddle/issues/63927 + # "paddle": ["paddlepaddle==0.0.0 -f https://www.paddlepaddle.org.cn/whl/linux/cpu-mkl/develop.html"], + "paddle": ["paddlepaddle"], + "oneflow": ["oneflow==0.9.0"], + "pytensor": ["pytensor"], + } + + usage = f""" + Usage: python -m einops.tests.run_tests [--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) + return + else: + synonyms = { + "tf": "tensorflow", + "pytorch": "torch", + "paddlepaddle": "paddle", + } + frameworks = [synonyms.get(f, f) for f in frameworks] + wrong_frameworks = [f for f in frameworks if f not in framework_name2installation] + if wrong_frameworks: + print(usage) + raise RuntimeError(f"Unrecognized frameworks: {wrong_frameworks}") + + 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))) + + 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 an envvar EINOPS_TEST_BACKENDS + from einops.tests import unparse_backends + + envvar_name, envvar_value = unparse_backends(backend_names=frameworks) + return_code = run( + "python -m pytest .", + **{envvar_name: envvar_value}, + ) + assert return_code == 0 + + +if __name__ == "__main__": + main() diff --git a/einops/src/tests/test_einsum.py b/einops/src/tests/test_einsum.py index 97a1257ef..addf117b3 100644 --- a/einops/src/tests/test_einsum.py +++ b/einops/src/tests/test_einsum.py @@ -1,20 +1,12 @@ -__kernel_port_einops_root = None -def __kernel_port_einops(module=""): - global __kernel_port_einops_root - if __kernel_port_einops_root is None: - __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) - root = __kernel_port_einops_root - if not module: - return root - return __import__("importlib").import_module(root.__name__ + "." + module) - from typing import Any, Callable from . import collect_test_backends -_compactify_pattern_for_einsum, einsum, EinopsError = getattr(__kernel_port_einops("einops"), "_compactify_pattern_for_einsum"), getattr(__kernel_port_einops("einops"), "einsum"), getattr(__kernel_port_einops("einops"), "EinopsError") +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +_compactify_pattern_for_einsum = einops.einops._compactify_pattern_for_einsum +einsum = einops.einops.einsum +EinopsError = einops.einops.EinopsError import numpy as np import pytest - -pytestmark = pytest.mark.kernels_ci import string @@ -176,6 +168,7 @@ def __call__(self, function: Callable): ] +@pytest.mark.kernels_ci def test_layer(): for backend in collect_test_backends(layers=True, symbolic=False): if backend.framework_name in ["tensorflow", "torch", "oneflow", "paddle"]: @@ -203,6 +196,7 @@ def test_layer(): ] +@pytest.mark.kernels_ci def test_functional(): # Functional tests: backends = filter(lambda x: x.framework_name in valid_backends_functional, collect_test_backends()) @@ -238,6 +232,7 @@ def test_functional(): np.testing.assert_array_almost_equal(predicted_out_array, true_out_array, decimal=5) +@pytest.mark.kernels_ci def test_functional_symbolic(): backends = filter( lambda x: x.framework_name in valid_backends_functional, collect_test_backends(symbolic=True, layers=False) @@ -270,6 +265,7 @@ def test_functional_symbolic(): np.testing.assert_array_almost_equal(predicted_out_data, expected_out_data, decimal=5) +@pytest.mark.kernels_ci def test_functional_errors(): # Specific backend does not matter, as errors are raised # during the pattern creation. diff --git a/einops/src/tests/test_examples.py b/einops/src/tests/test_examples.py index 8cda27f01..4a66431bc 100644 --- a/einops/src/tests/test_examples.py +++ b/einops/src/tests/test_examples.py @@ -1,23 +1,16 @@ -__kernel_port_einops_root = None -def __kernel_port_einops(module=""): - global __kernel_port_einops_root - if __kernel_port_einops_root is None: - __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) - root = __kernel_port_einops_root - if not module: - return root - return __import__("importlib").import_module(root.__name__ + "." + module) - import numpy import pytest -pytestmark = pytest.mark.kernels_ci - -rearrange, parse_shape, reduce = getattr(__kernel_port_einops(), "rearrange"), getattr(__kernel_port_einops(), "parse_shape"), getattr(__kernel_port_einops(), "reduce") +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +rearrange = einops.rearrange +parse_shape = einops.parse_shape +reduce = einops.reduce from . import is_backend_tested from .test_ops import imp_op_backends +@pytest.mark.kernels_ci def test_rearrange_examples(): def test1(x): # transpose @@ -203,6 +196,7 @@ def tensor_train_example_numpy(): assert numpy.allclose(y1, y3) +@pytest.mark.kernels_ci def test_pytorch_yolo_fragment(): if not is_backend_tested("torch"): pytest.skip() diff --git a/einops/src/tests/test_layers.py b/einops/src/tests/test_layers.py index 46afdd069..eab53c908 100644 --- a/einops/src/tests/test_layers.py +++ b/einops/src/tests/test_layers.py @@ -1,22 +1,14 @@ -__kernel_port_einops_root = None -def __kernel_port_einops(module=""): - global __kernel_port_einops_root - if __kernel_port_einops_root is None: - __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) - root = __kernel_port_einops_root - if not module: - return root - return __import__("importlib").import_module(root.__name__ + "." + module) - import pickle from collections import namedtuple import numpy import pytest -pytestmark = pytest.mark.kernels_ci - -rearrange, reduce, EinopsError = getattr(__kernel_port_einops(), "rearrange"), getattr(__kernel_port_einops(), "reduce"), getattr(__kernel_port_einops(), "EinopsError") +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +rearrange = einops.rearrange +reduce = einops.reduce +EinopsError = einops.EinopsError from . import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS __author__ = "Alex Rogozhnikov" @@ -45,6 +37,7 @@ def __kernel_port_einops(module=""): ] +@pytest.mark.kernels_ci def test_rearrange_imperative(): for backend in collect_test_backends(symbolic=False, layers=True): print("Test layer for ", backend.framework_name) @@ -77,6 +70,7 @@ def test_rearrange_imperative(): assert numpy.allclose(backend.to_numpy(variable.grad), 1) +@pytest.mark.kernels_ci def test_rearrange_symbolic(): for backend in collect_test_backends(symbolic=True, layers=True): print("Test layer for ", backend.framework_name) @@ -117,6 +111,7 @@ def test_rearrange_symbolic(): ] +@pytest.mark.kernels_ci def test_reduce_imperative(): for backend in collect_test_backends(symbolic=False, layers=True): print("Test layer for ", backend.framework_name) @@ -158,6 +153,7 @@ def test_reduce_imperative(): assert numpy.sum(grad) > 0.5 +@pytest.mark.kernels_ci def test_reduce_symbolic(): for backend in collect_test_backends(symbolic=True, layers=True): print("Test layer for ", backend.framework_name) @@ -189,7 +185,7 @@ def create_torch_model(use_reduce=False, add_scripted_layer=False): pytest.skip() else: from torch.nn import Sequential, Conv2d, MaxPool2d, Linear, ReLU - Rearrange, Reduce, EinMix = getattr(__kernel_port_einops("layers.torch"), "Rearrange"), getattr(__kernel_port_einops("layers.torch"), "Reduce"), getattr(__kernel_port_einops("layers.torch"), "EinMix") + Rearrange, Reduce, EinMix = einops.layers.torch.Rearrange, einops.layers.torch.Reduce, einops.layers.torch.EinMix import torch.jit return Sequential( @@ -210,6 +206,7 @@ def create_torch_model(use_reduce=False, add_scripted_layer=False): ) +@pytest.mark.kernels_ci def test_torch_layer(): if not is_backend_tested("torch"): pytest.skip() @@ -236,6 +233,7 @@ def test_torch_layer(): torch.testing.assert_close(model1(input + 1), model4(input + 1), atol=1e-3, rtol=1e-3) +@pytest.mark.kernels_ci def test_torch_layers_scripting(): if not is_backend_tested("torch"): pytest.skip() @@ -250,6 +248,7 @@ def test_torch_layers_scripting(): torch.testing.assert_close(model1(input), model2(input), atol=1e-3, rtol=1e-3) +@pytest.mark.kernels_ci def test_keras_layer(): if not is_backend_tested("tensorflow"): pytest.skip() @@ -261,7 +260,7 @@ def test_keras_layer(): pytest.skip() from tensorflow.keras.models import Sequential from tensorflow.keras.layers import Conv2D as Conv2d, Dense as Linear, ReLU - Rearrange, Reduce, EinMix, keras_custom_objects = getattr(__kernel_port_einops("layers.keras"), "Rearrange"), getattr(__kernel_port_einops("layers.keras"), "Reduce"), getattr(__kernel_port_einops("layers.keras"), "EinMix"), getattr(__kernel_port_einops("layers.keras"), "keras_custom_objects") + Rearrange, Reduce, EinMix, keras_custom_objects = einops.layers.keras.Rearrange, einops.layers.keras.Reduce, einops.layers.keras.EinMix, einops.layers.keras.keras_custom_objects def create_keras_model(): return Sequential( @@ -310,6 +309,7 @@ def create_keras_model(): numpy.testing.assert_allclose(model1.predict_on_batch(input), model4.predict_on_batch(input)) +@pytest.mark.kernels_ci def test_flax_layers(): """ One-off simple tests for Flax layers. @@ -323,7 +323,7 @@ def test_flax_layers(): import flax from flax import linen as nn - EinMix, Reduce, Rearrange = getattr(__kernel_port_einops("layers.flax"), "EinMix"), getattr(__kernel_port_einops("layers.flax"), "Reduce"), getattr(__kernel_port_einops("layers.flax"), "Rearrange") + EinMix, Reduce, Rearrange = einops.layers.flax.EinMix, einops.layers.flax.Reduce, einops.layers.flax.Rearrange class NN(nn.Module): @nn.compact @@ -357,11 +357,12 @@ def eval_at_point(params): _loaded = flax.serialization.from_bytes(params, fbytes) +@pytest.mark.kernels_ci def test_einmix_decomposition(): """ Testing that einmix correctly decomposes into smaller transformations. """ - _EinmixDebugger = getattr(__kernel_port_einops("layers._einmix"), "_EinmixDebugger") + _EinmixDebugger = einops.layers._einmix._EinmixDebugger mixin1 = _EinmixDebugger( "a b c d e -> e d c b a", @@ -447,11 +448,12 @@ def test_einmix_decomposition(): assert mixin7.saved_bias_shape == [1, 4, 2] # (a) d b, ellipsis does not participate +@pytest.mark.kernels_ci def test_einmix_restrictions(): """ Testing different cases """ - _EinmixDebugger = getattr(__kernel_port_einops("layers._einmix"), "_EinmixDebugger") + _EinmixDebugger = einops.layers._einmix._EinmixDebugger with pytest.raises(EinopsError): _EinmixDebugger( diff --git a/einops/src/tests/test_ops.py b/einops/src/tests/test_ops.py index ef0196242..1873524e1 100644 --- a/einops/src/tests/test_ops.py +++ b/einops/src/tests/test_ops.py @@ -1,23 +1,16 @@ -__kernel_port_einops_root = None -def __kernel_port_einops(module=""): - global __kernel_port_einops_root - if __kernel_port_einops_root is None: - __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) - root = __kernel_port_einops_root - if not module: - return root - return __import__("importlib").import_module(root.__name__ + "." + module) - import itertools import numpy import numpy as np import pytest -pytestmark = pytest.mark.kernels_ci - -EinopsError = getattr(__kernel_port_einops(), "EinopsError") -rearrange, reduce, repeat, _enumerate_directions = getattr(__kernel_port_einops("einops"), "rearrange"), getattr(__kernel_port_einops("einops"), "reduce"), getattr(__kernel_port_einops("einops"), "repeat"), getattr(__kernel_port_einops("einops"), "_enumerate_directions") +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +rearrange = einops.einops.rearrange +reduce = einops.einops.reduce +repeat = einops.einops.repeat +_enumerate_directions = einops.einops._enumerate_directions from . import collect_test_backends, is_backend_tested, FLOAT_REDUCTIONS as REDUCTIONS imp_op_backends = collect_test_backends(symbolic=False, layers=False) @@ -51,6 +44,7 @@ def __kernel_port_einops(module=""): ] +@pytest.mark.kernels_ci def test_collapsed_ellipsis_errors_out(): x = numpy.zeros([1, 1, 1, 1, 1]) rearrange(x, "a b c d ... -> a b c ... d") @@ -62,6 +56,7 @@ def test_collapsed_ellipsis_errors_out(): rearrange(x, "(...) -> (...)") +@pytest.mark.kernels_ci def test_ellipsis_ops_numpy(): x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) for pattern in identity_patterns: @@ -107,6 +102,7 @@ def operation(x): check_equal(numpy_result, backend_result) +@pytest.mark.kernels_ci def test_ellipsis_ops_imperative(): """Checking various patterns against numpy""" x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) @@ -124,9 +120,10 @@ def test_ellipsis_ops_imperative(): ) +@pytest.mark.kernels_ci def test_rearrange_array_api(): import numpy as xp - AA = getattr(__kernel_port_einops(), "array_api") + AA = einops.array_api if xp.__version__ < "2.0.0": pytest.skip() @@ -138,9 +135,10 @@ def test_rearrange_array_api(): assert numpy.array_equal(AA.asnumpy(result + 0), expected) +@pytest.mark.kernels_ci def test_reduce_array_api(): import numpy as xp - AA = getattr(__kernel_port_einops(), "array_api") + AA = einops.array_api if xp.__version__ < "2.0.0": pytest.skip() @@ -153,6 +151,7 @@ def test_reduce_array_api(): assert numpy.array_equal(AA.asnumpy(np.asarray(result + 0)), expected) +@pytest.mark.kernels_ci def test_rearrange_consistency_numpy(): shape = [1, 2, 3, 5, 7, 11] x = numpy.arange(numpy.prod(shape)).reshape(shape) @@ -191,6 +190,7 @@ def test_rearrange_consistency_numpy(): assert x2[0, 1, 2] == result[1, 2, 0] +@pytest.mark.kernels_ci def test_rearrange_permutations_numpy(): # tests random permutation of axes against two independent numpy ways for n_axes in range(1, 10): @@ -219,6 +219,7 @@ def test_rearrange_permutations_numpy(): assert numpy.array_equal(result, expected_result) +@pytest.mark.kernels_ci def test_reduction_imperatives(): for backend in imp_op_backends: print("Reduction tests for ", backend.framework_name) @@ -255,6 +256,7 @@ def test_reduction_imperatives(): assert numpy.allclose(result, expected_result), f"Failed at {pattern}" +@pytest.mark.kernels_ci def test_reduction_symbolic(): for backend in sym_op_backends: print("Reduction tests for ", backend.framework_name) @@ -308,6 +310,7 @@ def test_reduction_symbolic(): assert numpy.allclose(result, expected_numpy_result) +@pytest.mark.kernels_ci def test_reduction_stress_imperatives(): for backend in imp_op_backends: print("Stress-testing reduction for ", backend.framework_name) @@ -340,6 +343,7 @@ def test_reduction_stress_imperatives(): check_op_against_numpy(backend, x, pattern, reduction=reduction, axes_lengths={}, is_symbolic=False) +@pytest.mark.kernels_ci def test_reduction_with_callable_imperatives(): x_numpy = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]).astype("float32") x_numpy /= x_numpy.max() @@ -365,7 +369,7 @@ def logsumexp_numpy(x, tuple_of_axes): y = numpy.sum(y, axis=tuple_of_axes) return numpy.log(y) + minused - TorchBackend, TensorflowBackend, TFKerasBackend, NumpyBackend = getattr(__kernel_port_einops("_backends"), "TorchBackend"), getattr(__kernel_port_einops("_backends"), "TensorflowBackend"), getattr(__kernel_port_einops("_backends"), "TFKerasBackend"), getattr(__kernel_port_einops("_backends"), "NumpyBackend") + TorchBackend, TensorflowBackend, TFKerasBackend, NumpyBackend = einops._backends.TorchBackend, einops._backends.TensorflowBackend, einops._backends.TFKerasBackend, einops._backends.NumpyBackend backend2callback = { TorchBackend.framework_name: logsumexp_torch, @@ -391,6 +395,7 @@ def logsumexp_numpy(x, tuple_of_axes): ) +@pytest.mark.kernels_ci def test_enumerating_directions(): for backend in imp_op_backends: print("testing directions for", backend.framework_name) @@ -405,6 +410,7 @@ def test_enumerating_directions(): assert numpy.allclose(ax1, ax2) +@pytest.mark.kernels_ci def test_concatenations_and_stacking(): for backend in imp_op_backends: print("testing shapes for ", backend.framework_name) @@ -424,6 +430,7 @@ def test_concatenations_and_stacking(): assert numpy.array_equal(result1, backend.to_numpy(result2)) +@pytest.mark.kernels_ci def test_gradients_imperatives(): # lazy - just checking reductions for reduction in REDUCTIONS: @@ -451,6 +458,7 @@ def test_gradients_imperatives(): assert numpy.allclose(grad1, grad2), [name1, name2, "provided different gradients"] +@pytest.mark.kernels_ci def test_tiling_imperatives(): for backend in imp_op_backends: print("Tiling tests for ", backend.framework_name) @@ -468,6 +476,7 @@ def test_tiling_imperatives(): assert numpy.array_equal(result, expected) +@pytest.mark.kernels_ci def test_tiling_symbolic(): for backend in sym_op_backends: print("Tiling tests for ", backend.framework_name) @@ -513,6 +522,7 @@ def check_reversion(x, repeat_pattern, **sizes): assert numpy.array_equal(x, reduced_max) +@pytest.mark.kernels_ci def test_repeat_numpy(): # check repeat vs reduce. Repeat works ok if reverse reduction with min and max work well x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) @@ -522,6 +532,7 @@ def test_repeat_numpy(): check_reversion(x, pattern, **axis_dimensions) +@pytest.mark.kernels_ci def test_repeat_imperatives(): x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) for backend in imp_op_backends: @@ -535,6 +546,7 @@ def test_repeat_imperatives(): assert numpy.array_equal(result, expected) +@pytest.mark.kernels_ci def test_repeat_symbolic(): x = numpy.arange(2 * 3 * 5).reshape([2, 3, 5]) @@ -549,9 +561,10 @@ def test_repeat_symbolic(): assert numpy.array_equal(result, expected) +@pytest.mark.kernels_ci def test_repeat_array_api(): import numpy as xp - AA = getattr(__kernel_port_einops(), "array_api") + AA = einops.array_api if xp.__version__ < "2.0.0": pytest.skip() @@ -576,12 +589,14 @@ def test_repeat_array_api(): ] +@pytest.mark.kernels_ci def test_anonymous_axes(): x = numpy.arange(1 * 2 * 4 * 6).reshape([1, 2, 4, 6]) for pattern, axis_dimensions in test_cases_repeat_anonymous: check_reversion(x, pattern, **axis_dimensions) +@pytest.mark.kernels_ci def test_list_inputs(): x = numpy.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6]) @@ -599,6 +614,7 @@ def test_list_inputs(): ) +@pytest.mark.kernels_ci def test_torch_compile_with_dynamic_shape(): if not is_backend_tested("torch"): pytest.skip() @@ -630,6 +646,7 @@ def bit_count(x): return sum((x >> i) & 1 for i in range(20)) +@pytest.mark.kernels_ci def test_reduction_imperatives_booleans(): """Checks that any/all reduction works in all frameworks""" x_np = numpy.asarray([(bit_count(x) % 2) == 0 for x in range(2**6)]).reshape([2] * 6) diff --git a/einops/src/tests/test_other.py b/einops/src/tests/test_other.py index ff4346851..063f3298a 100644 --- a/einops/src/tests/test_other.py +++ b/einops/src/tests/test_other.py @@ -1,37 +1,27 @@ -__kernel_port_einops_root = None -def __kernel_port_einops(module=""): - global __kernel_port_einops_root - if __kernel_port_einops_root is None: - __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) - root = __kernel_port_einops_root - if not module: - return root - return __import__("importlib").import_module(root.__name__ + "." + module) - from doctest import testmod -testmod.__test__ = False import numpy import pytest -pytestmark = pytest.mark.kernels_ci - -einops = __kernel_port_einops() -einops = __kernel_port_einops(); __kernel_port_einops("layers") -einops = __kernel_port_einops(); __kernel_port_einops("parsing") -AbstractBackend = getattr(__kernel_port_einops("_backends"), "AbstractBackend") -rearrange, parse_shape, _optimize_transformation = getattr(__kernel_port_einops("einops"), "rearrange"), getattr(__kernel_port_einops("einops"), "parse_shape"), getattr(__kernel_port_einops("einops"), "_optimize_transformation") +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +AbstractBackend = einops._backends.AbstractBackend +rearrange = einops.einops.rearrange +parse_shape = einops.einops.parse_shape +_optimize_transformation = einops.einops._optimize_transformation from . import collect_test_backends, is_backend_tested __author__ = "Alex Rogozhnikov" +@pytest.mark.kernels_ci def test_doctests_examples(): # tests docstrings, additionally testmod(einops.layers, raise_on_error=True, extraglobs=dict(np=numpy)) testmod(einops.einops, raise_on_error=True, extraglobs=dict(np=numpy)) +@pytest.mark.kernels_ci def test_backends_installed(): """ This test will fail if some of backends are not installed or can't be imported @@ -52,6 +42,7 @@ def test_backends_installed(): assert len(errors) == 0, errors +@pytest.mark.kernels_ci def test_optimize_transformations_numpy(): print("Testing optimizations") shapes = [[2] * n_dimensions for n_dimensions in range(14)] @@ -89,6 +80,7 @@ def test_optimize_transformations_numpy(): x_np = numpy.zeros([10, 20, 30, 40]) +@pytest.mark.kernels_ci def test_parse_shape_imperative(): for backend in _IMPERATIVE_BACKENDS: print("Shape parsing for ", backend.framework_name) @@ -98,6 +90,7 @@ def test_parse_shape_imperative(): assert parsed1 != dict(a=1, b=20, c=30, d=40) != parsed2 +@pytest.mark.kernels_ci def test_underscore(): for backend in _IMPERATIVE_BACKENDS: parsed1 = parse_shape(x_np, "_ _ _ _") @@ -105,6 +98,7 @@ def test_underscore(): assert parsed1 == parsed2 == dict() +@pytest.mark.kernels_ci def test_underscore_one(): for backend in _IMPERATIVE_BACKENDS: parsed1 = parse_shape(x_np, "_ _ _ hello") @@ -112,6 +106,7 @@ def test_underscore_one(): assert parsed1 == parsed2 == dict(hello=40) +@pytest.mark.kernels_ci def test_underscore_several(): for backend in _IMPERATIVE_BACKENDS: parsed1 = parse_shape(x_np, "_ _ a1 a1a111a") @@ -119,6 +114,7 @@ def test_underscore_several(): assert parsed1 == parsed2 == dict(a1=30, a1a111a=40) +@pytest.mark.kernels_ci def test_repeating(): with pytest.raises(einops.EinopsError): parse_shape(x_np, "a a b b") @@ -128,6 +124,7 @@ def test_repeating(): parse_shape(backend.from_numpy(x_np), "a a b b") +@pytest.mark.kernels_ci def test_ellipsis(): for backend in _IMPERATIVE_BACKENDS: for shape, pattern, expected in [ @@ -149,6 +146,7 @@ def test_ellipsis(): assert parsed1 == parsed2 == expected +@pytest.mark.kernels_ci def test_parse_with_anonymous_axes(): for backend in _IMPERATIVE_BACKENDS: for shape, pattern, expected in [ @@ -162,6 +160,7 @@ def test_parse_with_anonymous_axes(): assert parsed1 == parsed2 == expected +@pytest.mark.kernels_ci def test_failures(): for backend in _IMPERATIVE_BACKENDS: # every test should fail @@ -188,6 +187,7 @@ def test_failures(): _SYMBOLIC_BACKENDS = [backend for backend in _SYMBOLIC_BACKENDS if backend.framework_name != "tensorflow.keras"] +@pytest.mark.kernels_ci @pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) def test_parse_shape_symbolic(backend): for shape in [ @@ -218,6 +218,7 @@ def test_parse_shape_symbolic(backend): assert numpy.allclose(result, 0) +@pytest.mark.kernels_ci @pytest.mark.parametrize("backend", _SYMBOLIC_BACKENDS) def test_parse_shape_symbolic_ellipsis(backend): for static_shape, shape, pattern, expected in [ @@ -244,6 +245,7 @@ def test_parse_shape_symbolic_ellipsis(backend): assert out_shape == expected +@pytest.mark.kernels_ci def test_is_float_type(): backends = collect_test_backends(symbolic=False, layers=False) backends += collect_test_backends(symbolic=False, layers=True) @@ -255,6 +257,7 @@ def test_is_float_type(): assert backend.is_float_type(input) == is_float, (dtype, backend, input.dtype) +@pytest.mark.kernels_ci def test_torch_compile(): """ Test ensures that allow_ops_in_compiled_graph allows compiling in a single graph @@ -267,8 +270,8 @@ def test_torch_compile(): pytest.skip() import torch from torch import nn - repeat, reduce, pack, unpack, einsum = getattr(__kernel_port_einops(), "repeat"), getattr(__kernel_port_einops(), "reduce"), getattr(__kernel_port_einops(), "pack"), getattr(__kernel_port_einops(), "unpack"), getattr(__kernel_port_einops(), "einsum") - allow_ops_in_compiled_graph = getattr(__kernel_port_einops("_torch_specific"), "allow_ops_in_compiled_graph") + repeat, reduce, pack, unpack, einsum = einops.repeat, einops.reduce, einops.pack, einops.unpack, einops.einsum + allow_ops_in_compiled_graph = einops._torch_specific.allow_ops_in_compiled_graph allow_ops_in_compiled_graph() diff --git a/einops/src/tests/test_packing.py b/einops/src/tests/test_packing.py index 09869c3fb..d76553475 100644 --- a/einops/src/tests/test_packing.py +++ b/einops/src/tests/test_packing.py @@ -1,22 +1,15 @@ -__kernel_port_einops_root = None -def __kernel_port_einops(module=""): - global __kernel_port_einops_root - if __kernel_port_einops_root is None: - __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) - root = __kernel_port_einops_root - if not module: - return root - return __import__("importlib").import_module(root.__name__ + "." + module) - import dataclasses import typing import numpy as np import pytest -pytestmark = pytest.mark.kernels_ci - -EinopsError, asnumpy, pack, unpack = getattr(__kernel_port_einops(), "EinopsError"), getattr(__kernel_port_einops(), "asnumpy"), getattr(__kernel_port_einops(), "pack"), getattr(__kernel_port_einops(), "unpack") +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +asnumpy = einops.asnumpy +pack = einops.pack +unpack = einops.unpack from . import collect_test_backends @@ -71,6 +64,7 @@ def __exit__(self, exc_type, exc_val, exc_tb): return True +@pytest.mark.kernels_ci def test_numpy_trivial(H=13, W=17): def rand(*shape): return np.random.random(shape) @@ -154,6 +148,7 @@ def selfcheck(self): ] +@pytest.mark.kernels_ci def test_pack_unpack_with_numpy(): case: UnpackTestCase @@ -217,6 +212,7 @@ def test_pack_unpack_with_numpy(): unpack_and_pack(x, [[2, -1], [1, 5]], pattern) +@pytest.mark.kernels_ci def test_pack_unpack_against_numpy(): for backend in collect_test_backends(symbolic=False, layers=False): print(f"test packing against numpy for {backend.framework_name}") @@ -281,8 +277,9 @@ def test_pack_unpack_against_numpy(): unpack_and_pack(x, [[2, -1], [1, 5]], pattern) +@pytest.mark.kernels_ci def test_pack_unpack_array_api(): - AA = getattr(__kernel_port_einops(), "array_api") + AA = einops.array_api import numpy as xp if xp.__version__ < "2.0.0": diff --git a/einops/src/tests/test_parsing.py b/einops/src/tests/test_parsing.py index 6687adee8..491126793 100644 --- a/einops/src/tests/test_parsing.py +++ b/einops/src/tests/test_parsing.py @@ -1,19 +1,11 @@ -__kernel_port_einops_root = None -def __kernel_port_einops(module=""): - global __kernel_port_einops_root - if __kernel_port_einops_root is None: - __kernel_port_einops_root = __import__("kernels").get_kernel("kernels-community/einops", version=1) - root = __kernel_port_einops_root - if not module: - return root - return __import__("importlib").import_module(root.__name__ + "." + module) - import pytest -pytestmark = pytest.mark.kernels_ci - -EinopsError = getattr(__kernel_port_einops(), "EinopsError") -ParsedExpression, AnonymousAxis, _ellipsis = getattr(__kernel_port_einops("parsing"), "ParsedExpression"), getattr(__kernel_port_einops("parsing"), "AnonymousAxis"), getattr(__kernel_port_einops("parsing"), "_ellipsis") +import kernels +einops = kernels.get_kernel("kernels-community/einops", version=1) +EinopsError = einops.EinopsError +ParsedExpression = einops.parsing.ParsedExpression +AnonymousAxis = einops.parsing.AnonymousAxis +_ellipsis = einops.parsing._ellipsis __author__ = "Alex Rogozhnikov" @@ -27,6 +19,7 @@ def __eq__(self, other): return isinstance(other, AnonymousAxis) and self.value == other.value +@pytest.mark.kernels_ci def test_anonymous_axes(): a, b = AnonymousAxis("2"), AnonymousAxis("2") assert a != b @@ -36,6 +29,7 @@ def test_anonymous_axes(): assert [a, 2, b] == [c, 2, c] +@pytest.mark.kernels_ci def test_elementary_axis_name(): for name in [ "a", @@ -57,6 +51,7 @@ def test_elementary_axis_name(): assert not ParsedExpression.check_axis_name(name) +@pytest.mark.kernels_ci def test_invalid_expressions(): # double ellipsis should raise an error ParsedExpression("... a b c d") @@ -90,6 +85,7 @@ def test_invalid_expressions(): ParsedExpression("pre...") +@pytest.mark.kernels_ci def test_parse_expression(): parsed = ParsedExpression("a1 b1 c1 d1") assert parsed.identifiers == {"a1", "b1", "c1", "d1"} From 9c378a03c305b652a141f3fbf516b6e93a8f161a Mon Sep 17 00:00:00 2001 From: David Holtz Date: Mon, 5 Oct 2026 15:28:27 +0000 Subject: [PATCH 14/15] ci: run kernel-port from huggingface/kernels#877 until it merges --- .github/workflows/check-ports.yaml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/check-ports.yaml b/.github/workflows/check-ports.yaml index dab3574fd..8568cc46d 100644 --- a/.github/workflows/check-ports.yaml +++ b/.github/workflows/check-ports.yaml @@ -51,7 +51,8 @@ jobs: # Track huggingface/kernels main rather than pinning a rev, so that a # kernel-port change that breaks a port is caught immediately. # --refresh skips Nix's flake cache so every run sees the latest main. - KERNEL_PORT: "nix run --refresh github:huggingface/kernels#kernel-port --" + # TODO: switch back to main once huggingface/kernels#877 is merged. + KERNEL_PORT: "nix run --refresh github:huggingface/kernels/kernel-port-no-importlib#kernel-port --" BASE_REF: ${{ github.event_name == 'pull_request' && format('origin/{0}', github.base_ref) || '' }} run: | if [ -n "$BASE_REF" ]; then From 8707a621ff7a85a08fd39b361ad19a2307af2e56 Mon Sep 17 00:00:00 2001 From: David Holtz Date: Tue, 6 Oct 2026 15:16:03 +0000 Subject: [PATCH 15/15] fix: simplify comments and narrow test --- .github/workflows/check-ports.yaml | 22 ++++----------- einops/port/overlay/tests/conftest.py | 13 ++------- einops/port/port.kdl | 3 +-- einops/src/port-provenance.json | 4 +-- einops/src/tests/conftest.py | 13 ++------- scripts/check_ports.py | 39 +++------------------------ 6 files changed, 15 insertions(+), 79 deletions(-) diff --git a/.github/workflows/check-ports.yaml b/.github/workflows/check-ports.yaml index 8568cc46d..b4e57ced7 100644 --- a/.github/workflows/check-ports.yaml +++ b/.github/workflows/check-ports.yaml @@ -1,19 +1,11 @@ name: Check Ports -# A ported kernel's /src is generated from /port/port.kdl and -# committed. Regenerate it here and fail if the two disagree: a stale overlay -# reverting an in-tree fix, a hand-edited build.toml field the manifest op -# drops, an overlay file edited without its generated twin, or a recipe edited -# without rerunning the port (src/port-provenance.json records the recipe hash, -# so that case shows up as a diff in the provenance file). -# -# Pull requests only check the ports they touch; pushes to main, the daily -# run and manual runs check every port. +# Verify committed src trees match their port recipes. PRs check changed ports; +# checker changes, pushes to main, scheduled runs and manual runs check all ports. on: push: branches: [main] - # kernel-port tracks huggingface/kernels main, so also run without a change - # here: a kernel-port regression should surface even when no port PR is open. + # Catch kernel-port regressions even when no port changes. schedule: - cron: "0 6 * * *" pull_request: @@ -44,13 +36,9 @@ jobs: - uses: DeterminateSystems/nix-installer-action@ef8a148080ab6020fd15196c2084a2eea5ff2d25 # v22 - # Each recipe pins its own upstream commit; the runner verifies the - # origin URL, HEAD and cleanliness of the checkout it is handed. - - name: Regenerate every port and compare + - name: Regenerate selected ports and compare env: - # Track huggingface/kernels main rather than pinning a rev, so that a - # kernel-port change that breaks a port is caught immediately. - # --refresh skips Nix's flake cache so every run sees the latest main. + # Refresh the runner to catch upstream changes. # TODO: switch back to main once huggingface/kernels#877 is merged. KERNEL_PORT: "nix run --refresh github:huggingface/kernels/kernel-port-no-importlib#kernel-port --" BASE_REF: ${{ github.event_name == 'pull_request' && format('origin/{0}', github.base_ref) || '' }} diff --git a/einops/port/overlay/tests/conftest.py b/einops/port/overlay/tests/conftest.py index 33118d584..3fe86a864 100644 --- a/einops/port/overlay/tests/conftest.py +++ b/einops/port/overlay/tests/conftest.py @@ -1,13 +1,4 @@ import os -# Upstream's test suite refuses to run without EINOPS_TEST_BACKENDS. Default it -# to the frameworks that are installed, so `pytest` works out of the box. -if "EINOPS_TEST_BACKENDS" not in os.environ: - available = [] - for name in ["torch", "jax", "numpy", "tensorflow"]: - try: - __import__(name) - available.append(name) - except ImportError: - pass - os.environ["EINOPS_TEST_BACKENDS"] = ",".join(available) +# Use the CI backends by default, while allowing local overrides. +os.environ.setdefault("EINOPS_TEST_BACKENDS", "torch,numpy") diff --git a/einops/port/port.kdl b/einops/port/port.kdl index 88c980fc6..89ac6b50b 100644 --- a/einops/port/port.kdl +++ b/einops/port/port.kdl @@ -55,8 +55,7 @@ kernelize_imports \ kernel="kernels-community/einops" \ version=1 -// Pick the test backends from the installed frameworks; upstream requires -// EINOPS_TEST_BACKENDS to be set, which Kernel CI cannot do. +// Default to Torch and NumPy tests, with an EINOPS_TEST_BACKENDS override. overlay from="overlay" // Use the supported replacement for NumPy 2's removed in1d API. diff --git a/einops/src/port-provenance.json b/einops/src/port-provenance.json index f07fdb1c5..c683484df 100644 --- a/einops/src/port-provenance.json +++ b/einops/src/port-provenance.json @@ -1,6 +1,6 @@ { "format": 1, - "recipe": "sha256:a1a1accf96eb178d5a2ec8470b70741a3dad80fc60a30e8df37f4ce7999789ac", + "recipe": "sha256:238127312e85c9489a67de689d3b2e59c75c456f6839afb640f344ef0e7965d5", "runner": "kernel-port 0.1.0", "sources": { "upstream": { @@ -8,5 +8,5 @@ "commit": "ad2c8d67d08cbebb45e1aff2be21dc88635fc695" } }, - "tree": "sha256:efc73241cf8f4c4b74ea12b12906604665449c1fb40b696196397bf8390e55dd" + "tree": "sha256:418031b9de86c73cadcf391017cb3f666fd816137449513b481f4af2caa6fe53" } diff --git a/einops/src/tests/conftest.py b/einops/src/tests/conftest.py index 33118d584..3fe86a864 100644 --- a/einops/src/tests/conftest.py +++ b/einops/src/tests/conftest.py @@ -1,13 +1,4 @@ import os -# Upstream's test suite refuses to run without EINOPS_TEST_BACKENDS. Default it -# to the frameworks that are installed, so `pytest` works out of the box. -if "EINOPS_TEST_BACKENDS" not in os.environ: - available = [] - for name in ["torch", "jax", "numpy", "tensorflow"]: - try: - __import__(name) - available.append(name) - except ImportError: - pass - os.environ["EINOPS_TEST_BACKENDS"] = ",".join(available) +# Use the CI backends by default, while allowing local overrides. +os.environ.setdefault("EINOPS_TEST_BACKENDS", "torch,numpy") diff --git a/scripts/check_ports.py b/scripts/check_ports.py index a7ab25a2d..d6dd27644 100644 --- a/scripts/check_ports.py +++ b/scripts/check_ports.py @@ -1,30 +1,4 @@ #!/usr/bin/env python3 -"""Verify that every ported kernel's committed tree matches its recipe. - -A ported kernel is a directory holding `port/port.kdl`. The recipe generates -`/src`, which is committed so that what ships stays reviewable. This -script re-runs the port against a fresh upstream checkout at the pinned commit -and fails if the result differs from what is committed. - -Without this, four failure modes are silent: - - - an overlay file goes stale, so regenerating reverts a fix made in-tree - - a build.toml field is edited by hand, and the manifest op drops it - - an overlay file and its generated twin are edited independently and drift - - the recipe is edited without rerunning the port; src/port-provenance.json - records the recipe hash, so this shows up as a diff in that file - -Usage: - python3 scripts/check_ports.py [kernel ...] - python3 scripts/check_ports.py --changed-since - -Defaults to every port found. --changed-since limits the check to ports whose -port/ or src/ tree differs from (or every port, if the checker itself -changed), which is what CI uses for pull requests. -Set KERNEL_PORT to override the runner command (e.g. "nix run -github:huggingface/kernels#kernel-port --"). -""" - import filecmp import os import re @@ -35,8 +9,7 @@ import tempfile from pathlib import Path -# `source repo="..." commit="..."` is the recipe's first op; the runner already -# enforces a 40-char sha and a clean checkout, so a loose match is enough here. +# Extract source arguments; kernel-port validates them. SOURCE_RE = re.compile( r"^source\b(?P(?:[^\n\\]|\\\s*\n)*)", re.MULTILINE ) @@ -56,11 +29,7 @@ def find_ports(root: Path) -> list: def changed_ports(root: Path, ref: str) -> list: - """Ports whose port/ or src/ tree differs from `ref`. - - A change to the checker itself (this script, or the workflow that pins the - runner) can break any port, so it selects all of them. - """ + """Select changed ports, or all ports when the checker or workflow changes.""" out = subprocess.run( ["git", "diff", "--name-only", f"{ref}...HEAD"], capture_output=True, @@ -90,8 +59,6 @@ def parse_source(recipe: Path) -> tuple: def checkout(repo: str, commit: str, dest: Path) -> None: - # The `source` op verifies origin URL, HEAD, and cleanliness, so clone from - # the exact URL the recipe names and land on the exact commit. subprocess.run( ["git", "clone", "--quiet", "--filter=blob:none", "--no-checkout", repo, str(dest)], check=True, @@ -125,7 +92,7 @@ def check(kernel: str, root: Path, workdir: Path) -> list: repo, commit = parse_source(recipe) print(f" {kernel}: {repo} @ {commit[:12]}") - # One checkout per (repo, commit): two kernels may share an upstream. + # Reuse checkouts for kernels sharing an upstream commit. src = workdir / f"upstream-{commit[:12]}" if not src.exists(): checkout(repo, commit, src)