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_copyto(backend): backend, types = backend try: with ua.set_backend(backend, coerce=True): dst = np.asarray([1, 2]) src = np.asarray([3, 4]) np.copyto(dst, src) assert np.array_equal(dst, src) except ua.BackendNotImplementedError: if backend in FULLY_TESTED_BACKENDS and (backend, np.copyto) not in EXCEPTIONS: raise pytest.xfail( reason="The backend has no implementation for this ufunc." ) assert isinstance(dst, types)
def __ua_convert__(self, value, dispatch_type, coerce): if dispatch_type is np.ndarray: if value is None: return value if isinstance(value, DiffArray): return value if coerce: with ua.set_backend(self._inner, coerce=True): if self._mode == "vjp": return VJPDiffArray(np.asarray(value)) else: return JVPDiffArray(np.asarray(value)) return NotImplemented return value
def trim_zeros(self, filt, trim="fb"): nonzero_idxs = unumpy.nonzero(filt)[0].compute_chunk_sizes() if len(nonzero_idxs.compute()) == 0: return unumpy.asarray([], dtype=filt.dtype) start, end = None, None if "f" in trim: start = nonzero_idxs[0].compute() if "b" in trim: end = nonzero_idxs[-1].compute() + 1 return filt[start:end]
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 __init__(self, arr): if isinstance(arr, DiffArray): self._value = arr.value self._id = arr.id return with ua.determine_backend(arr, np.ndarray, domain="numpy", coerce=True): arr = np.asarray(arr) self._value = arr self._id = uuid.uuid4()
def test_separation_binary( backend, u, v, u_diff_ndim, v_diff_ndim, func, diff_u, diff_v ): try: with ua.set_backend(backend), ua.set_backend(udiff, coerce=True): u = np.asarray(u) u.var = udiff.Variable("u", diff_ndim=u_diff_ndim) v = np.asarray(v) v.var = udiff.Variable("v", diff_ndim=v_diff_ndim) ret = func(u, v) 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()) assert_allclose(ret.diffs[v].arr, diff_v.tolist())
def test_separation_binary(backend, mode, u, v, func, expect_u_jacobian, expect_v_jacobian): try: with ua.set_backend(udiff.DiffArrayBackend(backend, mode=mode), coerce=True): u = np.asarray(u) v = np.asarray(v) y = func(u, v) u_jacobian = y.to(u, jacobian=True) v_jacobian = y.to(v, jacobian=True) except ua.BackendNotImplementedError: if backend in FULLY_TESTED_BACKENDS: raise pytest.xfail( reason="The backend has no implementation for this ufunc.") if isinstance(y, da.Array): y.compute() assert_allclose(u_jacobian.value, expect_u_jacobian) assert_allclose(v_jacobian.value, expect_v_jacobian)
def __ua_convert__(value, dispatch_type, coerce): if dispatch_type is np.ndarray: if value is None: return value if isinstance(value, DiffArray): return value if coerce: import udiff with udiff.SKIP_SELF: return DiffArray(np.asarray(value)) return NotImplemented return value
def test_jacobian(backend, mode, x, func, expect_jacobian): try: with ua.set_backend(udiff.DiffArrayBackend(backend, mode=mode), coerce=True): x = np.asarray(x) y = func(x) x_jacobian = y.to(x, jacobian=True) except ua.BackendNotImplementedError: if backend in FULLY_TESTED_BACKENDS: raise pytest.xfail( reason="The backend has no implementation for this ufunc.") if isinstance(y, da.Array): y.compute() assert_allclose(x_jacobian.value, expect_jacobian)
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 __init__(self, arr): if isinstance(arr, VJPDiffArray): self._value = arr.value self._id = arr.id self._parents = arr._parents self._vjp = arr._vjp self._diff = arr._diff self._jacobian = arr._jacobian return with ua.determine_backend(arr, np.ndarray, domain="numpy", coerce=True): arr = np.asarray(arr) self._value = arr self._id = uuid.uuid4() self._parents = None self._vjp = None self._diff = None self._jacobian = None
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 __init__(self, arr, diffs=None): if isinstance(arr, DiffArray): self._diffs = arr.diffs self._arr = arr.arr self._var = arr.var from udiff import SKIP_SELF with SKIP_SELF: arr = np.asarray(arr) if diffs is None: diffs = ArrayDiffRegistry(arr.shape) if not isinstance(diffs, ArrayDiffRegistry): raise ValueError("diffs must be an ArrayDiffRegistry") if diffs.shape != arr.shape: raise ValueError("diffs didn't have the right shape") self._diffs = diffs self._arr = arr self._var = None
def test_high_order_diff(backend, mode, 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]) expect_diff = [y_d(xa) for xa in x_arr] try: with ua.set_backend(udiff.DiffArrayBackend(backend, mode=mode), coerce=True): x = np.asarray(x_arr) y = func(x) x_diff = y.to(x).to(x) except ua.BackendNotImplementedError: if backend in FULLY_TESTED_BACKENDS: raise pytest.xfail( reason="The backend has no implementation for this ufunc.") if isinstance(y, da.Array): y.compute() assert_allclose(x_diff.value, expect_diff)