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()