| 773 | |
| 774 | |
| 775 | def test_gradients_backends(): |
| 776 | rnd = np.random.RandomState(0) |
| 777 | v = rnd.randn(10) |
| 778 | c = rnd.randn() |
| 779 | e = rnd.randn() |
| 780 | |
| 781 | if torch: |
| 782 | nx = ot.backend.TorchBackend() |
| 783 | |
| 784 | v2 = torch.tensor(v, requires_grad=True) |
| 785 | c2 = torch.tensor(c, requires_grad=True) |
| 786 | |
| 787 | val = c2 * torch.sum(v2 * v2) |
| 788 | |
| 789 | val2 = nx.set_gradients(val, (v2, c2), (v2, c2)) |
| 790 | |
| 791 | val2.backward() |
| 792 | |
| 793 | assert torch.equal(v2.grad, v2) |
| 794 | assert torch.equal(c2.grad, c2) |
| 795 | |
| 796 | if jax: |
| 797 | nx = ot.backend.JaxBackend() |
| 798 | with jax.checking_leaks(): |
| 799 | |
| 800 | def fun(a, b, d): |
| 801 | val = b * nx.sum(a**4) + d |
| 802 | return nx.set_gradients(val, (a, b, d), (a, b, 2 * d)) |
| 803 | |
| 804 | grad_val = jax.grad(fun, argnums=(0, 1, 2))(v, c, e) |
| 805 | |
| 806 | np.testing.assert_almost_equal(fun(v, c, e), c * np.sum(v**4) + e, decimal=4) |
| 807 | np.testing.assert_allclose(grad_val[0], v, atol=1e-4) |
| 808 | np.testing.assert_allclose(grad_val[2], 2 * e, atol=1e-4) |
| 809 | |
| 810 | if tf: |
| 811 | nx = ot.backend.TensorflowBackend() |
| 812 | w = tf.Variable(tf.random.normal((3, 2)), name="w") |
| 813 | b = tf.Variable(tf.random.normal((2,), dtype=tf.float32), name="b") |
| 814 | x = tf.random.normal((1, 3), dtype=tf.float32) |
| 815 | |
| 816 | with tf.GradientTape() as tape: |
| 817 | y = x @ w + b |
| 818 | loss = tf.reduce_mean(y**2) |
| 819 | manipulated_loss = nx.set_gradients(loss, (w, b), (w, b)) |
| 820 | [dl_dw, dl_db] = tape.gradient(manipulated_loss, [w, b]) |
| 821 | assert nx.allclose(dl_dw, w) |
| 822 | assert nx.allclose(dl_db, b) |
| 823 | |
| 824 | |
| 825 | def test_get_backend_none(): |