code: def test_struct():
      struct = jax.ShapeDtypeStruct((2,), jnp.float32)
  
      def f(x):
          x1 = x[None]  # (1, 2)
          x2 = x[:, None]  # (2, 1)
          # That is, what if a ShapeDtypeStruct is itself part of the static return value?
          return x1 * x2, jax.ShapeDtypeStruct((3,), jnp.float32)
  
      (
          closed_jaxpr,
          dynamic_out_struct,
          out_static,
      ) = eqx.filter_make_jaxpr(f)(struct)
      jaxpr = closed_jaxpr.jaxpr
      consts = closed_jaxpr.consts
  
      prims = co.defaultdict(int)
      for eqn in jaxpr.eqns:
          prims[eqn.primitive] += 1
      assert prims == {lax.broadcast_in_dim_p: 2, lax.mul_p: 1}
      assert consts == []
      assert dynamic_out_struct == (
          jax.ShapeDtypeStruct((2, 2), jnp.float32),
          None,
      )
      assert out_static == (
          None,
          jax.ShapeDtypeStruct((3,), jnp.float32),
      )
  
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 2 for calling a global function with example usage: test_struct() and returns: unknown,  jax