import importlib.util import pytest import autoray as ar from autoray import shape import numpy as np # find backends to tests BACKENDS = [pytest.param("numpy")] for lib in [ "cupy", "dask", "tensorflow", "torch", "mars", "jax", "sparse", "paddle", ]: if importlib.util.find_spec(lib): BACKENDS.append(pytest.param(lib)) if lib == "jax": import os import jax jax.config.update("jax_enable_x64", True) jax.config.update("jax_platform_name", "cpu") os.environ["XLA_PYTHON_CLIENT_ALLOCATOR"] = "platform" else: BACKENDS.append( pytest.param( lib, marks=pytest.mark.skipif(True, reason=f"No {lib}.") ) ) JAX_RANDOM_KEY = None def gen_rand(shape, backend, dtype="float64"): if "complex" in dtype: re = gen_rand(shape, backend) im = gen_rand(shape, backend) return ar.astype(ar.do("complex", re, im), dtype) if backend == "jax": from jax import random as jrandom global JAX_RANDOM_KEY if JAX_RANDOM_KEY is None: JAX_RANDOM_KEY = jrandom.PRNGKey(42) JAX_RANDOM_KEY, subkey = jrandom.split(JAX_RANDOM_KEY) return jrandom.uniform(subkey, shape=shape, dtype=dtype) elif backend == "sparse": x = ar.do( "random.uniform", size=shape, like=backend, density=0.5, format="coo", fill_value=0, ) else: x = ar.do("random.uniform", size=shape, like=backend) x = ar.astype(x, ar.to_backend_dtype(dtype, backend)) assert ar.get_dtype_name(x) == dtype return x @pytest.mark.parametrize( "f,args,xfail_backends", ( ("all", (), ()), ("clip", (0.2, 0.7), ()), ("conj", (), ()), ("cos", (), ()), ("cosh", (), ()), ("count_nonzero", (), ()), ("cumsum", (0,), ("sparse",)), ("diag", (), ("sparse",)), ("diag", (1,), ("sparse",)), ("diag", (-1,), ("sparse",)), ("exp", (), ()), ("imag", (), ()), ("log", (), ()), ("log10", (), ()), ("max", (-1,), ("sparse",)), ("max", (), ()), ("mean", (), ()), ("mean", (0,), ("sparse",)), ("min", (), ()), ("power", (2,), ("sparse",)), ("prod", (), ()), ("ravel", (), ("sparse",)), ("real", (), ()), ("sin", (), ()), ("sinh", (), ()), ("sqrt", (), ()), ("sum", (), ()), ("sum", (1,), ()), ("tan", (), ()), ("tanh", (), ()), ("trace", (), ("sparse",)), ("tril", (), ()), ("tril", (), (1,)), ("tril", (), (-1,)), ("triu", (), ()), ("triu", (), (1,)), ("triu", (), (-1,)), ), ) @pytest.mark.parametrize("backend", BACKENDS) def test_unary_functions(f, args, xfail_backends, backend): if backend in xfail_backends: pytest.xfail(f"{backend} doesn't support {f}.") xn = ar.do("random.uniform", size=(4, 5), like="numpy") yn = ar.do(f, xn, *args) x = ar.do("asarray", xn, like=backend) y = ar.do(f, x, *args) yt = ar.do("to_numpy", y) assert ar.do("allclose", yt, yn) @pytest.mark.parametrize( "f,args,xfail_backends", ( ("add", (), ()), ("allclose", (), ("sparse",)), ("divide", (), ()), ("matmul", (), ()), ("multiply", (), ()), ("subtract", (), ()), ), ) @pytest.mark.parametrize("backend", BACKENDS) def test_binary_functions(f, args, xfail_backends, backend): if backend in xfail_backends: pytest.xfail(f"{backend} doesn't support {f}.") xan = ar.do("random.uniform", size=(3, 3), like="numpy") xbn = ar.do("random.uniform", size=(3, 3), like="numpy") yn = ar.do(f, xan, xbn, *args) xa = ar.do("asarray", xan, like=backend) xb = ar.do("asarray", xbn, like=backend) y = ar.do(f, xa, xb, *args) yt = ar.do("to_numpy", y) assert ar.do("allclose", yt, yn) @pytest.mark.parametrize( "f", [ "sum", "prod", "max", "min", "mean", ], ) @pytest.mark.parametrize( "kwargs", [ {}, {"axis": 1}, {"axis": 1, "keepdims": True}, {"axis": (0, 2)}, ], ) @pytest.mark.parametrize("backend", BACKENDS) def test_reduce_functions(f, kwargs, backend): if ( backend == "torch" and f == "prod" and isinstance(kwargs.get("axis"), tuple) ): pytest.xfail("Pytorch doesn't support prod with tuple axis.") x = ar.do("random.normal", size=(2, 3, 4), like="numpy") y = ar.do(f, x, **kwargs) xb = ar.do("asarray", x, like=backend) yb = ar.do(f, xb, **kwargs) yt = ar.do("to_numpy", yb) assert ar.do("allclose", yt, y) @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("fn", ["sqrt", "exp", "sum"]) def test_basic(backend, fn): if (backend == "ctf") and fn in ("sqrt", "sum"): pytest.xfail("ctf doesn't have sqrt, and converts sum output to numpy") x = gen_rand((2, 3, 4), backend) y = ar.do(fn, x) if (backend == "sparse") and (fn == "sum"): pytest.xfail("Sparse 'sum' outputs dense.") assert ar.infer_backend(x) == ar.infer_backend(y) == backend def test_infer_backend_multi(): x = 1.0 y = gen_rand((2, 3), "numpy") z = ar.lazy.Variable((4, 5)) assert ar.infer_backend_multi(x) == "builtins" assert ar.infer_backend_multi(x, y) == "numpy" assert ar.infer_backend_multi(x, y, z) == "autoray.lazy" def test_raises_import_error_when_missing(): with pytest.raises(ImportError): ar.do("anonexistantfunction", 1, like="numpy") with pytest.raises(ImportError): ar.do("ones", 1, like="anonexistantbackend") @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize( "fn,args", [ (ar.conj, []), (ar.transpose, []), (ar.real, []), (ar.imag, []), (ar.reshape, [(5, 3)]), ], ) def test_attribute_prefs(backend, fn, args): if (backend == "torch") and fn in (ar.real, ar.imag): pytest.xfail("Pytorch doesn't support complex numbers yet...") x = gen_rand((3, 5), backend) y = fn(x, *args) assert ar.infer_backend(x) == ar.infer_backend(y) == backend def modified_gram_schmidt(X): Q = [] for j in range(0, shape(X)[0]): q = X[j, :] for i in range(0, j): rij = ar.do("tensordot", ar.do("conj", Q[i]), q, 1) q = q - rij * Q[i] rjj = ar.do("linalg.norm", q, 2) Q.append(q / rjj) return ar.do("stack", Q, axis=0) @pytest.mark.parametrize("backend", BACKENDS) def test_mgs(backend): if backend == "sparse": pytest.xfail("Sparse doesn't support linear algebra yet...") if backend == "ctf": pytest.xfail("ctf does not have 'stack' function.") x = gen_rand((3, 5), backend) Ux = modified_gram_schmidt(x) y = ar.do("sum", Ux @ ar.dag(Ux)) assert ar.to_numpy(y) == pytest.approx(3) def modified_gram_schmidt_np_mimic(X): from autoray import numpy as np print(np) Q = [] for j in range(0, shape(X)[0]): q = X[j, :] for i in range(0, j): rij = np.tensordot(np.conj(Q[i]), q, 1) q = q - rij * Q[i] rjj = np.linalg.norm(q, 2) Q.append(q / rjj) return np.stack(Q, axis=0) def test_numpy_mimic_dunder_methods(): from abc import ABC from autoray import numpy as np class Base(ABC): pass assert isinstance(np, object) assert not isinstance(np, Base) print(np) dir(np) @pytest.mark.parametrize("backend", BACKENDS) def test_mgs_np_mimic(backend): if backend == "sparse": pytest.xfail("Sparse doesn't support linear algebra yet...") if backend == "ctf": pytest.xfail("ctf does not have 'stack' function.") x = gen_rand((3, 5), backend) Ux = modified_gram_schmidt_np_mimic(x) y = ar.do("sum", Ux @ ar.dag(Ux)) assert ar.to_numpy(y) == pytest.approx(3) @pytest.mark.parametrize("backend", BACKENDS) def test_linalg_svd_square(backend): if backend == "sparse": pytest.xfail("Sparse doesn't support linear algebra yet...") x = gen_rand((5, 4), backend) U, s, V = ar.do("linalg.svd", x) assert ( ar.infer_backend(x) == ar.infer_backend(U) == ar.infer_backend(s) == ar.infer_backend(V) == backend ) y = U @ ar.do("diag", s, like=x) @ V diff = ar.do("sum", ar.do("abs", y - x)) assert ar.to_numpy(diff) < 1e-8 @pytest.mark.parametrize("backend", BACKENDS) def test_translator_random_uniform(backend): from autoray import numpy as anp if backend == "sparse": pytest.xfail("Sparse will have zeros") x = anp.random.uniform(low=-10, size=(4, 5), like=backend) assert (ar.to_numpy(x) > -10).all() assert (ar.to_numpy(x) < 1.0).all() # test default single scalar x = anp.random.uniform(low=1000, high=2000, like=backend) assert 1000 <= ar.to_numpy(x) < 2000 @pytest.mark.parametrize("backend", BACKENDS) def test_translator_random_normal(backend): if backend == "ctf": pytest.xfail() from autoray import numpy as anp x = anp.random.normal(100.0, 0.1, size=(4, 5), like=backend) if backend == "sparse": assert (x.data > 90.0).all() assert (x.data < 110.0).all() return assert (ar.to_numpy(x) > 90.0).all() assert (ar.to_numpy(x) < 110.0).all() if backend == "tensorflow": x32 = ar.do( "random.normal", 100.0, 0.1, dtype="float32", size=(4, 5), like=backend, ) assert x32.dtype == "float32" assert (ar.to_numpy(x32) > 90.0).all() assert (ar.to_numpy(x32) < 110.0).all() # test default single scalar x = anp.random.normal(loc=1500, scale=10, like=backend) assert 1000 <= ar.to_numpy(x) < 2000 @pytest.mark.parametrize("backend", BACKENDS) def test_tril(backend): x = gen_rand((4, 4), backend) xl = ar.do("tril", x) xln = ar.to_numpy(xl) assert xln[0, 1] == 0.0 if backend != "sparse": # this won't work for sparse because density < 1 assert (xln > 0.0).sum() == 10 xl = ar.do("tril", x, k=1) xln = ar.to_numpy(xl) if backend != "sparse": # this won't work for sparse because density < 1 assert xln[0, 1] != 0.0 assert xln[0, 2] == 0.0 if backend != "sparse": # this won't work for sparse because density < 1 assert (xln > 0.0).sum() == 13 @pytest.mark.parametrize("backend", BACKENDS) def test_triu(backend): x = gen_rand((4, 4), backend) xl = ar.do("triu", x) xln = ar.to_numpy(xl) assert xln[1, 0] == 0.0 if backend != "sparse": # this won't work for sparse because density < 1 assert (xln > 0.0).sum() == 10 xl = ar.do("triu", x, k=-1) xln = ar.to_numpy(xl) if backend != "sparse": # this won't work for sparse because density < 1 assert xln[1, 0] != 0.0 assert xln[2, 0] == 0.0 if backend != "sparse": # this won't work for sparse because density < 1 assert (xln > 0.0).sum() == 13 @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("shape", [(4, 3), (4, 4), (3, 4)]) def test_qr_thin_square_fat(backend, shape): if backend == "sparse": pytest.xfail("Sparse doesn't support linear algebra yet...") x = gen_rand(shape, backend) Q, R = ar.do("linalg.qr", x) xn, Qn, Rn = map(ar.to_numpy, (x, Q, R)) assert ar.do("allclose", xn, Qn @ Rn) @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("array_dtype", ["int", "float", "bool"]) def test_count_nonzero(backend, array_dtype): if backend == "mars": import mars if tuple(map(int, mars.__version__.split("."))) < (0, 4, 0): pytest.xfail("mars count_nonzero bug fixed in version 0.4.") if backend == "ctf" and array_dtype == "bool": pytest.xfail("ctf doesn't support bool array dtype") if array_dtype == "int": x = ar.do("asarray", [0, 1, 2, 0, 3], like=backend) elif array_dtype == "float": x = ar.do("asarray", [0.0, 1.0, 2.0, 0.0, 3.0], like=backend) elif array_dtype == "bool": x = ar.do("asarray", [False, True, True, False, True], like=backend) nz = ar.do("count_nonzero", x) assert ar.to_numpy(nz) == 3 def test_pseudo_submodules(): x = gen_rand((2, 3), "numpy") xT = ar.do("numpy.transpose", x, like="autoray") assert shape(xT) == (3, 2) @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("creation", ["ones", "zeros"]) @pytest.mark.parametrize( "dtype", ["float32", "float64", "complex64", "complex128"] ) def test_dtype_specials(backend, creation, dtype): import numpy as np x = ar.do(creation, shape=(2, 3), like=backend) if backend == "torch" and "complex" in dtype: pytest.xfail("Pytorch doesn't support complex numbers yet...") x = ar.astype(x, dtype) assert ar.get_dtype_name(x) == dtype x = ar.to_numpy(x) assert isinstance(x, np.ndarray) assert ar.get_dtype_name(x) == dtype @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("real_dtype", ["float32", "float64"]) def test_complex_creation(backend, real_dtype): if backend == "torch": pytest.xfail("Pytorch doesn't support complex numbers yet...") if (backend == "sparse") and (real_dtype == "float32"): pytest.xfail( "Bug in sparse where single precision isn't maintained " "after scalar multiplication." ) if (backend == "ctf") and (real_dtype == "float32"): pytest.xfail( "ctf currently doesn't preserve single precision when " "multiplying by python scalars." ) x = ar.do( "complex", ar.astype( ar.do("random.uniform", size=(3, 4), like=backend), real_dtype ), ar.astype( ar.do("random.uniform", size=(3, 4), like=backend), real_dtype ), ) assert ( ar.get_dtype_name(x) == {"float32": "complex64", "float64": "complex128"}[real_dtype] ) @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize( "dtype_in,dtype_out", [ ("float32", "float32"), ("float64", "float64"), ("complex64", "float32"), ("complex128", "float64"), ], ) def test_real_imag(backend, dtype_in, dtype_out): x = gen_rand((3, 4), backend, dtype_in) re = ar.do("real", x) im = ar.do("imag", x) assert ar.infer_backend(re) == backend assert ar.infer_backend(im) == backend assert ar.get_dtype_name(re) == dtype_out assert ar.get_dtype_name(im) == dtype_out assert ar.do("allclose", ar.to_numpy(x).real, ar.to_numpy(re)) assert ar.do("allclose", ar.to_numpy(x).imag, ar.to_numpy(im)) @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize( "dtype", ["float32", "float64", "complex64", "complex128"], ) def test_linalg_solve(backend, dtype): if backend == "sparse": pytest.xfail("Sparse doesn't support linear algebra yet...") if backend == "paddle" and "complex" in dtype: pytest.xfail( "Paddle `linalg.solve` doesn't support complex numbers yet..." ) A = gen_rand((4, 4), backend, dtype) b = gen_rand((4, 1), backend, dtype) x = ar.do("linalg.solve", A, b) assert ar.do( "allclose", ar.to_numpy(A @ x), ar.to_numpy(b), rtol=1e-3, atol=1e-6 ) @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize( "dtype", ["float32", "float64", "complex64", "complex128"], ) def test_linalg_eigh(backend, dtype): if backend == "sparse": pytest.xfail("sparse doesn't support linalg.eigh yet.") if backend == "dask": pytest.xfail("dask doesn't support linalg.eigh yet.") if backend == "mars": pytest.xfail("mars doesn't support linalg.eigh yet.") if (backend == "torch") and ("complex" in dtype): pytest.xfail("Pytorch doesn't fully support complex yet.") A = gen_rand((4, 4), backend, dtype) A = A + ar.dag(A) el, ev = ar.do("linalg.eigh", A) B = (ev * ar.reshape(el, (1, -1))) @ ar.dag(ev) assert ar.do("allclose", ar.to_numpy(A), ar.to_numpy(B), rtol=1e-3) @pytest.mark.parametrize("backend", BACKENDS) def test_pad(backend): if backend == "sparse": pytest.xfail("sparse doesn't support linalg.eigh yet.") if backend == "mars": pytest.xfail("mars doesn't support linalg.eigh yet.") A = gen_rand((3, 4, 5), backend) for pad_width, new_shape in [ # same pad before and after for every axis (2, (7, 8, 9)), # same pad for every axis (((1, 2),), (6, 7, 8)), # different pad for every axis (((4, 3), (2, 4), (3, 2)), (10, 10, 10)), ]: B = ar.do("pad", A, pad_width) assert shape(B) == new_shape assert ar.to_numpy(ar.do("sum", A)) == pytest.approx( ar.to_numpy(ar.do("sum", B)) ) @pytest.mark.parametrize("backend", BACKENDS) def test_register_function(backend): x = ar.do("ones", shape=(2, 3), like=backend) def direct_fn(x): return 1 # first test we can provide the function directly ar.register_function(backend, "test_register", direct_fn) assert ar.do("test_register", x) == 1 def wrap_fn(fn): def new_fn(*args, **kwargs): res = fn(*args, **kwargs) return res + 1 return new_fn # then check we can wrap the old (previous) function ar.register_function(backend, "test_register", wrap_fn, wrap=True) assert ar.do("test_register", x) == 2 @pytest.mark.parametrize("backend", BACKENDS) def test_take(backend): if backend in {"sparse", "paddle"}: pytest.xfail(f"{backend} doesn't fully support take yet") num_inds = 4 A = gen_rand((2, 3, 4), backend) if backend == "jax": # gen_rand doesn't work with ints for JAX ind = gen_rand((num_inds,), "numpy", dtype="int64") else: ind = gen_rand((num_inds,), backend, dtype="int64") # Ensure indices are of dtype 'int32' ind = ind.astype('int32') # Take along axis 1, and check if result makes sense B = ar.do("take", A, ind, axis=1) assert shape(B) == (2, 4, 4) for i in range(num_inds): assert ar.do( "allclose", ar.to_numpy(A[:, ind[0], :]), ar.to_numpy(B[:, 0, :]) ) assert ar.infer_backend(A) == ar.infer_backend(B) @pytest.mark.parametrize("backend", BACKENDS) def test_concatenate(backend): mats = [gen_rand((2, 3, 4), backend) for _ in range(3)] # Concatenate along axis 1, check if shape is correct # also check if automatically inferring backend works mats_concat1 = ar.do("concatenate", mats, axis=1) mats_concat2 = ar.do("concatenate", mats, axis=1, like=backend) assert shape(mats_concat1) == shape(mats_concat2) == (2, 9, 4) assert ( backend == ar.infer_backend(mats_concat1) == ar.infer_backend(mats_concat2) ) @pytest.mark.parametrize("backend", BACKENDS) def test_stack(backend): mats = [gen_rand((2, 3, 4), backend) for _ in range(3)] # stack, creating a new axis (at position 0) # also check if automatically inferring backend works mats_stack1 = ar.do("stack", mats) mats_stack2 = ar.do("stack", mats, like=backend) assert shape(mats_stack1) == shape(mats_stack2) == (3, 2, 3, 4) assert ( backend == ar.infer_backend(mats_stack1) == ar.infer_backend(mats_stack2) ) @pytest.mark.parametrize("backend", BACKENDS) def test_einsum(backend): A = gen_rand((2, 3, 4), backend) B = gen_rand((3, 4, 2), backend) C1 = ar.do("einsum", "ijk,jkl->il", A, B, like=backend) C2 = ar.do("einsum", "ijk,jkl->il", A, B) if backend not in ("torch", "tensorflow", "paddle"): # interleaved syntax is not supported C3 = ar.do("einsum", A, [0, 1, 2], B, [1, 2, 3], [0, 3]) else: C3 = C1 C4 = ar.do("reshape", A, (2, 12)) @ ar.do("reshape", B, (12, 2)) assert shape(C1) == shape(C2) == shape(C3) == (2, 2) assert ar.do("allclose", ar.to_numpy(C1), ar.to_numpy(C4)) assert ar.do("allclose", ar.to_numpy(C2), ar.to_numpy(C4)) assert ar.do("allclose", ar.to_numpy(C3), ar.to_numpy(C4)) assert ( ar.infer_backend(C1) == ar.infer_backend(C2) == ar.infer_backend(C3) == ar.infer_backend(C4) == backend ) @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("int_or_section", ["int", "section", "empty"]) def test_split(backend, int_or_section): if backend == "sparse": pytest.xfail("sparse doesn't support split yet") if backend == "dask": pytest.xfail("dask doesn't support split yet") A = ar.do("ones", (10, 20, 10), like=backend) if int_or_section == "section": sections = [2, 4, 14] splits = ar.do("split", A, sections, axis=1) assert len(splits) == 4 assert ar.shape(splits[3]) == (10, 6, 10) elif int_or_section == "empty": splits = ar.do("split", A, [], axis=1) assert len(splits) == 1 assert ar.shape(splits[0]) == (10, 20, 10) else: splits = ar.do("split", A, 5, axis=2) assert len(splits) == 5 assert ar.shape(splits[2]) == (10, 20, 2) @pytest.mark.parametrize("backend", BACKENDS) def test_where(backend): if backend in {"sparse", "paddle"}: pytest.xfail(f"{backend} doesn't fully support `where` yet") A = ar.do("arange", 10, like=backend) B = ar.do("arange", 10, like=backend) + 1 C = ar.do("stack", [A, B]) D = ar.do("where", C < 5) if backend == "dask": for x in D: x.compute_chunk_sizes() for x in D: assert ar.to_numpy(x).shape == (9,) @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize("dtype_str", ["float32", "float64"]) @pytest.mark.parametrize( "fn", ["random.normal", "random.uniform", "zeros", "ones", "eye"] ) @pytest.mark.parametrize("str_or_backend", ("str", "backend")) def test_dtype_kwarg(backend, dtype_str, fn, str_or_backend): if str_or_backend == "str": dtype = dtype_str else: dtype = ar.to_backend_dtype(dtype_str, like=backend) if fn in ("random.normal", "random.uniform"): A = ar.do(fn, size=(10, 5), dtype=dtype, like=backend) elif fn in ("zeros", "ones"): A = ar.do(fn, shape=(10, 5), dtype=dtype, like=backend) else: # fn = 'eye' A = ar.do(fn, 10, dtype=dtype, like=backend) assert shape(A) == (10, 10) A = ar.do(fn, 10, 5, dtype=dtype, like=backend) assert shape(A) == (10, 5) assert ar.get_dtype_name(A) == dtype_str @pytest.mark.parametrize("backend", BACKENDS) def test_get_common_dtype(backend): x = ar.do("ones", (1,), like=backend, dtype="complex64") y = ar.do("ones", (1,), like=backend, dtype="float64") assert ar.get_common_dtype(x, y) == "complex128" @pytest.mark.parametrize("backend", BACKENDS) def test_backend_like(backend): assert ar.get_backend() is None ar.set_backend("test") assert ar.get_backend() == "test" ar.set_backend(None) assert ar.get_backend() is None with ar.backend_like(backend): assert ar.get_backend() == backend x = ar.do("ones", (2,), like=backend) assert ar.infer_backend(x) == backend assert ar.get_backend() is None def test_nested_multihreaded_backend_like(): from autoray.autoray import choose_backend from concurrent.futures import ThreadPoolExecutor def foo(backend1, backend2): bs = [] bs.append( ( ar.get_backend(), choose_backend("test", 1), ) ) with ar.backend_like(backend1): bs.append( ( ar.get_backend(), choose_backend("test", 1), ) ) with ar.backend_like(backend2): bs.append( ( ar.get_backend(), choose_backend("test", 1), ) ) bs.append( ( ar.get_backend(), choose_backend("test", 1), ) ) bs.append((ar.get_backend(), choose_backend("test", 1))) return bs b_exp = [("A", "A"), ("B", "B"), ("C", "C"), ("B", "B"), ("A", "A")] with ar.backend_like("A"): b = foo("B", "C") assert b == b_exp b_exp = [ ("A", "A"), ("B", "B"), (None, "builtins"), ("B", "B"), ("A", "A"), ] with ar.backend_like("A"): b = foo("B", None) assert b == b_exp with ThreadPoolExecutor(3) as pool: b_exp = [(None, "A"), ("B", "B"), ("C", "C"), ("B", "B"), (None, "A")] with ar.backend_like("A"): bs = [pool.submit(foo, "B", "C") for _ in range(3)] for b in bs: assert b.result() == b_exp b_exp = [(None, "A"), ("B", "B"), (None, "A"), ("B", "B"), (None, "A")] with ar.backend_like("A"): bs = [pool.submit(foo, "B", None) for _ in range(3)] for b in bs: assert b.result() == b_exp def test_compose(): @ar.compose def mycomposedfn(x, backend): x = ar.do("exp", x, like=backend) x = ar.do("log", x, like=backend) return x x = ar.do("ones", (2,), like="numpy") y = ar.do("mycomposedfn", x) assert ar.do("allclose", x, y) y = mycomposedfn(x) assert ar.do("allclose", x, y) mycomposedfn.register("numpy", lambda x: 1) y = ar.do("mycomposedfn", x) assert y == 1 y = mycomposedfn(x) assert y == 1 @mycomposedfn.register("numpy") def f(x): return 2 y = ar.do("mycomposedfn", x) assert y == 2 y = mycomposedfn(x) assert y == 2 def test_builtins_complex(): re = 1.0 im = 2.0 z = ar.do("complex", re, im) assert z == 1.0 + 2.0j assert ar.infer_backend(z) == "builtins" def test_shape_ndim_builtins(): import numpy as np xs = [ 1, 4.0, 7j, (), [], [[]], [np.ones(3), np.ones(3)], np.ones((5, 4, 3)), ] for x in xs: assert ar.shape(x) == np.shape(x) assert ar.ndim(x) == np.ndim(x) @pytest.mark.parametrize("backend", BACKENDS) def test_scipy_dispatching(backend): if backend not in ["numpy", "cupy", "jax"]: pytest.xfail("backend doesn't suport scipy.") x = gen_rand((3, 3), backend=backend) ar.do("scipy.linalg.expm", x) def check_array_dtypes(x, y): assert x.dtype == y.dtype if hasattr(x, "device"): assert x.device == y.device @pytest.mark.parametrize("backend", BACKENDS) @pytest.mark.parametrize( "dtype", ["float32", "float64", "complex64", "complex128"] ) class TestCreationRoutines: def test_empty_passes_dtype_device(self, backend, dtype): if backend in ("tensorflow",): pytest.xfail(f"{backend} doesn't support empty yet.") x = gen_rand((1,), backend, dtype) y = ar.do("empty", (2, 3), like=x) check_array_dtypes(x, y) def test_eye_passes_dtype_device(self, backend, dtype): if backend == "paddle" and "complex" in dtype: pytest.xfail("Paddle doesn't support complex eye yet.") x = gen_rand((1,), backend, dtype) y = ar.do("eye", 3, like=x) check_array_dtypes(x, y) def test_full_passes_dtype_device(self, backend, dtype): if backend in ("tensorflow",): pytest.xfail(f"{backend} doesn't support full yet.") x = gen_rand((1,), backend, dtype) y = ar.do("full", (2, 3), 7, like=x) check_array_dtypes(x, y) def test_identity_passes_dtype_device(self, backend, dtype): if backend == "paddle" and "complex" in dtype: pytest.xfail("Paddle doesn't support complex identity yet.") x = gen_rand((1,), backend, dtype) y = ar.do("identity", 4, like=x) check_array_dtypes(x, y) def test_ones_passes_dtype_device(self, backend, dtype): x = gen_rand((1,), backend, dtype) y = ar.do("ones", (2, 3), like=x) check_array_dtypes(x, y) def test_zeros_passes_dtype_device(self, backend, dtype): x = gen_rand((1,), backend, dtype) y = ar.do("zeros", (2, 3), like=x) check_array_dtypes(x, y) # def test_arange_passes_dtype_device(self, backend, dtype): # if backend in ("sparse",): # pytest.xfail("Sparse doesn't support arange yet.") # if backend == "torch" and "complex" in dtype: # pytest.xfail("torch.arange doesn't support complex numbers yet.") # if backend == "tensorflow" and "complex" in dtype: # pytest.xfail("torch.arange doesn't support complex numbers yet.") # x = gen_rand((1,), backend, dtype) # y = ar.do("arange", 1, 10, like=x) # check_array_dtypes(x, y) # def test_linspace_passes_dtype_device(self, backend, dtype): # if backend in ("sparse", "tensorflow"): # pytest.xfail(f"{backend} doesn't support linspace yet.") # x = gen_rand((1,), backend, dtype) # y = ar.do("linspace", 10, 20, 11, like=x) # check_array_dtypes(x, y) # def test_logspace_passes_dtype_device(self, backend, dtype): # if backend in ("sparse", "tensorflow"): # pytest.xfail(f"{backend} doesn't support logspace yet.") # x = gen_rand((1,), backend, dtype) # if backend not in {"dask"}: # y = ar.do("logspace", 10, 20, 11, like=x) # check_array_dtypes(x, y) # def test_geomspace_passes_dtype_device(self, backend, dtype): # if backend in ("sparse", "tensorflow"): # pytest.xfail(f"{backend} doesn't support logspace yet.") # x = gen_rand((1,), backend, dtype) # if backend not in {"dask"}: # y = ar.do("logspace", 10, 20, 11, like=x) # check_array_dtypes(x, y) creation_funcs_with_args = [ ("empty", ((2, 3),)), ("eye", (4,)), ("full", ((2, 3), 7)), ("identity", (4,)), ("ones", ((2, 3),)), ("zeros", ((2, 3),)), ] creation_builtins = [ (float, [np.float64]), (int, [np.int32, np.int64]), # np.int32 on Windows and np.int64 else (complex, [np.complex128]), ] @pytest.mark.parametrize("fn, args", creation_funcs_with_args) @pytest.mark.parametrize("dtype, expected", creation_builtins) def test_creation_with_builtins(fn, args, dtype, expected): x = dtype(4) y = ar.do(fn, *args, like=x) assert y.dtype in expected @pytest.mark.parametrize("backend", BACKENDS) def test_indices(backend): if backend == "sparse": pytest.xfail("Sparse doesn't support `indices` function yet.") from numpy.testing import assert_array_equal x = ar.do("indices", (3, 4), like=backend) xn = ar.to_numpy(x) xe = ar.do("indices", (3, 4), like="numpy") assert_array_equal(xn, xe)