code: def _CheckBatching(
      self,
      op,
      bdim_size,
      bdims,
      shapes,
      dtypes,
      rng,
      rtol=None,
      atol=None,
      multiple_results=False,
  ):
      batched_shapes = map(
          functools.partial(lax_test_util.add_bdim, bdim_size),
          bdims,
          shapes,
      )
      args = [
          rng(shape, dtype)
          for shape, dtype in zip(batched_shapes, dtypes)
      ]
      args_slice = lax_test_util.args_slicer(args, bdims)
      ans = jax.vmap(op, bdims)(*args)
      if bdim_size == 0:
          args = [
              rng(shape, dtype) for shape, dtype in zip(shapes, dtypes)
          ]
          out = op(*args)
          if not multiple_results:
              expected = np.zeros((0,) + out.shape, out.dtype)
          else:
              expected = [
                  np.zeros((0,) + o.shape, o.dtype) for o in out
              ]
      else:
          outs = [op(*args_slice(i)) for i in range(bdim_size)]
          if not multiple_results:
              expected = np.stack(outs)
          else:
              expected = [np.stack(xs) for xs in zip(*outs)]
      self.assertAllClose(ans, expected, rtol=rtol, atol=atol)
  
file_path: jax/tests/lax_vmap_op_test.py
imports: import functools
  import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  import numpy
  import jax
  from jax import lax
  from jax._src import test_util
  from jax._src.internal_test_util import lax_test_util
  from jax._src import util
  from jax import config
label_desc: Use image node_id 1 for calling the LaxVmapOpTest obj's underlying member method code with example usage: obj._CheckBatching(op, bdim_size, bdims, shapes, dtypes, rng, rtol, atol, multiple_results) without return types