code: def test_tensorflow_container(batch_axis: int):
      from bentoml._internal.frameworks.tensorflow_v2 import (
          TensorflowTensorContainer,
      )
  
      one_batch: ext.TensorLike = tf.reshape(
          tf.convert_to_tensor(np.arange(6)), (2, 3)
      )
      batch_list: list[ext.TensorLike] = [one_batch, one_batch + 1]
      merged_batch = tf.concat(batch_list, batch_axis)
  
      batches, indices = TensorflowTensorContainer.batches_to_batch(
          batch_list,
          batch_dim=batch_axis,
      )
      assert batches.shape == merged_batch.shape
      assert_tensor_equal(batches, merged_batch)
      assert_tensor_equal(
          TensorflowTensorContainer.batch_to_batches(
              merged_batch,
              indices=indices,
              batch_dim=batch_axis,
          )[0],
          one_batch,
      )
      assert_tensor_equal(
          TensorflowTensorContainer.from_payload(
              TensorflowTensorContainer.to_payload(one_batch)
          ),
          one_batch,
      )
  
      assert_tensor_equal(
          AutoContainer.from_payload(
              AutoContainer.to_payload(one_batch, batch_dim=0)
          ),
          one_batch,
      )
  
file_path: BentoML/tests/integration/frameworks/test_tensorflow_unit.py
imports: from __future__ import annotations
  import typing
  from typing import TYPE_CHECKING
  import numpy
  import pytest
  import tensorflow
  from bentoml._internal.runner.container import AutoContainer
  from bentoml._internal.runner.container import DataContainerRegistry
  from bentoml._internal.types import LazyType
label_desc: Use image node_id 3 for calling a global function with example usage: test_tensorflow_container(batch_axis) without return types