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
def calculate_reshape ( original_shape , new_shape , validate = False , name = None ) : 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).
Calculates the reshaped dimensions ( replacing up to one - 1 in reshape ) .