docstr: Tests forward and backward passes with different event shapes and axes.
  
  Args:
    input_shape: Tuple of shapes for input tensor.
    event_dims: Tuple of dimension indices that will be normalized.
    training: Boolean of whether bijector runs in training or inference mode.
docstr_tok: ["Tests","forward","and","backward","passes","with","different","event","shapes","and","axes",".","Args",":","input_shape",":","Tuple","of","shapes","for","input","tensor",".","event_dims",":","Tuple","of","dimension","indices","that","will","be","normalized",".","training",":","Boolean","of","whether","bijector","runs","in","training","or","inference","mode","."]
code: def testForwardInverse(self, input_shape, event_dims, training):
      """Tests forward and backward passes with different event shapes and axes.
  
      Args:
        input_shape: Tuple of shapes for input tensor.
        event_dims: Tuple of dimension indices that will be normalized.
        training: Boolean of whether bijector runs in training or inference mode.
      """
      x_ = np.arange(5 * 4 * 2).astype(np.float32).reshape(input_shape)
      x = tf1.placeholder_with_default(
          x_,
          input_shape if 0 in event_dims else (None,) + input_shape[1:],
      )
      # When training, memorize the exact mean of the last
      # minibatch that it normalized (instead of moving average assignment).
      layer = tf_keras.layers.BatchNormalization(
          axis=event_dims, momentum=0.0, epsilon=0.0
      )
      batch_norm = batch_normalization.BatchNormalization(
          batchnorm_layer=layer, training=training
      )
      # Minibatch statistics are saved only after norm_x has been computed.
      norm_x = batch_norm.inverse(x)
      with tf.control_dependencies(batch_norm.batchnorm.updates):
          moving_mean = tf.identity(batch_norm.batchnorm.moving_mean)
          moving_var = tf.identity(batch_norm.batchnorm.moving_variance)
          denorm_x = batch_norm.forward(tf.identity(norm_x))
          fldj = batch_norm.forward_log_det_jacobian(
              x, event_ndims=len(event_dims)
          )
          # Use identity to invalidate cache.
          ildj = batch_norm.inverse_log_det_jacobian(
              tf.identity(denorm_x), event_ndims=len(event_dims)
          )
      self.evaluate(tf1.global_variables_initializer())
      # Update variables.
      norm_x_ = self.evaluate(norm_x)
      [
          norm_x_,
          moving_mean_,
          moving_var_,
          denorm_x_,
          ildj_,
          fldj_,
      ] = self.evaluate(
          [
              norm_x,
              moving_mean,
              moving_var,
              denorm_x,
              ildj,
              fldj,
          ]
      )
      self.assertStartsWith(batch_norm.name, "batch_normalization")
  
      reduction_axes = self._reduction_axes(input_shape, event_dims)
      keepdims = len(event_dims) > 1
  
      expected_batch_mean = np.mean(
          x_, axis=reduction_axes, keepdims=keepdims
      )
      expected_batch_var = np.var(
          x_, axis=reduction_axes, keepdims=keepdims
      )
  
      if training:
          # When training=True, values become normalized across batch dim and
          # original values are recovered after de-normalizing.
          zeros = np.zeros_like(norm_x_)
          self.assertAllClose(
              np.mean(zeros, axis=reduction_axes),
              np.mean(norm_x_, axis=reduction_axes),
          )
  
          self.assertAllClose(expected_batch_mean, moving_mean_)
          self.assertAllClose(expected_batch_var, moving_var_)
          self.assertAllClose(x_, denorm_x_, atol=1e-5)
          # Since moving statistics are set to batch statistics after
          # normalization, ildj and -fldj should match.
          self.assertAllClose(ildj_, -fldj_)
          # ildj is computed with minibatch statistics.
          expected_ildj = np.sum(
              np.log(1.0)
              - 0.5
              * np.log(
                  expected_batch_var + batch_norm.batchnorm.epsilon
              )
          )
          self.assertAllClose(expected_ildj, np.squeeze(ildj_))
      else:
          # When training=False, moving_mean, moving_var remain at their
          # initialized values (0., 1.), resulting in no scale/shift (a small
          # shift occurs if epsilon > 0.)
          self.assertAllClose(x_, norm_x_)
          self.assertAllClose(x_, denorm_x_, atol=1e-5)
          # ildj is computed with saved statistics.
          expected_ildj = np.sum(
              np.log(1.0)
              - 0.5 * np.log(1.0 + batch_norm.batchnorm.epsilon)
          )
          self.assertAllClose(expected_ildj, np.squeeze(ildj_))
  
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 2 for calling the BatchNormTest obj's underlying member method code with example usage: obj.testForwardInverse(input_shape, event_dims, training) without return types