code: def test_pmap_grad():
      num_traces = 0
  
      def f(x):
          nonlocal num_traces
          num_traces += 1
          return x + 1
  
      grad = filter_pmap(eqx.filter_grad(f))(jnp.array([1.0]))
      assert shaped_allclose(grad, jnp.array([1.0]))
      assert num_traces == 2  # eval_shape + pmap
  
      grad = filter_pmap(eqx.filter_grad(f))(jnp.array([2.0]))
      assert shaped_allclose(grad, jnp.array([1.0]))
      assert num_traces == 2  # (eval_shape cached, pmap cached)
  
      value, grad = filter_pmap(eqx.filter_value_and_grad(f))(
          jnp.array([1.0])
      )
      assert shaped_allclose(value, jnp.array([2.0]))
      assert shaped_allclose(grad, jnp.array([1.0]))
      assert num_traces == 4  # eval_shape + pmap
  
      value, grad = filter_pmap(eqx.filter_value_and_grad(f))(
          jnp.array([2.0])
      )
      assert shaped_allclose(value, jnp.array([3.0]))
      assert shaped_allclose(grad, jnp.array([1.0]))
      assert num_traces == 4
  
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 9 for calling a global function with example usage: test_pmap_grad() and returns: unknown