code: def batches_to_batch(
      cls,
      batches: t.Sequence[jax.Array],
      batch_dim: int = 0,
  ) -> tuple[jax.Array, list[int]]:
      batch: jax.Array = jnp.concatenate(batches, axis=batch_dim)
      indices: list[int] = list(
          itertools.accumulate(
              subbatch.shape[0] for subbatch in batches
          )
      )
      indices = [0] + indices
      return batch, indices
  
file_path: BentoML/src/bentoml/_internal/frameworks/common/jax.py
imports: from __future__ import annotations
  import itertools
  import pickle
  import typing
  from typing import TYPE_CHECKING
  from simple_di import inject
  from ....exceptions import MissingDependencyException
  from ...runner.container import DataContainer
  from ...runner.container import DataContainerRegistry
  from ...runner.container import Payload
  from ...types import LazyType
  from ...utils import LazyLoader
label_desc: Use image node_id 1 for calling the JaxArrayContainer obj's underlying member method code with example usage: obj.batches_to_batch(cls, batches, batch_dim) and returns: batch,  indices