diff --git a/RELEASES.md b/RELEASES.md index 20a45aef2..908b909e6 100644 --- a/RELEASES.md +++ b/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) diff --git a/ot/backend.py b/ot/backend.py index 1de3a60c8..955662b22 100644 --- a/ot/backend.py +++ b/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 diff --git a/test/test_backend.py b/test/test_backend.py index 7fa0010b4..89084103a 100644 --- a/test/test_backend.py +++ b/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]))