diff --git a/RELEASES.md b/RELEASES.md index 20a45aef2..03093ac6f 100644 --- a/RELEASES.md +++ b/RELEASES.md @@ -18,6 +18,7 @@ #### Closed issues +- Fix ordinary JAX gradients for dense `ot.emd2` and preserve the magnitude and higher-order dependence of gradients supplied to `JaxBackend.set_gradients`, including the envelope gradients in `ot.solve` (Issue #625) - Fix `ot.semidiscrete_wasserstein2_unif_circle` pairing sorted sample locations with unsorted weights, which broke permutation and rotation invariance for nonuniform weights (Issue #921) - Fix an out-of-bounds read in `ot.bsp.merge_bijections`: the number of plans was taken from `perms.shape[1]` (the number of samples) instead of `perms.shape[0]`, so merging fewer plans than samples read past the array and merging more silently dropped plans. A `perms` outside the documented `(n_plans, n_samples)` layout, an empty plan stack and rows that are not permutations are now rejected (PR #911, Issue #910) - Fix overflow to NaN in `ot.stochastic.coordinate_grad_semi_dual` by shifting before exponentiating, as `ot.stochastic.c_transform_entropic` already does. The factor cancels in the normalisation, so results are unchanged where the old code did not overflow (PR #866, Issue #264) diff --git a/ot/backend.py b/ot/backend.py index 1de3a60c8..b74377041 100644 --- a/ot/backend.py +++ b/ot/backend.py @@ -1708,16 +1708,21 @@ def _from_numpy(self, a, type_as=None): def set_gradients(self, val, inputs, grads): from jax.flatten_util import ravel_pytree - (val,) = jax.lax.stop_gradient((val,)) - ravelled_inputs, _ = ravel_pytree(inputs) ravelled_grads, _ = ravel_pytree(grads) - aux = jnp.sum(ravelled_inputs * ravelled_grads) / 2 - aux = aux - jax.lax.stop_gradient(aux) + @jax.custom_jvp + def auxiliary(args, gradients): + return jnp.zeros((), dtype=jnp.result_type(args, gradients, 0.0)) + + @auxiliary.defjvp + def auxiliary_jvp(primals, tangents): + args, gradients = primals + args_dot, _ = tangents + return auxiliary(args, gradients), jnp.sum(args_dot * gradients) - (val,) = jax.tree_util.tree_map(lambda z: z + aux, (val,)) - return val + aux = auxiliary(ravelled_inputs, ravelled_grads) + return jax.tree_util.tree_map(lambda z: jax.lax.stop_gradient(z) + aux, val) def _detach(self, a): return jax.lax.stop_gradient(a) diff --git a/ot/lp/_network_simplex.py b/ot/lp/_network_simplex.py index a3f3d548d..e8b08cdc8 100644 --- a/ot/lp/_network_simplex.py +++ b/ot/lp/_network_simplex.py @@ -691,6 +691,9 @@ def emd2( else: type_as = a if a is not None else b + if type_as is not None: + type_as = nx.detach(type_as) + # Set n1, n2 if not already set (dense case) if n1 is None: n1, n2 = M.shape @@ -713,7 +716,7 @@ def emd2( if is_sparse: a, b = nx.to_numpy(a, b) else: - M, a, b = nx.to_numpy(M, a, b) + M, a, b = nx.to_numpy(*nx.detach(M, a, b)) M = np.asarray(M, dtype=np.float64, order="C") a = np.asarray(a, dtype=np.float64) diff --git a/test/test_backend.py b/test/test_backend.py index 7fa0010b4..0b784aa62 100644 --- a/test/test_backend.py +++ b/test/test_backend.py @@ -1008,3 +1008,60 @@ def test_no_cuda_context_for_cpu_only_work(): f"interpreter exited with returncode {result.returncode}: " f"{result.stderr.decode(errors='replace')[-2000:]}" ) + + +@pytest.mark.parametrize("power,expected,second", [(0, 3., 0.), (1, 6., 3.), (2, 12., 12.)]) +def test_jax_supplied_gradients(power, expected, second): + jax = pytest.importorskip("jax") + nx = ot.backend.JaxBackend() + + def loss(x): + return nx.set_gradients(x**4, (x,), (3 * x**power,)) + + np.testing.assert_allclose(loss(2.), 16.) + with jax.checking_leaks(): + np.testing.assert_allclose(jax.grad(loss)(2.), expected) + np.testing.assert_allclose(jax.jvp(loss, (2.,), (1.,))[1], expected) + np.testing.assert_allclose(jax.grad(lambda x: 7 * loss(x))(2.), 7 * expected) + np.testing.assert_allclose(jax.grad(jax.grad(loss))(2.), second) + np.testing.assert_allclose(jax.hessian(loss)(2.), second) + + +def test_jax_supplied_gradients_nested_inputs(): + jax = pytest.importorskip("jax") + jnp = jax.numpy + nx = ot.backend.JaxBackend() + + def loss(x, y): + inputs = {"x": x, "y": y} + grads = {"x": 2 * x + jnp.sum(y), "y": jnp.ones_like(y) * jnp.sum(x)} + return nx.set_gradients(jnp.sum(x**4) + jnp.sum(y**4), inputs, grads) + + x, y = jnp.array([1., 2.]), jnp.array([3., 4., 5.]) + with jax.checking_leaks(): + dx, dy = jax.grad(loss, argnums=(0, 1))(x, y) + hessian = jax.hessian(loss, argnums=(0, 1))(x, y) + np.testing.assert_allclose(dx, [14., 16.]) + np.testing.assert_allclose(dy, [3., 3., 3.]) + np.testing.assert_allclose(hessian[0][0], [[2., 0.], [0., 2.]]) + np.testing.assert_allclose(hessian[0][1], np.ones((2, 3))) + np.testing.assert_allclose(hessian[1][0], np.ones((3, 2))) + np.testing.assert_allclose(hessian[1][1], np.zeros((3, 3))) + + +@pytest.mark.parametrize("value_dtype", [np.float32, np.float64, np.int32]) +def test_jax_supplied_gradients_value_dtype(value_dtype): + jax = pytest.importorskip("jax") + jnp = jax.numpy + nx = ot.backend.JaxBackend() + + def loss(x): + value = {"a": jnp.array(2, dtype=value_dtype), "b": jnp.array([4, 5], dtype=value_dtype)} + result = nx.set_gradients(value, (x,), (3 * x,)) + return result["a"] + jnp.sum(result["b"]) + + x = jnp.array(2., dtype=jnp.float64) + np.testing.assert_allclose(loss(x), 11.) + with jax.checking_leaks(): + np.testing.assert_allclose(jax.grad(loss)(x), 18.) + np.testing.assert_allclose(jax.hessian(loss)(x), 9.) diff --git a/test/test_ot.py b/test/test_ot.py index 9b4c9545a..0e189cc00 100644 --- a/test/test_ot.py +++ b/test/test_ot.py @@ -190,6 +190,76 @@ def test_emd2_gradients(): ) +@pytest.mark.parametrize("dtype", [np.float32, np.float64]) +@pytest.mark.parametrize("log,return_matrix", [(False, False), (True, False), (False, True)]) +@pytest.mark.parametrize("argnums", [0, 1, 2, (0, 1, 2)]) +def test_emd2_jax_gradients(dtype, log, return_matrix, argnums): + jax = pytest.importorskip("jax") + jnp = jax.numpy + a = jnp.array([0.2, 0.35, 0.45], dtype=dtype) + b = jnp.array([0.4, 0.3, 0.3], dtype=dtype) + M = jnp.array([[0, 1, 4], [2, 0.2, 1.5], [5, 2, 0.1]], dtype=dtype) + plan = np.array([[0.2, 0, 0], [0.2, 0.15, 0], [0, 0.15, 0.3]]) + expected = (np.array([-29, 1, 28]) / 15, np.array([55, 1, -56]) / 30, plan) + tolerance = 2e-4 if dtype == np.float32 else 1e-8 + step = 1e-3 if dtype == np.float32 else 1e-5 + + def loss(a, b, M): + result = ot.emd2(a, b, M, log=log, return_matrix=return_matrix) + return result[0] if log or return_matrix else result + + value = loss(a, b, M) + np.testing.assert_allclose(value, 0.76, atol=tolerance) + assert value.dtype == a.dtype + assert value.devices() == a.devices() + if return_matrix: + _, info = ot.emd2(a, b, M, return_matrix=True) + np.testing.assert_allclose(info["G"], plan, atol=tolerance) + + with jax.checking_leaks(): + gradients = jax.grad(loss, argnums=argnums)(a, b, M) + scaled = jax.grad(lambda a, b, M: 7 * loss(a, b, M), argnums=argnums)(a, b, M) + indices = argnums if isinstance(argnums, tuple) else (argnums,) + if not isinstance(argnums, tuple): + gradients, scaled = (gradients,), (scaled,) + for index, gradient, scaled_gradient in zip(indices, gradients, scaled): + np.testing.assert_allclose(gradient, expected[index], atol=tolerance) + np.testing.assert_allclose(scaled_gradient, 7 * expected[index], atol=7 * tolerance) + directions = np.eye(9).reshape(9, 3, 3) if index == 2 else np.array([[1, -1, 0], [0, 1, -1]]) + for direction in directions: + plus, minus = [a, b, M], [a, b, M] + plus[index] = plus[index] + step * direction + minus[index] = minus[index] - step * direction + finite_difference = (float(loss(*plus)) - float(loss(*minus))) / (2 * step) + np.testing.assert_allclose(np.sum(np.asarray(gradient) * direction), finite_difference, atol=tolerance) + + +@pytest.mark.parametrize("empty_a,empty_b", [(True, False), (False, True), (True, True)]) +def test_emd2_jax_empty_marginals(empty_a, empty_b): + jax = pytest.importorskip("jax") + jnp = jax.numpy + a = jnp.array([] if empty_a else [0.2, 0.35, 0.45]) + b = jnp.array([] if empty_b else [0.4, 0.3, 0.3]) + M = jnp.array([[0., 1., 4.], [2., 0.2, 1.5], [5., 2., 0.1]]) + reference_a = np.ones(3) / 3 if empty_a else np.asarray(a) + reference_b = np.ones(3) / 3 if empty_b else np.asarray(b) + plan = ot.emd(reference_a, reference_b, np.asarray(M)) + with jax.checking_leaks(): + gradients = jax.grad(ot.emd2, argnums=(0, 1, 2))(a, b, M) + assert gradients[0].shape == a.shape + assert gradients[1].shape == b.shape + np.testing.assert_allclose(gradients[2], plan, atol=1e-8) + for index, empty in enumerate((empty_a, empty_b)): + if empty: + continue + for direction in np.array([[1., -1., 0.], [0., 1., -1.]]): + plus, minus = [a, b, M], [a, b, M] + plus[index] = plus[index] + 1e-5 * direction + minus[index] = minus[index] - 1e-5 * direction + finite_difference = (float(ot.emd2(*plus)) - float(ot.emd2(*minus))) / 2e-5 + np.testing.assert_allclose(np.sum(gradients[index] * direction), finite_difference, atol=1e-8) + + def test_emd_emd2(): # test emd and emd2 for simple identity n = 100 diff --git a/test/test_solvers.py b/test/test_solvers.py index ba16a07a3..b08c594c3 100644 --- a/test/test_solvers.py +++ b/test/test_solvers.py @@ -1233,3 +1233,34 @@ def custom_metric(X_b, X_a): metric=custom_metric, max_iter_bary=1, ) + + +@pytest.mark.parametrize("reg_type", ["entropy", "KL"]) +def test_solve_jax_envelope_gradients(reg_type): + jax = pytest.importorskip("jax") + jnp = jax.numpy + M = np.array([[0.1, 0.3, 0.5], [0.4, 0.2, 0.3], [0.6, 0.4, 0.1]]) + a = b = np.ones(3) / 3 + + def loss(M, a, b): + return ot.solve( + M, a, b, reg=0.1, reg_type=reg_type, grad="envelope", tol=1e-12 + ).value + + inputs = (M, a, b) + with jax.checking_leaks(): + gradients = jax.grad(loss, argnums=(0, 1, 2))(*map(jnp.asarray, inputs)) + for index, gradient in enumerate(gradients): + directions = ( + np.eye(9).reshape(9, 3, 3) + if index == 0 + else np.array([[1., -1., 0.], [0., 1., -1.]]) + ) + for direction in directions: + plus, minus = list(inputs), list(inputs) + plus[index] = plus[index] + 1e-5 * direction + minus[index] = minus[index] - 1e-5 * direction + finite_difference = (loss(*plus) - loss(*minus)) / 2e-5 + np.testing.assert_allclose( + np.sum(np.asarray(gradient) * direction), finite_difference, atol=1e-7 + )