tensorflow_probability/python/distributions/batch_reshape.py
def calculate_reshape(original_shape, new_shape, validate=False, name=None):
  """Calculates the reshaped dimensions (replacing up to one -1 in reshape)."""
  batch_shape_static = tensorshape_util.constant_value_as_shape(new_shape)
  if tensorshape_util.is_fully_defined(batch_shape_static):
    return np.int32(batch_shape_static), batch_shape_static, []
  with tf.name_scope(name or "calculate_reshape"):
    original_size = tf.reduce_prod(input_tensor=original_shape)
    implicit_dim = tf.equal(new_shape, -1)
    size_implicit_dim = (
        original_size // tf.maximum(1, -tf.reduce_prod(input_tensor=new_shape)))
    new_ndims = tf.shape(input=new_shape)
    expanded_new_shape = tf.where(  # Assumes exactly one `-1`.
        implicit_dim, tf.fill(new_ndims, size_implicit_dim), new_shape)
    validations = [] if not validate else [  # pylint: disable=g-long-ternary
        assert_util.assert_rank(
            original_shape, 1, message="Original shape must be a vector."),
        assert_util.assert_rank(
            new_shape, 1, message="New shape must be a vector."),
        assert_util.assert_less_equal(
            tf.math.count_nonzero(implicit_dim, dtype=tf.int32),
            1,
            message="At most one dimension can be unknown."),
        assert_util.assert_positive(
            expanded_new_shape, message="Shape elements must be >=-1."),
        assert_util.assert_equal(
            tf.reduce_prod(input_tensor=expanded_new_shape),
            original_size,
            message="Shape sizes do not match."),
    ]
    return expanded_new_shape, batch_shape_static, validations
def calculate_reshape(original_shape, new_shape, validate=False, name=None):
  """Calculates the reshaped dimensions (replacing up to one -1 in reshape)."""
  batch_shape_static = tensorshape_util.constant_value_as_shape(new_shape)
  if tensorshape_util.is_fully_defined(batch_shape_static):
    return np.int32(batch_shape_static), batch_shape_static, []
  with tf.name_scope(name or "calculate_reshape"):
    original_size = tf.reduce_prod(input_tensor=original_shape)
    implicit_dim = tf.equal(new_shape, -1)
    size_implicit_dim = (
        original_size // tf.maximum(1, -tf.reduce_prod(input_tensor=new_shape)))
    new_ndims = tf.shape(input=new_shape)
    expanded_new_shape = tf.where(  # Assumes exactly one `-1`.
        implicit_dim, tf.fill(new_ndims, size_implicit_dim), new_shape)
    validations = [] if not validate else [  # pylint: disable=g-long-ternary
        assert_util.assert_rank(
            original_shape, 1, message="Original shape must be a vector."),
        assert_util.assert_rank(
            new_shape, 1, message="New shape must be a vector."),
        assert_util.assert_less_equal(
            tf.math.count_nonzero(implicit_dim, dtype=tf.int32),
            1,
            message="At most one dimension can be unknown."),
        assert_util.assert_positive(
            expanded_new_shape, message="Shape elements must be >=-1."),
        assert_util.assert_equal(
            tf.reduce_prod(input_tensor=expanded_new_shape),
            original_size,
            message="Shape sizes do not match."),
    ]
    return expanded_new_shape, batch_shape_static, validations
Calculates the reshaped dimensions (replacing up to one -1 in reshape).