code: def test_ndarray_container(batch_dim: int):
      arr1 = np.ones((3, 3))
      if batch_dim == 0:
          arr2 = np.arange(6).reshape(2, 3)
      else:
          arr2 = np.arange(6).reshape(3, 2)
  
      batches = [arr1, arr2]
      batch, indices = c.NdarrayContainer.batches_to_batch(
          batches, batch_dim=batch_dim
      )
      assert (batch == np.concatenate(batches, axis=batch_dim)).all()
      (
          restored_arr1,
          restored_arr2,
      ) = c.NdarrayContainer.batch_to_batches(
          batch, indices, batch_dim=batch_dim
      )
      assert (arr1 == restored_arr1).all()
      assert (arr2 == restored_arr2).all()
  
      assert (
          c.NdarrayContainer.from_payload(
              c.NdarrayContainer.to_payload(arr1, batch_dim)
          )
          == arr1
      ).all()
  
      (
          restored_batch,
          restored_indices,
      ) = c.NdarrayContainer.from_batch_payloads(
          c.NdarrayContainer.batch_to_payloads(
              batch, indices, batch_dim=batch_dim
          ),
          batch_dim=batch_dim,
      )
      assert restored_indices == indices
      assert (restored_batch == batch).all()
  
file_path: BentoML/tests/unit/_internal/runner/test_container.py
imports: from __future__ import annotations
  import typing
  import numpy
  import pandas
  import pytest
  import bentoml._internal.runner.container
label_desc: Use image node_id 2 for calling a global function with example usage: test_ndarray_container(batch_dim) without return types