code: def bench_jax():
      import jax
      import jax.numpy as jnp
  
      jax_a = jnp.array(a.detach().numpy())
      jax_b = jnp.array(b.detach().numpy())
  
      def f(a):
          return jnp.sin((a * jax_b).sum(axis=[0])).sum()
  
      jit_f = jax.jit(jax.grad(f))
      jit_f(jax_a)
      begin = time.time()
      for _ in range(iters):
          out = jit_f(jax_a)
      out.block_until_ready()
      print(time.time() - begin)
  
imports: import time
  import torch
  import torch.utils
  from functorch.compile import aot_function, tvm_compile
label_desc: Use image node_id 3 for calling a global function with example usage: bench_jax() and returns: jnp