code: def test_batch_shape_inference_is_correct(
      self, value_fn, expected_batch_shape_parts, expected_batch_shape
  ):
      value = (
          value_fn()
      )  # Defer construction until we're in the right graph.
  
      parts = batch_shape_lib.batch_shape_parts(value)
      self.assertAllEqualNested(
          parts,
          nest.map_structure_up_to(
              parts, tf.TensorShape, expected_batch_shape_parts
          ),
      )
  
      self.assertAllEqual(
          expected_batch_shape,
          batch_shape_lib.inferred_batch_shape_tensor(value),
      )
  
      batch_shape = batch_shape_lib.inferred_batch_shape(value)
      self.assertIsInstance(batch_shape, tf.TensorShape)
      self.assertTrue(
          batch_shape.is_compatible_with(expected_batch_shape)
      )
  
imports: from absl import logging
  from absl.testing import parameterized
  import mock
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import exp
  from tensorflow_probability.python.bijectors import joint_map
  from tensorflow_probability.python.bijectors import sigmoid
  from tensorflow_probability.python.distributions import categorical
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import mixture_same_family
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import batch_shape_lib
  from tensorflow_probability.python.internal import parameter_properties
  from tensorflow_probability.python.internal import test_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 1 for calling the BatchShapeInferenceTests obj's underlying member method code with example usage: obj.test_batch_shape_inference_is_correct(value_fn, expected_batch_shape_parts, expected_batch_shape) without return types