code: def test_vprim():
      def impl(x):
          assert x.shape == (2,)
          return 2 * x, jnp.concatenate([x, jnp.flip(x)])
  
      def abstract(x):
          assert type(x) is jax.core.ShapedArray
          return x, jax.core.ShapedArray((4,), x.dtype)
  
      def jvp(primals, tangents):
          (x,) = primals
          (tx,) = tangents
          primals_out = newprim_p.bind(x)
          tangents_out = (jnp.flip(tx), jnp.concatenate([x, x + 1]))
          return primals_out, tangents_out
  
      def transpose(cts_out, x):
          ct_out1, ct_out2 = cts_out
          assert ct_out1.shape == (2,)
          assert ct_out2.shape == (4,)
          return [ct_out1 + ct_out2[:2]]
  
      newprim_p = eqxi.create_vprim(
          "newprim", impl, abstract, jvp, transpose
      )
      bind = newprim_p.bind
      y = jnp.array([1.0, 3.5])
      y2 = jnp.array([[1.0, 3.5], [2.0, 1.5], [0.0, 0.2]])
  
      o1 = [jnp.array([2.0, 7.0]), jnp.array([1.0, 3.5, 3.5, 1.0])]
      assert tree_allclose(bind(y), o1)
      assert tree_allclose(jax.jit(bind)(y), o1)
  
      dtype = jnp.array(1.0).dtype  # default floating-point dtype
      o2 = [
          jax.ShapeDtypeStruct((2,), dtype),
          jax.ShapeDtypeStruct((4,), dtype),
      ]
      assert tree_allclose(jax.eval_shape(bind, y), o2)
  
      t_o1 = [jnp.array([5.5, 3.0]), jnp.array([1.0, 3.5, 2.0, 4.5])]
      o3 = (o1, t_o1)
      assert tree_allclose(jax.jvp(bind, (y,), (y + 2,)), o3)
  
      o4 = [
          jnp.array([[2.0, 7.0], [4.0, 3.0], [0.0, 0.4]]),
          jnp.array(
              [
                  [1.0, 3.5, 3.5, 1.0],
                  [2.0, 1.5, 1.5, 2.0],
                  [0.0, 0.2, 0.2, 0.0],
              ]
          ),
      ]
      assert tree_allclose(jax.vmap(bind)(y2), o4)
      assert tree_allclose(jax.jit(jax.vmap(bind))(y2), o4)
      assert tree_allclose(jax.vmap(jax.jit(bind))(y2), o4)
  
      assert tree_allclose(
          jax.vmap(jax.vmap(bind))(y2[None]), [o4[0][None], o4[1][None]]
      )
  
      o5 = [
          jax.ShapeDtypeStruct((3, 2), dtype),
          jax.ShapeDtypeStruct((3, 4), dtype),
      ]
      assert tree_allclose(jax.eval_shape(jax.vmap(bind), y2), o5)
  
      t_o4 = [
          jnp.array([[5.5, 3.0], [3.5, 4.0], [2.2, 2.0]]),
          jnp.array(
              [
                  [1.0, 3.5, 2.0, 4.5],
                  [2.0, 1.5, 3.0, 2.5],
                  [0.0, 0.2, 1.0, 1.2],
              ]
          ),
      ]
      o6 = (o4, t_o4)
      assert tree_allclose(
          jax.jvp(jax.vmap(bind), (y2,), (y2 + 2,)), o6
      )
  
      o7 = jnp.array([3.0, 10.5])
      assert tree_allclose(jax.linear_transpose(bind, y)(o1), (o7,))
  
      o8 = jnp.array([[3.0, 10.5], [6.0, 4.5], [0.0, 0.6]])
      assert tree_allclose(
          jax.linear_transpose(jax.vmap(bind), y2)(o4), (o8,)
      )
  
file_path: equinox/tests/test_primitive.py
imports: import functools
  import equinox
  import equinox.internal
  import jax
  import jax.core
  import jax.interpreters.ad
  import jax.interpreters.batching
  import jax.interpreters.mlir
  import jax.numpy
  import jax.tree_util
  from .helpers import tree_allclose
label_desc: Use image node_id 2 for calling a global function with example usage: test_vprim() and returns: unknown,  jnp, x,  jax, primals_out,  tangents_out, list