code: def testLogProb(self, event_shape, event_dims, training, layer_cls):
      training = tf1.placeholder_with_default(training, (), "training")
      layer = layer_cls(axis=event_dims, epsilon=0.0)
      batch_norm = batch_normalization.BatchNormalization(
          batchnorm_layer=layer, training=training
      )
      base_dist = mvn_diag.MultivariateNormalDiag(
          loc=np.zeros(np.prod(event_shape), dtype=np.float32)
      )
      # Reshape the events.
      if isinstance(event_shape, int):
          event_shape = [event_shape]
      base_dist = transformed_distribution.TransformedDistribution(
          distribution=base_dist,
          bijector=reshape.Reshape(event_shape_out=event_shape),
      )
      dist = transformed_distribution.TransformedDistribution(
          distribution=base_dist,
          bijector=batch_norm,
          validate_args=True,
      )
      samples = dist.sample(int(1e5))
      # No volume distortion since training=False, bijector is initialized
      # to the identity transformation.
      base_log_prob = base_dist.log_prob(samples)
      dist_log_prob = dist.log_prob(samples)
      self.evaluate(tf1.global_variables_initializer())
      base_log_prob_, dist_log_prob_ = self.evaluate(
          [base_log_prob, dist_log_prob]
      )
      self.assertAllClose(base_log_prob_, dist_log_prob_)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import batch_normalization
  from tensorflow_probability.python.bijectors import bijector_test_util
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import reshape
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
label_desc: Use image node_id 3 for calling the BatchNormTest obj's underlying member method code with example usage: obj.testLogProb(event_shape, event_dims, training, layer_cls) without return types