code: def test_grad():
      # Checks that `jax.grad` can be applied to functions
      # that use operators.
      def loss(x, y):
          z = x**2 + y * 5 + x * y
          return jnp.sum(z.params)
  
      gradfn = grad(loss, argnums=(0, 1))
      x_grad, y_grad = gradfn(X, Y)
      _ = x_grad
      assert x_grad.params == 4.0
      assert y_grad.params == 6.0
  
imports: import jax
  from jax import grad, jit, make_jaxpr
  from jax import numpy
  from plum.function import NotFoundLookupError
  from pytest import mark, raises
  from jaxdf import *
  from jaxdf.operators.dummy import dummy, yummy
  from jaxdf.signatures import SignatureError
label_desc: Use image node_id 12 for calling a global function with example usage: test_grad() and returns: jnp