code: def get_batch_shape_tensor_part(x, event_ndims):
      """Extracts an object's runtime (Tensor) shape for batch shape inference."""
      if hasattr(
          x, "experimental_batch_shape_tensor"
      ):  # `x` is a bijector
          try:
              return x.experimental_batch_shape_tensor(
                  x_event_ndims=event_ndims
              )
          except NotImplementedError:
              # Backwards compatibility with bijector-like instances that don't
              # inherit from tfb.Bijector (e.g., Distrax bijectors) and/or don't
              # implement `_parameter_properties`.
              return []
      if hasattr(
          x, "batch_shape_tensor"
      ):  # `x` is a distribution or linop.
          base_shape = x.batch_shape_tensor()
          if tf.nest.is_nested(base_shape):
              # Attempt to collapse non-autobatched JDs to a coherent batch shape.
              base_shape = functools.reduce(
                  ps.broadcast_shape, tf.nest.flatten(base_shape)
              )
      elif hasattr(x, "shape") and tensorshape_util.is_fully_defined(
          x.shape
      ):
          base_shape = x.shape
      else:
          base_shape = tf.shape(x)
      return _truncate_shape_tensor(base_shape, event_ndims)
  
file_path: probability/tensorflow_probability/python/internal/batch_shape_lib.py
imports: import functools
  from absl import logging
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow.python.util import nest
label_desc: Use image node_id 3 for calling a global function with example usage: get_batch_shape_tensor_part(x, event_ndims) and returns: _truncate_shape_tensor, x, list