[WIP] Propagate upstream gradients in the TensorFlow backend

RELEASES.md

@@ -18,6 +18,8 @@
 
 #### Closed issues
 
+- Fix TensorFlow custom gradients ignoring the upstream gradient, so scaling an OT loss now scales its supplied gradients consistently.
+
 - 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

@@ -3234,7 +3234,7 @@ def set_gradients(self, val, inputs, grads):
         @tf.custom_gradient
         def tmp(input):
             def grad(upstream):
-                return grads
+                return tf.nest.map_structure(lambda g: upstream * g, grads)
 
             return val, grad
 

test/test_backend.py

@@ -1008,3 +1008,40 @@ 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.skipif(not tf, reason="TensorFlow not installed")
+@pytest.mark.parametrize("coefficient", [0.0, 1.0, -2.0, 7.0])
+@pytest.mark.parametrize("compiled", [False, True])
+def test_tensorflow_set_gradients_chain_rule(coefficient, compiled):
+    nx = ot.backend.TensorflowBackend()
+    x = tf.constant([1.0, 2.0], dtype=tf.float64)
+    y = tf.constant(3.0, dtype=tf.float64)
+
+    def differentiate(x, y):
+        with tf.GradientTape() as tape:
+            tape.watch((x, y))
+            value = nx.set_gradients(
+                tf.constant(5.0, dtype=x.dtype), (x, y), (2 * x, 3 * y)
+            )
+            loss = coefficient * value
+        return value, tape.gradient(loss, (x, y))
+
+    if compiled:
+        differentiate = tf.function(differentiate)
+    value, gradients = differentiate(x, y)
+    np.testing.assert_allclose(value, 5.0)
+    np.testing.assert_allclose(gradients[0], coefficient * 2 * x)
+    np.testing.assert_allclose(gradients[1], coefficient * 3 * y)
+
+
+@pytest.mark.skipif(not tf, reason="TensorFlow not installed")
+@pytest.mark.parametrize("coefficient", [0.0, -2.0, 7.0])
+def test_tensorflow_emd2_scaled_gradient(coefficient):
+    a = tf.constant([0.5, 0.5], dtype=tf.float64)
+    cost = tf.constant([[0.0, 2.0], [3.0, 0.0]], dtype=tf.float64)
+    with tf.GradientTape() as tape:
+        tape.watch(cost)
+        loss = coefficient * ot.emd2(a, a, cost)
+    gradient = tape.gradient(loss, cost)
+    np.testing.assert_allclose(gradient, coefficient * np.diag([0.5, 0.5]))