From f70c5f275bbd8782913928f1a0c108304a272031 Mon Sep 17 00:00:00 2001 From: Zerui Ge Date: Mon, 9 Nov 2026 15:47:20 +0800 Subject: [PATCH] Fix JAX gradients for dense emd2 and envelope losses --- RELEASES.md | 1 + ot/backend.py | 17 ++++++---- ot/lp/_network_simplex.py | 5 ++- test/test_backend.py | 57 +++++++++++++++++++++++++++++++ test/test_ot.py | 70 +++++++++++++++++++++++++++++++++++++++ test/test_solvers.py | 31 +++++++++++++++++ 6 files changed, 174 insertions(+), 7 deletions(-) diff --git a/RELEASES.md b/RELEASES.md index 85be4654d..5253c3352 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 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) - Fix `ot.bsp.compute_bspot_bijection` (and `random_state` of `ot.solve_sample(..., method="bsp")`) not being reproducible for a given `seed`: its candidate plans came from a `thread_local` `std::mt19937` seeded once per OpenMP thread, so the result depended on the call history and on the number of threads. Every thread started from the same state, so the previous version could underperform with multiple threads. Each plan now draws its own stream, derived from the seed and the plan index (PR #898, Issue #897) 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 + )