[MRG] Fix JAX gradients for dense emd2

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)

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)

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)

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

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

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