code: def test_basic(getkey):
      mlp_pytree = eqx.nn.MLP(2, 2, 2, 2, key=getkey())
      mlp_arg = eqx.nn.MLP(2, 2, 2, 2, key=getkey())
      mlp_closure = eqx.nn.MLP(2, 2, 2, 2, key=getkey())
      int_pytree = jnp.array([1])
      int_closure = jnp.array([1])
  
      sentinel = object()
  
      class _M(eqx.Module):
          mlp_pytree: eqx.nn.MLP
          int_pytree: jax.Array
  
          def __call__(self, mlp_arg, x):
              x = self.mlp_pytree(x)
              x = mlp_arg(x)
              x = mlp_closure(x)
              with jax.numpy_dtype_promotion("standard"):
                  x = x + self.int_pytree + int_closure
              return x, sentinel
  
      m = _M(mlp_pytree, int_pytree)
      x = jnp.array([1.0, 2.0])
      (
          closed_jaxpr,
          dynamic_out_struct,
          out_static,
      ) = eqx.filter_make_jaxpr(m)(mlp_arg, x)
      jaxpr = closed_jaxpr.jaxpr
      consts = closed_jaxpr.consts
  
      prims = co.defaultdict(int)
      for eqn in jaxpr.eqns:
          prims[eqn.primitive] += 1
      assert prims == {
          lax.dot_general_p: 9,
          lax.add_p: 11,
          custom_derivatives.custom_jvp_call_p: 6,
          lax.convert_element_type_p: 2,
      }
  
      params = {
          id(x)
          for x in jtu.tree_leaves(
              (mlp_pytree, mlp_closure, int_pytree, int_closure)
          )
          if eqx.is_array(x)
      }
      assert {id(x) for x in consts} == params
  
      assert dynamic_out_struct == (
          jax.ShapeDtypeStruct((2,), jnp.float32),
          None,
      )
      assert out_static == (None, sentinel)
  
file_path: equinox/tests/test_make_jaxpr.py
imports: import collections
  import equinox
  import jax
  import jax.custom_derivatives
  import jax.lax
  import jax.numpy
  import jax.tree_util
label_desc: Use image node_id 1 for calling a global function with example usage: test_basic(getkey) and returns: x,  sentinel