code: def testSquashedOuterWrapperSimple(self):
      bn = tf.keras.layers.BatchNormalization(axis=-1)
      layer = squashed_outer_wrapper.SquashedOuterWrapper(
          bn, inner_rank=3
      )
  
      inputs_flat = tf.range(3 * 4 * 5 * 6 * 7, dtype=tf.float32)
      inputs_2_batch = tf.reshape(inputs_flat, [3, 4, 5, 6, 7])
      outputs_2_batch = layer(inputs_2_batch)
  
      inputs_1_batch = tf.reshape(inputs_flat, [3 * 4, 5, 6, 7])
      outputs_1_batch = layer(inputs_1_batch)
      outputs_1_batch_reshaped = tf.reshape(
          outputs_1_batch, [3, 4, 5, 6, 7]
      )
  
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.assertAllClose(
          self.evaluate(outputs_2_batch),
          self.evaluate(outputs_1_batch_reshaped),
      )
  
file_path: agents/tf_agents/keras_layers/squashed_outer_wrapper_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tf_agents.keras_layers import squashed_outer_wrapper
  from tf_agents.utils import common
  from tf_agents.utils import test_utils
label_desc: Use image node_id 2 for calling the SquashedOuterWrapperTest obj's underlying member method code with example usage: obj.testSquashedOuterWrapperSimple() without return types