code: def test_unvmap():
      unvmap_all = eqxi.unvmap_all
      unvmap_any = eqxi.unvmap_any
      jit_unvmap_all = jax.jit(unvmap_all)
      jit_unvmap_any = jax.jit(unvmap_any)
      vmap_unvmap_all = jax.vmap(unvmap_all, out_axes=None)
      vmap_unvmap_any = jax.vmap(unvmap_any, out_axes=None)
  
      tt = jnp.array([True, True])
      tf = jnp.array([True, False])
      ff = jnp.array([False, False])
  
      assert jnp.array_equal(unvmap_all(tt), jnp.array(True))
      assert jnp.array_equal(unvmap_all(tf), jnp.array(False))
      assert jnp.array_equal(unvmap_all(ff), jnp.array(False))
      assert jnp.array_equal(unvmap_any(tt), jnp.array(True))
      assert jnp.array_equal(unvmap_any(tf), jnp.array(True))
      assert jnp.array_equal(unvmap_any(ff), jnp.array(False))
  
      assert jnp.array_equal(jit_unvmap_all(tt), jnp.array(True))
      assert jnp.array_equal(jit_unvmap_all(tf), jnp.array(False))
      assert jnp.array_equal(jit_unvmap_all(ff), jnp.array(False))
      assert jnp.array_equal(jit_unvmap_any(tt), jnp.array(True))
      assert jnp.array_equal(jit_unvmap_any(tf), jnp.array(True))
      assert jnp.array_equal(jit_unvmap_any(ff), jnp.array(False))
  
      assert jnp.array_equal(vmap_unvmap_all(tt), jnp.array(True))
      assert jnp.array_equal(vmap_unvmap_all(tf), jnp.array(False))
      assert jnp.array_equal(vmap_unvmap_all(ff), jnp.array(False))
      assert jnp.array_equal(vmap_unvmap_any(tt), jnp.array(True))
      assert jnp.array_equal(vmap_unvmap_any(tf), jnp.array(True))
      assert jnp.array_equal(vmap_unvmap_any(ff), jnp.array(False))
  
      unvmap_max = eqxi.unvmap_max
      jit_unvmap_max = jax.jit(unvmap_max)
      vmap_unvmap_max = jax.vmap(unvmap_max, out_axes=None)
  
      _21 = jnp.array([2, 1])
      _11 = jnp.array([1, 1])
  
      assert jnp.array_equal(unvmap_max(_21), jnp.array(2))
      assert jnp.array_equal(unvmap_max(_11), jnp.array(1))
  
      assert jnp.array_equal(jit_unvmap_max(_21), jnp.array(2))
      assert jnp.array_equal(jit_unvmap_max(_11), jnp.array(1))
  
      assert jnp.array_equal(vmap_unvmap_max(_21), jnp.array(2))
      assert jnp.array_equal(vmap_unvmap_max(_11), jnp.array(1))
  
file_path: equinox/tests/test_misc.py
imports: import equinox.internal
  import jax
  import jax.numpy
  import jax.tree_util
  import pytest
  from .helpers import random_pytree, tree_allclose, treedefs
label_desc: Use image node_id 4 for calling a global function with example usage: test_unvmap() without return types