code: def test_grad(getkey):
      mlp = eqx.nn.MLP(2, 2, 512, 2, key=getkey())
      mlp_noinline = eqx.internal.noinline(mlp)
      mlp_grad = jax.grad(lambda x: jnp.sum(mlp(x)))
      mlp_jit_grad = eqx.filter_jit(mlp_grad, donate="none")
      mlp_grad_noinline = jax.grad(lambda x: jnp.sum(mlp_noinline(x)))
      mlp_jit_grad_noinline = eqx.filter_jit(
          mlp_grad_noinline, donate="none"
      )
      x = jr.normal(getkey(), (2,))
      o1 = mlp_grad(x)
      o2 = mlp_jit_grad(x)
      o3 = mlp_grad_noinline(x)
      o4 = mlp_jit_grad_noinline(x)
      assert tree_allclose(o1, o2)
      assert tree_allclose(o1, o3)
      assert tree_allclose(o1, o4)
  
file_path: equinox/tests/test_noinline.py
imports: import equinox
  import jax
  import jax.numpy
  import jax.random
  from .helpers import tree_allclose
label_desc: Use image node_id 5 for calling a global function with example usage: test_grad(getkey) without return types