code: def test_methods(call, outer):
      num_traces = 0
  
      class M(eqx.Module):
          increment: Union[int, jax.Array]
  
          if call:
  
              def __call__(self, x):
                  nonlocal num_traces
                  num_traces += 1
                  return x + self.increment
  
              if not outer:
                  __call__ = filter_pmap(__call__)
          else:
  
              def method(self, x):
                  nonlocal num_traces
                  num_traces += 1
                  return x + self.increment
  
              if not outer:
                  method = filter_pmap(method)
  
      y = jnp.array([1])
  
      def run(_m):
          if call:
              if outer:
                  return filter_pmap(_m)(y)
              else:
                  return _m(y)
          else:
              if outer:
                  return filter_pmap(_m.method)(y)
              else:
                  return _m.method(y)
  
      m = M(1)
      assert shaped_allclose(run(m), jnp.array([2]))
      assert num_traces == 2
      assert shaped_allclose(run(m), jnp.array([2]))
      assert num_traces == 2
      n = M(2)
      assert shaped_allclose(run(n), jnp.array([3]))
      assert num_traces == 4
      assert shaped_allclose(run(n), jnp.array([3]))
      assert num_traces == 4
      o = M(jnp.array([5]))
      p = M(jnp.array([6]))
      if outer:
          assert shaped_allclose(run(o), jnp.array([[6]]))
          assert num_traces == 6
          assert shaped_allclose(run(p), jnp.array([[7]]))
      else:
          assert shaped_allclose(run(o), jnp.array([6]))
          assert num_traces == 6
          assert shaped_allclose(run(p), jnp.array([7]))
      assert num_traces == 6
  
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 8 for calling a global function with example usage: test_methods(call, outer) and returns: unknown, unknown, filter_pmap, _m, filter_pmap, _m