tensorflow_probability/python/internal/distribution_util.py
def dimension_size(x, axis):
  """Returns the size of a specific dimension."""
  # Since tf.gather isn't "constant-in, constant-out", we must first check the
  # static shape or fallback to dynamic shape.
  s = tf.compat.dimension_value(
      tensorshape_util.with_rank_at_least(x.shape, np.abs(axis))[axis])
  if s is not None:
    return s
  return tf.shape(input=x)[axis]
def dimension_size(x, axis):
  """Returns the size of a specific dimension."""
  # Since tf.gather isn't "constant-in, constant-out", we must first check the
  # static shape or fallback to dynamic shape.
  s = tf.compat.dimension_value(
      tensorshape_util.with_rank_at_least(x.shape, np.abs(axis))[axis])
  if s is not None:
    return s
  return tf.shape(input=x)[axis]
Returns the size of a specific dimension.