def test_nestedbackend():
obj = object()
be_outer = Backend()
be_outer.__ua_function__ = lambda f, a, kw: obj
mm1 = ua.generate_multimethod(lambda: (), lambda a, kw, d: (a, kw),
"ua_tests")
def default(*a, **kw):
return mm1(*a, **kw)
mm2 = ua.generate_multimethod(lambda: (),
lambda a, kw, d: (a, kw),
"ua_tests",
default=default)
be_inner = Backend()
def be2_ua_func(f, a, kw):
with ua.skip_backend(be_inner):
return f(*a, **kw)
be_inner.__ua_function__ = be2_ua_func
with ua.set_backend(be_outer), ua.set_backend(be_inner):
assert mm2() is obj
def __ua_function__(self, func, args, kwargs):
extracted_args = func.arg_extractor(*args, **kwargs)
arr_args = tuple(x.value for x in extracted_args
if x.type is np.ndarray)
with ua.set_backend(self._inner, only=True):
if len(arr_args) == 0:
out = func(*args, **kwargs)
else:
a, kw = self.replace_arrays(
func,
args,
kwargs,
(x.value
if x is not None and isinstance(x, DiffArray) else x
for x in arr_args),
)
out = func(*a, **kw)
real_func = func
if func is np.ufunc.__call__:
real_func = args[0]
if real_func not in raw_functions:
with ua.set_backend(self._inner, coerce=True):
if self._mode == "vjp":
out = VJPDiffArray(out)
else:
out = JVPDiffArray(out)
if real_func not in nograd_functions:
out.register_diff(func, args, kwargs)
return out
def test_class_overriding():
with ua.set_backend(NumpyBackend, coerce=True):
assert isinstance(onp.add, np.ufunc)
assert isinstance(onp.dtype("float64"), np.dtype)
assert np.dtype("float64") == onp.float64
assert isinstance(np.dtype("float64"), onp.dtype)
assert issubclass(onp.ufunc, np.ufunc)
with ua.set_backend(DaskBackend(), coerce=True):
assert isinstance(da.add, np.ufunc)
assert isinstance(onp.dtype("float64"), np.dtype)
assert np.dtype("float64") == onp.float64
assert isinstance(np.dtype("float64"), onp.dtype)
assert issubclass(da.ufunc.ufunc, np.ufunc)
with ua.set_backend(SparseBackend, coerce=True):
assert isinstance(onp.add, np.ufunc)
assert isinstance(onp.dtype("float64"), np.dtype)
assert np.dtype("float64") == onp.float64
assert isinstance(np.dtype("float64"), onp.dtype)
assert issubclass(onp.ufunc, np.ufunc)
if hasattr(CupyBackend, "__ua_function__"):
with ua.set_backend(CupyBackend, coerce=True):
assert isinstance(cp.add, np.ufunc)
assert isinstance(cp.dtype("float64"), np.dtype)
assert np.dtype("float64") == cp.float64
assert isinstance(np.dtype("float64"), cp.dtype)
assert issubclass(cp.ufunc, np.ufunc)
def test_raising_from_backend(nullary_mm):
def raise_(foo):
raise foo
Foo = ua.BackendNotImplementedError("Foo")
be = Backend()
be.__ua_function__ = lambda f, a, kw: raise_(Foo)
# BackendNotImplementedErrors are nested
with ua.set_backend(be):
with pytest.raises(ua.BackendNotImplementedError) as e:
nullary_mm()
assert (
e.value.args[0] ==
"No selected backends had an implementation for this function.")
assert type(e.value.args[1]) == tuple
assert e.value.args[1] == (be, Foo)
Bar = ua.BackendNotImplementedError("Bar")
be2 = Backend()
be2.__ua_function__ = lambda f, a, kw: raise_(Bar)
# Errors are in the order the backends were tried
with ua.set_backend(be), ua.set_backend(be2):
with pytest.raises(ua.BackendNotImplementedError) as e:
nullary_mm()
assert e.value.args[1] == (be2, Bar)
assert e.value.args[2] == (be, Foo)
be3 = Backend()
be3.__ua_function__ = lambda f, a, kw: "Success"
# Can succeed after a backend has raised BackendNotImplementedError
with ua.set_backend(be3), ua.set_backend(be):
assert nullary_mm() == "Success"
def test_determine_backend_coerce(nullary_mm):
class TypeA:
pass
class TypeB:
pass
mark = "determine_backend_test"
class TypeBackend:
__ua_domain__ = "ua_tests"
def __init__(self, my_type):
self.my_type = my_type
def __ua_convert__(self, dispatchables, coerce):
if len(dispatchables) > 0:
print(dispatchables[0], coerce)
if coerce and all(d.coercible for d in dispatchables):
return tuple(self.my_type() for _ in dispatchables)
if not all(
type(d.value) is self.my_type and d.type is mark
for d in dispatchables):
return NotImplemented
return tuple(d.value for d in dispatchables)
def __ua_function__(self, func, args, kwargs):
return self.my_type
BackendA = TypeBackend(TypeA)
BackendB = TypeBackend(TypeB)
unary_mm = ua.generate_multimethod(lambda a: (ua.Dispatchable(a, mark), ),
lambda a, kw, d: (d, kw), "ua_tests")
# coercion is not forced on the existing set backend
with ua.set_backend(BackendA), ua.set_backend(BackendB):
with ua.determine_backend(TypeA(),
mark,
domain="ua_tests",
coerce=True):
assert nullary_mm() is TypeA
assert unary_mm(TypeB()) is TypeA
# But is allowed if the backend was set with coerce in the first place
with ua.set_backend(BackendA), ua.set_backend(BackendB, coerce=True):
with ua.determine_backend(TypeA(),
mark,
domain="ua_tests",
coerce=True):
assert nullary_mm() is TypeB
assert unary_mm(TypeA()) is TypeB
def test_separation_unary(backend, u, diff_ndim, func, diff_u):
try:
with ua.set_backend(backend), ua.set_backend(udiff, coerce=True):
u = np.asarray(u)
u.var = udiff.Variable("u", diff_ndim=diff_ndim)
ret = func(u)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS:
raise
pytest.xfail(reason="The backend has no implementation for this ufunc.")
if isinstance(ret, da.Array):
ret.compute()
assert_allclose(ret.diffs[u].arr, diff_u.tolist())
def test_array_creation(backend, method, args, kwargs):
backend, types = backend
if method is np.frombuffer:
buffer = onp.array([1, 2, 3]).tobytes()
args = args + (buffer,)
for dtype in dtypes:
try:
with ua.set_backend(backend, coerce=True):
kwargs["dtype"] = dtype
ret = method(*args, **kwargs)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS and (backend, method) not in EXCEPTIONS:
raise
pytest.xfail(reason="The backend has no implementation for this ufunc.")
except TypeError:
if method is np.asanyarray:
raise pytest.xfail(
reason="The ufunc for this backend got an unexpected keyword."
)
else:
raise
assert isinstance(ret, types)
if isinstance(ret, da.Array):
ret.compute()
assert ret.dtype == dtype
def test_coercion():
torch = pytest.importorskip('torch')
arr = torch.eye(5)
with ua.set_backend(ula.numpy_backend.NumpyBackend, coerce=True):
assert isinstance(ula.svd(arr, compute_uv=False), np.ndarray)
assert isinstance(ula.svd(arr, compute_uv=False), torch.Tensor)
def test_binary_function(backend, mode, func, u_d, v_d, u_domain, v_domain):
if u_domain is None:
u_arr = generate_test_data()
else:
u_arr = generate_test_data(a=u_domain[0], b=u_domain[1])
if v_domain is None:
v_arr = generate_test_data()
else:
v_arr = generate_test_data(a=v_domain[0], b=v_domain[1])
expect_u_diff = [u_d(ua, va) for ua, va in zip(u_arr, v_arr)]
expect_v_diff = [v_d(ua, va) for ua, va in zip(u_arr, v_arr)]
try:
with ua.set_backend(udiff.DiffArrayBackend(backend, mode=mode),
coerce=True):
u = np.asarray(u_arr)
v = np.asarray(v_arr)
y = func(u, v)
u_diff = y.to(u)
v_diff = y.to(v)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS:
raise
pytest.xfail(
reason="The backend has no implementation for this ufunc.")
except NotImplementedError:
pytest.xfail(reason="The func has no implementation in the {} mode.".
format(mode))
if isinstance(y, da.Array):
y.compute()
assert_allclose(u_diff.value, expect_u_diff)
assert_allclose(v_diff.value, expect_v_diff)
def test_linalg(backend, method, args, kwargs):
backend, types = backend
try:
with ua.set_backend(backend, coerce=True):
ret = method(*args, **kwargs)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS and (backend, method) not in EXCEPTIONS:
raise
pytest.xfail(reason="The backend has no implementation for this ufunc.")
if method in {
np.linalg.qr,
np.linalg.svd,
np.linalg.eig,
np.linalg.eigh,
np.linalg.slogdet,
np.linalg.lstsq,
}:
assert all(isinstance(arr, types) for arr in ret)
for arr in ret:
if isinstance(arr, da.Array):
arr.compute()
else:
assert isinstance(ret, types)
if isinstance(ret, da.Array):
ret.compute()
def wrapped(shape, *args, **kwargs):
if isinstance(shape, collections.abc.Iterable):
shape = tuple(int(s) for s in shape)
else:
shape = (int(shape), )
# Estimate 100 Mi elements per block
blocksize = int((100 * (2**20))**(1 / len(shape)))
chunks = []
for l in shape:
chunks.append([])
while l > 0:
s = max(min(blocksize, l), 0)
chunks[-1].append(s)
l -= s
name = func.__name__ + "-" + hex(random.randrange(2**64))
dsk = {}
with set_backend(self._inner):
for chunk_id in itertools.product(
*map(lambda x: range(len(x)), chunks)):
shape = tuple(chunks[i][j] for i, j in enumerate(chunk_id))
dsk[(name, ) + chunk_id] = func(shape, *args, **kwargs)
meta = func(tuple(0 for _ in shape), *args, **kwargs)
dtype = str(meta.dtype)
return da.Array(dsk, name, chunks, dtype=dtype, meta=meta)
def test_arbitrary_function(backend, func, y_d):
x_arr = [0.2, 0.3]
try:
with ua.set_backend(backend), ua.set_backend(udiff, coerce=True):
x = np.asarray(x_arr)
x.var = udiff.Variable('x')
ret = func(x)
y_d_arr = y_d(x)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS:
raise
pytest.xfail(reason="The backend has no implementation for this ufunc.")
if isinstance(ret, da.Array):
ret.compute()
assert_allclose(ret.diffs[x].arr, y_d_arr.arr)
def test_nested():
be = Backend()
be.__ua_function__ = lambda f, a, kw: None
ctx = ua.set_backend(be)
with ctx, ctx:
pass
def test_default(nullary_mm):
obj = object()
be = Backend()
be.__ua_function__ = lambda f, a, kw: NotImplemented
# If a backend returns NotImplemented, the default is called
def default1(*a, **kw):
return obj
mm1 = ua.generate_multimethod(lambda: (),
lambda a, kw, d: (a, kw),
"ua_tests",
default=default1)
with ua.set_backend(be):
assert mm1() is obj
# If all backends fail, the default is called again without a specific backend
num_calls = [0]
def default2(*a, **kw):
num_calls[0] = num_calls[0] + 1
raise ua.BackendNotImplementedError()
mm2 = ua.generate_multimethod(lambda: (),
lambda a, kw, d: (a, kw),
"ua_tests",
default=default2)
with ua.set_backend(be), pytest.raises(ua.BackendNotImplementedError):
mm2()
assert num_calls[0] == 2
# If the last backend is set as only or coerce, the last default call is skipped
num_calls[0] = 0
with ua.set_backend(be, only=True), pytest.raises(
ua.BackendNotImplementedError):
mm2()
assert num_calls[0] == 1
num_calls[0] = 0
with ua.set_backend(be, coerce=True), pytest.raises(
ua.BackendNotImplementedError):
mm2()
assert num_calls[0] == 1
def test_functions_coerce(backend, method, args, kwargs):
backend, types = backend
try:
with ua.set_backend(backend, coerce=True):
ret = method(*args, **kwargs)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS and (backend, method) not in EXCEPTIONS:
raise
pytest.xfail(reason="The backend has no implementation for this ufunc.")
except TypeError:
if backend is CupyBackend:
if method is np.flip:
pytest.xfail(reason="CuPy requires axis argument")
elif method in {np.repeat, np.tile}:
pytest.xfail(reason="CuPy does not accept array repeats")
raise
except ValueError:
if isinstance(backend, DaskBackend) and method is np.place:
pytest.xfail(reason="Default relies on delete and copyto")
if backend is CupyBackend and method in {np.argwhere, np.block}:
pytest.xfail(reason="Default relies on array_like coercion")
raise
except NotImplementedError:
if backend is CupyBackend and method is np.sort_complex:
pytest.xfail(reason="CuPy cannot sort complex data")
raise
except AttributeError:
if backend is CupyBackend and method is np.lexsort:
pytest.xfail(reason="CuPy doesn't accept tuples of arrays")
raise
if method is np.shape:
assert isinstance(ret, tuple) and all(isinstance(s, int) for s in ret)
elif method in (np.ndim, np.size):
assert isinstance(ret, int)
elif method in (
np.allclose,
np.iscomplex,
np.iscomplexobj,
np.isreal,
np.isrealobj,
np.isscalar,
np.array_equal,
np.array_equiv,
):
assert isinstance(ret, (bool,) + types)
elif method in {np.place, np.put, np.put_along_axis, np.putmask, np.fill_diagonal}:
assert ret is None
elif method in {np.nditer, np.ndenumerate, np.ndindex}:
assert isinstance(ret, collections.abc.Iterator)
elif method is np.lib.Arrayterator:
assert isinstance(ret, collections.abc.Iterable)
else:
assert isinstance(ret, types)
if isinstance(ret, da.Array):
ret.compute()
def test_invalid():
be1 = Backend()
be1.__ua_function__ = lambda f, a, kw: None
be2 = Backend()
be2.__ua_function__ = lambda f, a, kw: None
ctx1 = ua.set_backend(be1)
ctx2 = ua.set_backend(be2)
with pytest.raises(RuntimeError):
try:
ctx1.__enter__()
try:
ctx2.__enter__()
finally:
ctx1.__exit__(None, None, None)
finally:
ctx2.__exit__(None, None, None)
def test_pickle_state():
ua.set_global_backend(ComparableBackend("a"))
ua.register_backend(ComparableBackend("b"))
with ua.set_backend(ComparableBackend("c")), ua.skip_backend(
ComparableBackend("d")):
state = ua.get_state()
state_loaded = pickle.loads(pickle.dumps(state))
assert state._pickle() == state_loaded._pickle()
def repeat_to_match_shape(g, shape, dtype, axis, keepdims):
"""Returns the array g repeated along axis to fit vector space vs.
Also returns the number of repetitions of the array."""
with ua.set_backend(numpy_backend, coerce=True):
if shape == ():
return g, 1
axis = list(axis) if isinstance(axis, tuple) else axis
new_shape = np.array(shape, dtype=int)
new_shape[axis] = 1
num_reps = np.prod(np.array(shape)[axis])
return np.broadcast_to(np.reshape(g, new_shape), shape), num_reps
def test_arbitrary_function(backend, func, y_d, domain):
if domain is None:
x_arr = generate_test_data()
else:
x_arr = generate_test_data(a=domain[0], b=domain[1])
try:
with ua.set_backend(backend), ua.set_backend(udiff, coerce=True):
x = np.asarray(x_arr)
x.var = udiff.Variable("x")
ret = func(x)
y_d_arr = y_d(x)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS:
raise
pytest.xfail(reason="The backend has no implementation for this ufunc.")
if isinstance(ret, da.Array):
ret.compute()
assert_allclose(ret.diffs[x].arr, y_d_arr.arr)
def test_ufuncs_coerce(backend, method, args, kwargs):
backend, types = backend
try:
with ua.set_backend(backend, coerce=True):
ret = method(*args, **kwargs)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS:
raise
pytest.xfail(
reason="The backend has no implementation for this ufunc.")
assert isinstance(ret, types)
def overridden_class(self, self2):
"""
Convert ndarray to VJPDiffArray or JVPDiffArray according to mode.
"""
if self is ndarray:
if self._mode == "vjp":
return VJPDiffArray
else:
return JVPDiffArray
with ua.set_backend(self._inner, only=True):
return self2.overridden_class
def test_ufuncs_results(backend, method, args, kwargs, res):
backend, types = backend
try:
with ua.set_backend(backend, coerce=True):
ret = method(*args, **kwargs)
res = np.asarray(res)
assert np.allclose(ret, res, equal_nan=True)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS:
raise
pytest.xfail(reason="The backend has no implementation for this ufunc.")
def ctx_before_global(nullary_mm):
obj = object()
obj2 = object()
be = Backend()
be.__ua_function__ = lambda f, a, kw: obj
be2 = Backend()
be2.__ua_function__ = lambda f, a, kw: obj2
ua.set_global_backend(be)
with ua.set_backend(be2):
assert nullary_mm() is obj2
def test_multiple_output(backend, method, args, kwargs):
backend, types = backend
try:
with ua.set_backend(backend, coerce=True):
ret = method(*args, **kwargs)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS and (backend,
method) not in EXCEPTIONS:
raise
pytest.xfail(
reason="The backend has no implementation for this ufunc.")
assert all(isinstance(arr, types) for arr in ret)
def test_getset_state(cleanup_backends):
ua.set_global_backend(Backend())
ua.register_backend(Backend())
with ua.set_backend(Backend()), ua.skip_backend(Backend()):
state = ua.get_state()
pstate = state._pickle()
assert pstate != ua.get_state()._pickle()
with ua.set_state(state):
assert pstate[:2] == ua.get_state()._pickle()[:2]
def test_hierarchical_backends():
mm = ua.generate_multimethod(lambda: (), lambda a, kw, d: (a, kw),
"ua_tests.foo.bar")
subdomains = "ua_tests.foo.bar".split(".")
depth = len(subdomains)
mms = [
ua.generate_multimethod(lambda: (), lambda a, kw, d: (a, kw),
".".join(subdomains[:i + 1]))
for i in range(depth)
]
be = [DisableBackend(".".join(subdomains[:i + 1])) for i in range(depth)]
ua.set_global_backend(be[1])
with pytest.raises(ua.BackendNotImplementedError):
mms[0]()
for i in range(1, depth):
assert mms[i]() is be[1].ret
ua.set_global_backend(be[0])
for i in range(depth):
assert mms[i]() is be[min(i, 1)].ret
ua.set_global_backend(be[2])
for i in range(depth):
assert mms[i]() is be[i].ret
be[2].active = False
for i in range(depth):
print(i)
assert mms[i]() is be[min(i, 1)].ret
be[1].active = False
for i in range(depth):
assert mms[i]() is be[0].ret
be[0].active = False
for i in range(depth):
with pytest.raises(ua.BackendNotImplementedError):
mms[i]()
# only=True prevents all further domain checking
be[0].active = True
be[1].active = True
with ua.set_backend(be[2], only=True), pytest.raises(
ua.BackendNotImplementedError):
mms[2]()
def __ua_convert__(self, value, dispatch_type, coerce):
if dispatch_type is not ufunc and value is None:
return None
if dispatch_type is ndarray:
if not coerce and not isinstance(value, da.Array):
return NotImplemented
ret = da.asarray(value)
with set_backend(self._inner):
ret = ret.map_blocks(self._wrap_current_state(unumpy.asarray))
return ret
return value
def _generic(method, args, kwargs):
try:
import numpy as np
import unumpy.numpy_backend as NumpyBackend
except ImportError:
return NotImplemented
with ua.set_backend(NumpyBackend, coerce=True):
try:
out = method(*args, **kwargs)
except TypeError:
return NotImplemented
return convert_out(out, coerce=False)
def test_functional(backend, method, args, kwargs):
backend, types = backend
try:
with ua.set_backend(backend, coerce=True):
ret = method(*args, **kwargs)
except ua.BackendNotImplementedError:
if backend in FULLY_TESTED_BACKENDS and (backend, method) not in EXCEPTIONS:
raise
pytest.xfail(reason="The backend has no implementation for this ufunc.")
assert isinstance(ret, types)
if isinstance(ret, da.Array):
ret.compute()
def test_skip_comparison(nullary_mm):
be1 = Backend()
be1.__ua_function__ = lambda f, a, kw: None
class Backend2(Backend):
@staticmethod
def __ua_function__(f, a, kw):
pass
def __eq__(self, other):
return other is self or other is be1
with pytest.raises(ua.BackendNotImplementedError):
with ua.set_backend(be1), ua.skip_backend(Backend2()):
nullary_mm()
