docstr: This tests if the model runs with different batch shapes.
docstr_tok: ["This","tests","if","the","model","runs","with","different","batch","shapes","."]
code: def testMultiDimensionalInput(self):
      """This tests if the model runs with different batch shapes."""
  
      def dummy_glow_net(input_shape):
          del input_shape
          return lambda x: identity.Identity()
  
      def dummy_exit_net(input_shape, output_chan):
          del input_shape, output_chan
          return lambda x: identity.Identity()
  
      bijection = glow.Glow(
          output_shape=self.output_shape,
          num_glow_blocks=2,
          num_steps_per_block=1,
          coupling_bijector_fn=dummy_glow_net,
          exit_bijector_fn=dummy_exit_net,
          grab_after_block=[0.5, 0.5],
          use_actnorm=False,
      )
  
      single_z = bijection.inverse(tf.zeros([16, 16, 3]))
      single_z2 = bijection.inverse(tf.zeros([1, 16, 16, 3]))
      single_z3 = bijection.inverse(tf.zeros([1, 1, 16, 16, 3]))
      self.assertShapeEqual(np.zeros([16 * 16 * 3]), single_z)
      self.assertShapeEqual(np.zeros([1, 16 * 16 * 3]), single_z2)
      self.assertShapeEqual(np.zeros([1, 1, 16 * 16 * 3]), single_z3)
  
imports: import functools
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import bijector_test_util
  from tensorflow_probability.python.bijectors import blockwise
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import composition
  from tensorflow_probability.python.bijectors import glow
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import pad
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.math.gradient import batch_jacobian
label_desc: Use image node_id 12 for calling the GlowTest obj's underlying member method code with example usage: obj.testMultiDimensionalInput() and returns: function, function