class_docstr: Temporarily merges leading dimensions of input tensors.
  
  Merges the leading dimensions of a tensor into a single dimension, runs the
  given callable, then splits the leading dimension of the result to match the
  input.
  
  Input arrays whose rank is smaller than the number of dimensions to collapse
  are passed unmodified.
  
  This may be useful for applying a module to each timestep of e.g. a
  ``[Time, Batch, ...]`` array.
  
  For some ``f``\ s and platforms, this may be more efficient than
  :func:`jax.vmap`, especially when combined with other transformations like
  :func:`jax.grad`.
docstr: Constructs a :class:`BatchApply` module.
  
  Args:
    f: The callable to be applied to the reshaped array.
    num_dims: The number of dimensions to merge.
docstr_tok: ["Constructs","a",":","class",":","`","BatchApply","`","module",".","Args",":","f",":","The","callable","to","be","applied","to","the","reshaped","array",".","num_dims",":","The","number","of","dimensions","to","merge","."]
code: def __init__(self, f, num_dims=2):
      """Constructs a :class:`BatchApply` module.
  
      Args:
        f: The callable to be applied to the reshaped array.
        num_dims: The number of dimensions to merge.
      """
      self._f = f
      self.num_dims = num_dims
  
imports: from collections.abc import Iterable, Sequence
  import functools
  from typing import Any, Callable, Optional
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import typing
  from haiku._src.typing import PRNGKey
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 to create a new BatchApply object  with example: obj = BatchApply(f, num_dims)