code: def test_correct_output_shape(
      layer_class: Type[tf.keras.layers.Layer],
      model_config: Dict[Text, Any],
      layer_args: Dict[Text, Any],
      layer_inputs: List[List[Union[tf.SparseTensor, tf.Tensor]]],
      expected_output_shapes_train: List[List[int]],
      expected_output_shapes_test: Union[Text, List[List[int]]],
  ) -> None:
      layer = layer_class(
          **layer_args, attribute=attribute_name, config=model_config
      )
  
      train_outputs = layer(layer_inputs, training=True)
      if not isinstance(train_outputs, tuple):
          train_outputs = [train_outputs]
      for i, expected_shape in enumerate(expected_output_shapes_train):
          assert train_outputs[i].shape == expected_shape
  
      if expected_output_shapes_test == "same_as_train":
          expected_output_shapes_test = expected_output_shapes_train
      test_outputs = layer(layer_inputs, training=False)
      if not isinstance(test_outputs, tuple):
          test_outputs = [test_outputs]
      for i, expected_shape in enumerate(expected_output_shapes_test):
          assert test_outputs[i].shape == expected_shape
  
imports: import pytest
  import tensorflow
  import numpy
  from typing import Text, Union, Any, Dict, List, Type
  from rasa.shared.nlu.constants import TEXT, FEATURE_TYPE_SENTENCE, FEATURE_TYPE_SEQUENCE
  from rasa.utils.tensorflow import layers
  from rasa.utils.tensorflow.rasa_layers import ConcatenateSparseDenseFeatures, RasaFeatureCombiningLayer, RasaSequenceLayer, RasaCustomLayer
  from rasa.utils.tensorflow.constants import DENSE_INPUT_DROPOUT, SPARSE_INPUT_DROPOUT, DROP_RATE, DENSE_DIMENSION, REGULARIZATION_CONSTANT, CONCAT_DIMENSION, CONNECTION_DENSITY, DROP_RATE_ATTENTION, KEY_RELATIVE_ATTENTION, VALUE_RELATIVE_ATTENTION, MAX_RELATIVE_POSITION, UNIDIRECTIONAL_ENCODER, HIDDEN_LAYERS_SIZES, NUM_TRANSFORMER_LAYERS, TRANSFORMER_SIZE, NUM_HEADS, SENTENCE, SEQUENCE, MASKED_LM, LABEL
  from rasa.utils.tensorflow.exceptions import TFLayerConfigException
  from rasa.utils.tensorflow.model_data import FeatureSignature
label_desc: Use image node_id 2 for calling a global function with example usage: test_correct_output_shape(layer_class, model_config, layer_args, layer_inputs, expected_output_shapes_train, expected_output_shapes_test) without return types