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)
