code: def test_readme_example():
      # Defining operator
      @operator
      def custom_op(u, *, params=None):
          grad_u = jops.gradient(u)
          diag_jacobian = jops.diag_jacobian(grad_u)
          laplacian = jops.sum_over_dims(diag_jacobian)
          sin_u = jops.compose(u)(jnp.sin)
          return laplacian + sin_u
  
      # Defining discretizations
      domain = Domain((128, 128), (1.0, 1.0))
      parameters = jnp.ones((128, 128, 1))
      u = FourierSeries(parameters, domain)
  
      # Define a differentiable loss function
      @jit
      def loss(u):
          v = custom_op(u)
          return jnp.mean(jnp.abs(v.on_grid) ** 2)
  
      gradient = grad(loss)(u)
      assert gradient.on_grid.shape == (128, 128, 1)
      assert type(gradient) == FourierSeries
  
imports: from jax import grad, jit
  from jax import numpy
  from jaxdf import FourierSeries, operator
  from jaxdf import operators
  from jaxdf.geometry import Domain
label_desc: Use image node_id 1 for calling a global function with example usage: test_readme_example() and returns: unknown, jnp