Function fun

test/test_backend.py:800–802

(a, b, d)

        with jax.checking_leaks():

            def fun(a, b, d):
                val = b * nx.sum(a**4) + d
                return nx.set_gradients(val, (a, b, d), (a, b, 2 * d))

            grad_val = jax.grad(fun, argnums=(0, 1, 2))(v, c, e)