code: def f(p, x):
      return jnp.expand_dims(jnp.sum(p * (x**2)), -1)
  
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 1 for calling a global function with example usage: f(p, x) and returns: jnp