code: def test_continuous_diag_jacobian():
      # A function and its corresponding diagonal Jacobian.
      def f(params, x):
          return jnp.array([params[0] * x[0], params[1] * x[1]])
  
      def diag_jac(params, x):
          return jnp.array([params[0], params[1]])
  
      domain = geometry.Domain((8, 8), (1.0, 1.0))
      p = jnp.array([2.0, 3.0])
      a = Continuous(p, domain, f)
      b = operators.diag_jacobian(a)
  
      x = jnp.array([1.0, 1.0])  # sample point in the domain
  
      f1 = b(x)
      f2 = diag_jac(p, x)
  
      assert jnp.allclose(f1, f2)
  
file_path: jaxdf/tests/test_differential.py
imports: import jax
  from jax import jit, make_jaxpr
  from jax import numpy
  from jaxdf import *
label_desc: Use image node_id 8 for calling a global function with example usage: test_continuous_diag_jacobian() and returns: jnp, jnp