code: def get_batch_shape_part(x, event_ndims):
      """Extracts an object's shape for batch shape inference."""
      if hasattr(x, "experimental_batch_shape"):  # `x` is a bijector
          if event_ndims is None:
              # If `event_ndims` is None, then we don't know the rank of the event for
              # inputs to `x`, so we cannot determine the batch shape of bijector `x`.
              return tf.TensorShape(None)
          try:
              return x.experimental_batch_shape(
                  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 tf.TensorShape([])
      if hasattr(
          x, "batch_shape"
      ):  # `x` is a distribution or linear operator.
          base_shape = x.batch_shape
          # Hack to attempt to collapse non-autobatched JDs to a coherent batch shape.
          # Distrax distributions have tuple-valued batch shapes, which are
          # always spuriously 'nested', so to avoid confusion we check that the
          # distribution is actually a joint distribution (Distrax
          # distributions are always single-part).
          if tf.nest.is_nested(x.dtype):
              base_shape = functools.reduce(
                  tf.broadcast_static_shape, tf.nest.flatten(base_shape)
              )
          base_shape = tf.TensorShape(base_shape)
      elif hasattr(x, "shape"):  # `x` is a Tensor or ndarray.
          base_shape = tf.TensorShape(x.shape)
          # `x` is a Python list, tuple, or literal.
      else:
          base_shape = tf.TensorShape(np.array(x).shape)
      return _truncate_shape(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 4 for calling a global function with example usage: get_batch_shape_part(x, event_ndims) and returns: _truncate_shape, tf, x, tf