Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,109 @@
|
||||
"""
|
||||
Common utils for testing.
|
||||
These functions allow testing only some frameworks, not all.
|
||||
"""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import warnings
|
||||
from functools import lru_cache
|
||||
from typing import List, Tuple
|
||||
|
||||
from einops import _backends
|
||||
|
||||
__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 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
|
||||
|
||||
|
||||
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(f"backend could not be initialized for tests: {backend_type}", stacklevel=1)
|
||||
return result
|
||||
@@ -0,0 +1,89 @@
|
||||
"""
|
||||
Runs tests that are appropriate for framework.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from subprocess import Popen
|
||||
|
||||
__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"],
|
||||
}
|
||||
if sys.platform == "darwin":
|
||||
framework_name2installation["mlx"] = ["mlx"]
|
||||
if sys.platform.startswith("linux"):
|
||||
framework_name2installation["mlx"] = ["mlx[cpu]"]
|
||||
|
||||
usage = f"""
|
||||
Usage: python -m einops.tests.run_tests <frameworks> [--pip-install]
|
||||
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()
|
||||
@@ -0,0 +1,356 @@
|
||||
import string
|
||||
from typing import Any, Callable
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from einops.einops import EinopsError, _compactify_pattern_for_einsum, einsum
|
||||
from einops.tests import collect_test_backends
|
||||
|
||||
|
||||
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,)),
|
||||
(()),
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
def test_layer():
|
||||
for backend in collect_test_backends(layers=True, symbolic=False):
|
||||
rng = np.random.default_rng()
|
||||
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 = rng.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",
|
||||
"mlx",
|
||||
]
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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.
|
||||
@@ -0,0 +1,297 @@
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from einops import parse_shape, rearrange, reduce
|
||||
from einops.tests import is_backend_tested
|
||||
from einops.tests.test_ops import imp_op_backends
|
||||
|
||||
|
||||
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 = np.arange(10 * 20 * 30 * 40).reshape([10, 20, 30, 40])
|
||||
result1 = test(x)
|
||||
result2 = backend.to_numpy(test(backend.from_numpy(x)))
|
||||
assert np.array_equal(result1, result2)
|
||||
|
||||
# now with strides
|
||||
x = np.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 = np.index_exp[::2, ::3, ::1, ::last_step]
|
||||
result1 = test(x[indexing_expression])
|
||||
result2 = backend.to_numpy(test(backend.from_numpy(x)[indexing_expression]))
|
||||
assert np.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 = np.ones([3, 4, 5, 6])
|
||||
rank = 4
|
||||
if np.__version__ < "1.15.0":
|
||||
# numpy.einsum fails here, skip test
|
||||
return
|
||||
# creating appropriate Gs
|
||||
Gs = [np.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 = np.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 np.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 np.allclose(y1, y3)
|
||||
|
||||
|
||||
def test_pytorch_yolo_fragment():
|
||||
if not is_backend_tested("torch"):
|
||||
pytest.skip()
|
||||
|
||||
import torch
|
||||
|
||||
def old_way(tensor, num_classes, num_anchors, anchors, stride_h, stride_w):
|
||||
# https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L28-L44
|
||||
bs = tensor.size(0)
|
||||
in_h = tensor.size(2)
|
||||
in_w = tensor.size(3)
|
||||
scaled_anchors = [(a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors]
|
||||
|
||||
prediction = tensor.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(tensor, num_classes, num_anchors, anchors, stride_h, stride_w):
|
||||
raw_predictions = rearrange(tensor, " b (anchor prediction) h w -> prediction b anchor h w", anchor=num_anchors)
|
||||
|
||||
anchors = torch.FloatTensor(anchors).to(tensor.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(tensor.device)
|
||||
grid_w = rearrange(torch.arange(in_w).float(), "w -> () () () w").to(tensor.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)
|
||||
|
||||
x = torch.randn([batch_size, num_anchors * (5 + num_classes), 1, 1])
|
||||
result1 = old_way(
|
||||
tensor=x,
|
||||
num_anchors=num_anchors,
|
||||
num_classes=num_classes,
|
||||
stride_h=stride_h,
|
||||
stride_w=stride_w,
|
||||
anchors=anchors,
|
||||
)
|
||||
result2 = new_way(
|
||||
tensor=x,
|
||||
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)
|
||||
@@ -0,0 +1,480 @@
|
||||
import pickle
|
||||
from collections import namedtuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from einops import EinopsError, rearrange, reduce
|
||||
from einops.tests import FLOAT_REDUCTIONS as REDUCTIONS
|
||||
from einops.tests import collect_test_backends, is_backend_tested
|
||||
|
||||
__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)],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
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 = np.arange(np.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(np.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 np.allclose(result_numpy, result1)
|
||||
assert np.allclose(result1, result2)
|
||||
|
||||
just_sum = backend.layers().Reduce("...->", reduction="sum")
|
||||
|
||||
variable = backend.from_numpy(x)
|
||||
result = just_sum(layer(variable))
|
||||
|
||||
result.backward()
|
||||
assert np.allclose(backend.to_numpy(variable.grad), 1)
|
||||
|
||||
|
||||
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 = np.arange(np.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 np.allclose(result_numpy, result1)
|
||||
|
||||
layer2 = pickle.loads(pickle.dumps(layer))
|
||||
result_symbol2 = layer2(symbol)
|
||||
result2 = backend.eval_symbol(result_symbol2, eval_inputs)
|
||||
assert np.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 = np.sum(x)
|
||||
|
||||
assert np.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)]),
|
||||
]
|
||||
|
||||
|
||||
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 = np.arange(1, 1 + np.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(np.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 np.allclose(result_numpy, result1)
|
||||
assert np.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 np.allclose(grad, 1)
|
||||
if reduction == "mean":
|
||||
assert np.allclose(grad, grad.min())
|
||||
if reduction in ["max", "min"]:
|
||||
assert np.all(np.isin(grad, [0, 1]))
|
||||
assert np.sum(grad) > 0.5
|
||||
|
||||
|
||||
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 = np.arange(1, 1 + np.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 np.allclose(result_numpy, result1)
|
||||
|
||||
layer2 = pickle.loads(pickle.dumps(layer))
|
||||
result_symbol2 = layer2(symbol)
|
||||
result2 = backend.eval_symbol(result_symbol2, eval_inputs)
|
||||
assert np.allclose(result1, result2)
|
||||
|
||||
|
||||
def create_torch_model(use_reduce=False, add_scripted_layer=False):
|
||||
if not is_backend_tested("torch"):
|
||||
pytest.skip()
|
||||
else:
|
||||
import torch.jit
|
||||
from torch.nn import Conv2d, Linear, MaxPool2d, ReLU, Sequential
|
||||
|
||||
from einops.layers.torch import EinMix, Rearrange, Reduce
|
||||
|
||||
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),
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
def test_keras_layer():
|
||||
rng = np.random.default_rng()
|
||||
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.layers import Conv2D as Conv2d
|
||||
from tensorflow.keras.layers import Dense as Linear
|
||||
from tensorflow.keras.layers import ReLU
|
||||
from tensorflow.keras.models import Sequential
|
||||
|
||||
from einops.layers.keras import EinMix, Rearrange, Reduce, 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 = rng.normal(size=[10, 32, 32, 3]).astype("float32")
|
||||
# two randomly init models should provide different outputs
|
||||
assert not np.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)
|
||||
|
||||
np.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
|
||||
np.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
|
||||
np.testing.assert_allclose(model1.predict_on_batch(input), model4.predict_on_batch(input))
|
||||
|
||||
|
||||
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 flax
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
from flax import linen as nn
|
||||
|
||||
from einops.layers.flax import EinMix, Rearrange, Reduce
|
||||
|
||||
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)
|
||||
if jax.__version__ < "0.6.0":
|
||||
tree_map = jax.tree_map
|
||||
else:
|
||||
tree_map = jax.tree.map
|
||||
|
||||
params2 = 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)
|
||||
|
||||
|
||||
def test_einmix_decomposition():
|
||||
"""
|
||||
Testing that einmix correctly decomposes into smaller transformations.
|
||||
"""
|
||||
from einops.layers._einmix import _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
|
||||
|
||||
|
||||
def test_einmix_restrictions():
|
||||
"""
|
||||
Testing different cases
|
||||
"""
|
||||
from einops.layers._einmix import _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
|
||||
@@ -0,0 +1,658 @@
|
||||
import itertools
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from einops import EinopsError
|
||||
from einops.einops import _enumerate_directions, rearrange, reduce, repeat
|
||||
from einops.tests import FLOAT_REDUCTIONS as REDUCTIONS
|
||||
from einops.tests import collect_test_backends, is_backend_tested
|
||||
|
||||
imp_op_backends = collect_test_backends(symbolic=False, layers=False)
|
||||
sym_op_backends = collect_test_backends(symbolic=True, layers=False)
|
||||
|
||||
rng = np.random.default_rng()
|
||||
|
||||
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 -> (...) "),
|
||||
]
|
||||
|
||||
|
||||
def test_collapsed_ellipsis_errors_out():
|
||||
x = np.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, "(...) -> (...)")
|
||||
|
||||
|
||||
def test_ellipsis_ops_numpy():
|
||||
x = np.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
|
||||
for pattern in identity_patterns:
|
||||
assert np.array_equal(x, rearrange(x, pattern)), pattern
|
||||
|
||||
for pattern1, pattern2 in equivalent_rearrange_patterns:
|
||||
assert np.array_equal(rearrange(x, pattern1), rearrange(x, pattern2))
|
||||
|
||||
for reduction in ["min", "max", "sum"]:
|
||||
for pattern1, pattern2 in equivalent_reduction_patterns:
|
||||
assert np.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 = np.array_equal
|
||||
p_none_dimension = 0.5
|
||||
if is_symbolic:
|
||||
symbol_shape = [d if rng.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)
|
||||
|
||||
|
||||
def test_ellipsis_ops_imperative():
|
||||
"""Checking various patterns against numpy"""
|
||||
x = np.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
|
||||
)
|
||||
|
||||
|
||||
def test_rearrange_array_api():
|
||||
import numpy as xp
|
||||
|
||||
from einops import array_api as AA
|
||||
|
||||
if xp.__version__ < "2.0.0":
|
||||
pytest.skip()
|
||||
|
||||
x = np.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 np.array_equal(AA.asnumpy(result + 0), expected)
|
||||
|
||||
|
||||
def test_reduce_array_api():
|
||||
import numpy as xp
|
||||
|
||||
from einops import array_api as AA
|
||||
|
||||
if xp.__version__ < "2.0.0":
|
||||
pytest.skip()
|
||||
|
||||
x = np.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 np.array_equal(AA.asnumpy(np.asarray(result + 0)), expected)
|
||||
|
||||
|
||||
def test_rearrange_consistency_numpy():
|
||||
shape = [1, 2, 3, 5, 7, 11]
|
||||
x = np.arange(np.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(np.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 np.array_equal(x.flatten(), result.flatten())
|
||||
|
||||
result = rearrange(x, "a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11")
|
||||
assert np.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 np.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 np.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 np.array_equal(x, result)
|
||||
|
||||
x2 = np.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]
|
||||
|
||||
|
||||
def test_rearrange_permutations_numpy():
|
||||
# tests random permutation of axes against two independent numpy ways
|
||||
for n_axes in range(1, 10):
|
||||
input = np.arange(2**n_axes).reshape([2] * n_axes)
|
||||
permutation = rng.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 rng.integers(0, 2, [10, n_axes]):
|
||||
assert input[tuple(pick)] == result[tuple(pick[permutation])]
|
||||
|
||||
for n_axes in range(1, 10):
|
||||
input = np.arange(2**n_axes).reshape([2] * n_axes)
|
||||
permutation = rng.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 = np.zeros_like(input)
|
||||
for original_axis, result_axis in enumerate(permutation):
|
||||
expected_result |= ((input >> original_axis) & 1) << result_axis
|
||||
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
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 = np.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 np.allclose(result, expected_result), f"Failed at {pattern}"
|
||||
|
||||
|
||||
def test_reduction_symbolic():
|
||||
for backend in sym_op_backends:
|
||||
print("Reduction tests for ", backend.framework_name)
|
||||
for reduction in REDUCTIONS:
|
||||
input = np.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 np.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 np.allclose(result, expected_numpy_result)
|
||||
|
||||
|
||||
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 = np.array_equal
|
||||
if reduction in ["mean", "prod"]:
|
||||
dtype = "float64"
|
||||
coincide = np.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 = rng.integers(2, 4, size=n_axes)
|
||||
permutation = rng.permutation(n_axes)
|
||||
skipped = 0 if reduction == "rearrange" else rng.integers(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 = np.arange(1, 1 + np.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)
|
||||
|
||||
|
||||
def test_reduction_with_callable_imperatives():
|
||||
x_numpy = np.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 = np.exp(y)
|
||||
y = np.sum(y, axis=tuple_of_axes)
|
||||
return np.log(y) + minused
|
||||
|
||||
from einops._backends import NumpyBackend, TensorflowBackend, TFKerasBackend, TorchBackend
|
||||
|
||||
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 np.allclose(
|
||||
output_numpy,
|
||||
backend.to_numpy(output_backend),
|
||||
)
|
||||
|
||||
|
||||
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 = np.arange(np.prod(shape)).reshape(shape)
|
||||
axes1 = _enumerate_directions(x)
|
||||
axes2 = _enumerate_directions(backend.from_numpy(x))
|
||||
assert len(axes1) == len(axes2) == len(shape)
|
||||
axes2 = [backend.to_numpy(ax) for ax in axes2]
|
||||
for ax1, ax2 in zip(axes1, axes2):
|
||||
assert ax1.shape == ax2.shape
|
||||
assert np.allclose(ax1, ax2)
|
||||
|
||||
|
||||
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 = [np.arange(i, i + np.prod(shape)).reshape(shape) for i in range(n_arrays)]
|
||||
arrays2 = [backend.from_numpy(array) for array in arrays1]
|
||||
result0 = np.asarray(arrays1)
|
||||
result1 = rearrange(arrays1, "...->...")
|
||||
result2 = rearrange(arrays2, "...->...")
|
||||
assert np.array_equal(result0, result1)
|
||||
assert np.array_equal(result1, backend.to_numpy(result2))
|
||||
|
||||
result1 = rearrange(arrays1, "b ... -> ... b")
|
||||
result2 = rearrange(arrays2, "b ... -> ... b")
|
||||
assert np.array_equal(result1, backend.to_numpy(result2))
|
||||
|
||||
|
||||
def test_gradients_imperatives():
|
||||
# lazy - just checking reductions
|
||||
for reduction in REDUCTIONS:
|
||||
if reduction in ("any", "all"):
|
||||
continue # non-differentiable ops
|
||||
x = np.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 np.allclose(grad1, grad2), [name1, name2, "provided different gradients"]
|
||||
|
||||
|
||||
def test_tiling_imperatives():
|
||||
for backend in imp_op_backends:
|
||||
print("Tiling tests for ", backend.framework_name)
|
||||
input = np.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 = np.tile(input, repeats)
|
||||
converted = backend.from_numpy(input)
|
||||
repeated = backend.tile(converted, repeats)
|
||||
result = backend.to_numpy(repeated)
|
||||
assert np.array_equal(result, expected)
|
||||
|
||||
|
||||
def test_tiling_symbolic():
|
||||
for backend in sym_op_backends:
|
||||
print("Tiling tests for ", backend.framework_name)
|
||||
input = np.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 = np.tile(input, repeats)
|
||||
sym = backend.create_symbol(input.shape)
|
||||
result = backend.eval_symbol(backend.tile(sym, repeats), [[sym, input]])
|
||||
assert np.array_equal(result, expected)
|
||||
|
||||
sym = backend.create_symbol([None] * len(input.shape))
|
||||
result = backend.eval_symbol(backend.tile(sym, repeats), [[sym, input]])
|
||||
assert np.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 np.array_equal(x, reduced_min)
|
||||
assert np.array_equal(x, reduced_max)
|
||||
|
||||
|
||||
def test_repeat_numpy():
|
||||
# check repeat vs reduce. Repeat works ok if reverse reduction with min and max work well
|
||||
x = np.arange(2 * 3 * 5).reshape([2, 3, 5])
|
||||
x1 = repeat(x, "a b c -> copy a b c ", copy=1)
|
||||
assert np.array_equal(x[None], x1)
|
||||
for pattern, axis_dimensions in repeat_test_cases:
|
||||
check_reversion(x, pattern, **axis_dimensions)
|
||||
|
||||
|
||||
def test_repeat_imperatives():
|
||||
x = np.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 np.array_equal(result, expected)
|
||||
|
||||
|
||||
def test_repeat_symbolic():
|
||||
x = np.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 np.array_equal(result, expected)
|
||||
|
||||
|
||||
def test_repeat_array_api():
|
||||
import numpy as xp
|
||||
|
||||
from einops import array_api as AA
|
||||
|
||||
if xp.__version__ < "2.0.0":
|
||||
pytest.skip()
|
||||
|
||||
x = np.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 np.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()),
|
||||
]
|
||||
|
||||
|
||||
def test_anonymous_axes():
|
||||
x = np.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)
|
||||
|
||||
|
||||
def test_list_inputs():
|
||||
x = np.arange(2 * 3 * 4 * 5 * 6).reshape([2, 3, 4, 5, 6])
|
||||
|
||||
assert np.array_equal(
|
||||
rearrange(list(x), "... -> (...)"),
|
||||
rearrange(x, "... -> (...)"),
|
||||
)
|
||||
assert np.array_equal(
|
||||
reduce(list(x), "a ... e -> (...)", "min"),
|
||||
reduce(x, "a ... e -> (...)", "min"),
|
||||
)
|
||||
assert np.array_equal(
|
||||
repeat(list(x), "... -> b (...)", b=3),
|
||||
repeat(x, "... -> b (...)", b=3),
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
func1_compiled_dynamic = torch.compile(func1, dynamic=True, fullgraph=True)
|
||||
|
||||
x = torch.randn(size=[4, 5, 6, 3])
|
||||
assert torch.allclose(func1_compiled_static(x), func1(x), atol=1e-5)
|
||||
assert torch.allclose(func1_compiled_dynamic(x), func1(x), atol=1e-5)
|
||||
# check with input of different dimensionality, and with all shape elements changed
|
||||
x = torch.randn(size=[6, 3, 4, 2, 3])
|
||||
assert torch.allclose(func1_compiled_static(x), func1(x), atol=1e-5)
|
||||
assert torch.allclose(func1_compiled_dynamic(x), func1(x), atol=1e-5)
|
||||
|
||||
|
||||
def bit_count(x):
|
||||
return sum((x >> i) & 1 for i in range(20))
|
||||
|
||||
|
||||
def test_reduction_imperatives_booleans():
|
||||
"""Checks that any/all reduction works in all frameworks"""
|
||||
x_np = np.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 = np.any(x_np, axis=axis, keepdims=True)
|
||||
expected_result_all = np.all(x_np, axis=axis, keepdims=True)
|
||||
assert not np.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 np.array_equal(expected_result_any, backend.to_numpy(res_any))
|
||||
assert np.array_equal(expected_result_all, backend.to_numpy(res_all))
|
||||
|
||||
# expected result: any/all
|
||||
expected_result_any = np.any(x_np, axis=(0, 1), keepdims=True)
|
||||
expected_result_all = np.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 np.array_equal(expected_result_any, backend.to_numpy(res_any))
|
||||
assert np.array_equal(expected_result_all, backend.to_numpy(res_all))
|
||||
@@ -0,0 +1,363 @@
|
||||
import subprocess
|
||||
import tempfile
|
||||
from doctest import testmod
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
import einops
|
||||
import einops.layers
|
||||
from einops._backends import AbstractBackend
|
||||
from einops.einops import _optimize_transformation, parse_shape, rearrange
|
||||
from einops.tests import collect_test_backends, is_backend_tested
|
||||
|
||||
__author__ = "Alex Rogozhnikov"
|
||||
|
||||
rng = np.random.default_rng()
|
||||
|
||||
|
||||
def test_doctests_examples():
|
||||
# tests docstrings, additionally
|
||||
testmod(einops.layers, raise_on_error=True, extraglobs=dict(np=np))
|
||||
testmod(einops.einops, raise_on_error=True, extraglobs=dict(np=np))
|
||||
|
||||
|
||||
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 = set(parse_backends_to_test())
|
||||
errors = []
|
||||
# Find backend subclasses recursively
|
||||
backend_subclasses = []
|
||||
backends = AbstractBackend.__subclasses__()
|
||||
while backends:
|
||||
backend = backends.pop()
|
||||
backends += backend.__subclasses__()
|
||||
backend_subclasses.append(backend)
|
||||
|
||||
for backend_type in backend_subclasses:
|
||||
if backend_type.framework_name not in backends_to_test:
|
||||
continue
|
||||
try:
|
||||
# instantiate
|
||||
backend_type()
|
||||
backends_to_test.remove(backend_type.framework_name)
|
||||
except Exception as e:
|
||||
errors.append((backend_type.framework_name, e))
|
||||
assert len(errors) == 0, errors
|
||||
assert len(backends_to_test) == 0, f"did not instantiate {backends_to_test=}, they won't be tested"
|
||||
|
||||
|
||||
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 = rng.integers(0, 2**12, size=shape).reshape([-1])
|
||||
init_shape = shape[:]
|
||||
n_reduced = rng.integers(0, n_dimensions + 1)
|
||||
reduced_axes = tuple(rng.permutation(n_dimensions)[:n_reduced])
|
||||
axes_reordering = rng.permutation(n_dimensions - n_reduced)
|
||||
final_shape = rng.integers(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 np.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 np.array_equal(result1, result2)
|
||||
|
||||
# testing we can't optimize this formula again
|
||||
combination3 = _optimize_transformation(*combination2)
|
||||
for a, b in zip(combination2, combination3):
|
||||
assert np.array_equal(a, b)
|
||||
|
||||
|
||||
_IMPERATIVE_BACKENDS = collect_test_backends(symbolic=False, layers=False)
|
||||
|
||||
x_np = np.zeros([10, 20, 30, 40])
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
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()
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
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 = np.ones(shape)
|
||||
parsed1 = parse_shape(x, pattern)
|
||||
parsed2 = parse_shape(backend.from_numpy(x), pattern)
|
||||
assert parsed1 == parsed2 == expected
|
||||
|
||||
|
||||
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 = np.ones(shape)
|
||||
parsed1 = parse_shape(x, pattern)
|
||||
parsed2 = parse_shape(backend.from_numpy(x), pattern)
|
||||
assert parsed1 == parsed2 == expected
|
||||
|
||||
|
||||
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 = np.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.parametrize("backend", _SYMBOLIC_BACKENDS)
|
||||
def test_parse_shape_symbolic(backend):
|
||||
for input_shape in [
|
||||
[10, 20, 30, 40],
|
||||
[10, 20, None, None],
|
||||
[None, None, None, None],
|
||||
]:
|
||||
print(f"special shape parsing {backend.framework_name=} {input_shape=}")
|
||||
input_symbol = backend.create_symbol(input_shape)
|
||||
|
||||
shape_placeholder = parse_shape(input_symbol, "a b c d")
|
||||
out_shape = {}
|
||||
for name, symbol in shape_placeholder.items():
|
||||
out_shape[name] = (
|
||||
symbol
|
||||
if isinstance(symbol, int)
|
||||
else backend.eval_symbol(symbol, [(input_symbol, np.zeros([10, 20, 30, 40]))])
|
||||
) # out shape element is either int, or symbol that we are able to eval
|
||||
print(out_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, np.zeros([10, 20, 30, 40]))])
|
||||
print(result.shape)
|
||||
assert result.shape == (10 * 20 * 20, 30, 1 * 2)
|
||||
assert np.allclose(result, 0)
|
||||
|
||||
|
||||
@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, np.zeros(static_shape))])
|
||||
assert out_shape == expected
|
||||
|
||||
|
||||
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 = np.zeros([3, 4, 5], dtype=dtype)
|
||||
input = backend.from_numpy(input)
|
||||
assert backend.is_float_type(input) == is_float, (dtype, backend, input.dtype)
|
||||
|
||||
|
||||
def test_torch_compile_for_functions():
|
||||
"""
|
||||
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
|
||||
|
||||
from einops import einsum, pack, reduce, repeat, unpack
|
||||
from einops._torch_specific import 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)
|
||||
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.double(), suffix).float()
|
||||
|
||||
torch.testing.assert_close(result1, result2, atol=1e-5, rtol=1e-5)
|
||||
|
||||
|
||||
def test_torch_compile_for_layers():
|
||||
"""
|
||||
Einops layers are in general very friendly towards tracing/compiling,
|
||||
but we still want to make sure we can compile them.
|
||||
"""
|
||||
if not is_backend_tested("torch"):
|
||||
pytest.skip()
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from einops.layers.torch import EinMix, Rearrange, Reduce
|
||||
|
||||
original = nn.Sequential(
|
||||
Rearrange("b (t c) -> b t c", c=16),
|
||||
EinMix("b t c -> qkv b t cout", weight_shape="qkv c cout", bias_shape="qkv cout", qkv=3, c=16, cout=8),
|
||||
Reduce("qkv b t cout -> b t qkv", "min", cout=8),
|
||||
)
|
||||
|
||||
compiled = torch.compile(original, fullgraph=True)
|
||||
|
||||
for size in [16, 32, 64]:
|
||||
x = torch.rand([size, size])
|
||||
result1 = original(x)
|
||||
result2 = compiled(x)
|
||||
assert torch.allclose(result1, result2)
|
||||
|
||||
|
||||
src = """
|
||||
import einops
|
||||
import numpy as np
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import torch
|
||||
|
||||
def f():
|
||||
return einops.rearrange(np.ndarray((20, 150, 150)), "... i j -> ... j i")
|
||||
with ThreadPoolExecutor(max_workers=2) as ex:
|
||||
fs = []
|
||||
for i in range(20):
|
||||
fs.append(ex.submit(f))
|
||||
for fut in fs:
|
||||
fut.result()
|
||||
"""
|
||||
|
||||
|
||||
def test_einops_threading():
|
||||
# requires both. Reproduces problem from https://github.com/arogozhnikov/einops/issues/391
|
||||
if not is_backend_tested("torch"):
|
||||
pytest.skip()
|
||||
if not is_backend_tested("numpy"):
|
||||
pytest.skip()
|
||||
|
||||
with tempfile.TemporaryDirectory() as d:
|
||||
testfile = Path(d).joinpath("test.py")
|
||||
testfile.write_text(src)
|
||||
subprocess.run(["python", testfile.absolute().as_posix()], check=True)
|
||||
@@ -0,0 +1,312 @@
|
||||
import dataclasses
|
||||
import typing
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from einops import EinopsError, asnumpy, pack, unpack
|
||||
from einops.tests import collect_test_backends
|
||||
|
||||
rng = np.random.default_rng()
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def test_numpy_trivial(H=13, W=17):
|
||||
def rand(*shape):
|
||||
return rng.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 *"),
|
||||
]
|
||||
|
||||
|
||||
def test_pack_unpack_with_numpy():
|
||||
case: UnpackTestCase
|
||||
|
||||
for case in cases:
|
||||
shape = case.shape
|
||||
pattern = case.pattern
|
||||
|
||||
x = rng.random(shape)
|
||||
# all correct, no minus 1
|
||||
unpack_and_pack(x, [[2], [1], [2]], pattern)
|
||||
# no -1, asking for wrong shapes
|
||||
with pytest.raises(EinopsError):
|
||||
unpack_and_pack(x, [[2], [1], [2]], pattern + " non_existent_axis")
|
||||
with pytest.raises(EinopsError):
|
||||
unpack_and_pack(x, [[2], [1], [1]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
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(EinopsError):
|
||||
unpack(x, [[2], [4], [-1]], pattern)
|
||||
# this one does not raise, because indexing x[2:1] just returns zero elements
|
||||
# with pytest.raises(EinopsError):
|
||||
# unpack(x, [[2], [-1], [4]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
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(EinopsError):
|
||||
unpack(x, [[-1, 2], [1], [5]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
unpack(x, [[2, 2], [2], [5, -1]], pattern)
|
||||
|
||||
# asking for non-divisible number of elements
|
||||
with pytest.raises(EinopsError):
|
||||
unpack(x, [[2, 1], [1], [3, -1]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
unpack(x, [[2, 1], [3, -1], [1]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
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)
|
||||
|
||||
|
||||
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 = rng.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(EinopsError):
|
||||
unpack(x, [[2], [1], [1]], pattern)
|
||||
|
||||
with pytest.raises(EinopsError):
|
||||
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(EinopsError):
|
||||
unpack(x, [[2], [4], [-1]], pattern)
|
||||
# this one does not raise, because indexing x[2:1] just returns zero elements
|
||||
# with pytest.raises(EinopsError):
|
||||
# unpack(x, [[2], [-1], [4]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
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(EinopsError):
|
||||
unpack(x, [[-1, 2], [1], [5]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
unpack(x, [[2, 2], [2], [5, -1]], pattern)
|
||||
|
||||
# asking for non-divisible number of elements
|
||||
with pytest.raises(EinopsError):
|
||||
unpack(x, [[2, 1], [1], [3, -1]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
unpack(x, [[2, 1], [3, -1], [1]], pattern)
|
||||
with pytest.raises(EinopsError):
|
||||
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)
|
||||
|
||||
|
||||
def test_pack_unpack_array_api():
|
||||
import numpy as xp
|
||||
|
||||
from einops import array_api as AA
|
||||
|
||||
if xp.__version__ < "2.0.0":
|
||||
pytest.skip()
|
||||
|
||||
for case in cases:
|
||||
shape = case.shape
|
||||
pattern = case.pattern
|
||||
x_np = rng.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(EinopsError):
|
||||
unpack(x_np, ps, pattern)
|
||||
@@ -0,0 +1,126 @@
|
||||
import pytest
|
||||
|
||||
from einops import EinopsError
|
||||
from einops.parsing import AnonymousAxis, ParsedExpression, _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
|
||||
|
||||
|
||||
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]
|
||||
|
||||
|
||||
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)
|
||||
|
||||
|
||||
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...")
|
||||
|
||||
|
||||
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
|
||||
Reference in New Issue
Block a user