code: def test_batch_inference(
      batch_size: int,
      number_of_data_points: int,
      expected_number_of_batch_iterations: int,
  ):
      model = RasaModel()
  
      def _batch_predict(
          batch_in: Tuple[np.ndarray],
      ) -> Dict[Text, Union[np.ndarray, Dict[Text, np.ndarray]]]:
          dummy_output = batch_in[0]
          output = {
              "dummy_output": dummy_output,
              "non_input_affected_output": tf.constant(
                  np.array([[1, 2]]), dtype=tf.int32
              ),
          }
          return output
  
      # Monkeypatch batch predict so that run_inference interface can be tested
      model.batch_predict = _batch_predict
  
      # Create dummy model data to pass to model
      model_data = RasaModelData(
          label_key=LABEL,
          label_sub_key=IDS,
          data={
              TEXT: {
                  SENTENCE: [
                      FeatureArray(
                          np.random.rand(number_of_data_points, 2),
                          number_of_dimensions=2,
                      )
                  ]
              }
          },
      )
      output = model.run_inference(model_data, batch_size=batch_size)
  
      # Firstly, the number of data points in dummy_output should be equal
      # to the number of data points sent as input.
      assert output["dummy_output"].shape[0] == number_of_data_points
  
      # Secondly, the number of data points inside diagnostic_data should be
      # equal to the number of batches passed to the model because for every
      # batch passed as input, it would have created a
      # corresponding diagnostic data entry.
      assert output["non_input_affected_output"].shape == (
          expected_number_of_batch_iterations,
          2,
      )
  
file_path: rasa/tests/utils/tensorflow/test_models.py
imports: import pytest
  from typing import Dict, Text, Union, Tuple, List
  import numpy
  import tensorflow
  from rasa.utils.tensorflow.models import RasaModel, TransformerRasaModel
  from rasa.utils.tensorflow.model_data import RasaModelData
  from rasa.utils.tensorflow.model_data import FeatureArray
  from rasa.utils.tensorflow.constants import LABEL, IDS, SENTENCE
  from rasa.shared.nlu.constants import TEXT, FEATURE_TYPE_SENTENCE, FEATURE_TYPE_SEQUENCE
label_desc: Use image node_id 2 for calling a global function with example usage: test_batch_inference(batch_size, number_of_data_points, expected_number_of_batch_iterations) and returns: output