code: def test_double_if_mapped():
      out_axes = eqx.internal.if_mapped(1)
  
      def f(x):
          assert x.shape == (3, 1)
  
          def g(y):
              assert y.shape == (1,)
              return y + 1, x + 1
  
          a, b = eqx.filter_vmap(g, out_axes=out_axes)(x)
          assert a.shape == (1, 3)
          assert b.shape == (3, 1)
  
      filter_pmap(f)(jnp.arange(3).reshape(1, 3, 1))
  
imports: import functools
  from typing import Any, Union
  import equinox
  import jax
  import jax.numpy
  import jax.random
  import jax.tree_util
  import pytest
  from .helpers import tree_allclose
label_desc: Use image node_id 16 for calling a global function with example usage: test_double_if_mapped() and returns: unknown,  unknown