tensorflow_probability/python/distributions/finite_discrete.py
def _maybe_validate_args(outcomes, logits, probs, validate_args):
  """Validate `outcomes`, `logits` and `probs`'s shapes."""
  assertions = []

  def validate_equal_last_dim(tensor_a, tensor_b, message):
    if tensor_a.shape.is_fully_defined() and tensor_b.shape.is_fully_defined():
      if tensor_a.shape[-1] != tensor_b.shape[-1]:
        raise ValueError(message)
    elif validate_args:
      assertions.append(
          tf.compat.v1.assert_equal(
              tf.shape(input=tensor_a)[-1],
              tf.shape(input=tensor_b)[-1],
              message=message))

  if logits is not None:
    validate_equal_last_dim(
        outcomes,
        logits,
        message='Last dimension of outcomes and logits must be equal size.')
  if probs is not None:
    validate_equal_last_dim(
        outcomes,
        probs,
        message='Last dimension of outcomes and probs must be equal size.')

  message = 'Rank of outcomes must be 1.'
  if outcomes.shape.ndims is not None:
    if outcomes.shape.ndims != 1:
      raise ValueError(message)
  elif validate_args:
    assertions.append(tf.compat.v1.assert_rank(outcomes, 1, message=message))

  message = 'Size of outcomes must be greater than 0.'
  if outcomes.shape.num_elements() is not None:
    if outcomes.shape.num_elements() == 0:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        tf.compat.v1.assert_greater(
            tf.size(input=outcomes), 0, message=message))

  if validate_args:
    assertions.append(
        tf.compat.v1.assert_equal(
            tf.math.is_strictly_increasing(outcomes),
            True,
            message='outcomes is not strictly increasing.'))

  return assertions
def _maybe_validate_args(outcomes, logits, probs, validate_args):
  """Validate `outcomes`, `logits` and `probs`'s shapes."""
  assertions = []

  def validate_equal_last_dim(tensor_a, tensor_b, message):
    if tensor_a.shape.is_fully_defined() and tensor_b.shape.is_fully_defined():
      if tensor_a.shape[-1] != tensor_b.shape[-1]:
        raise ValueError(message)
    elif validate_args:
      assertions.append(
          tf.compat.v1.assert_equal(
              tf.shape(input=tensor_a)[-1],
              tf.shape(input=tensor_b)[-1],
              message=message))

  if logits is not None:
    validate_equal_last_dim(
        outcomes,
        logits,
        message='Last dimension of outcomes and logits must be equal size.')
  if probs is not None:
    validate_equal_last_dim(
        outcomes,
        probs,
        message='Last dimension of outcomes and probs must be equal size.')

  message = 'Rank of outcomes must be 1.'
  if outcomes.shape.ndims is not None:
    if outcomes.shape.ndims != 1:
      raise ValueError(message)
  elif validate_args:
    assertions.append(tf.compat.v1.assert_rank(outcomes, 1, message=message))

  message = 'Size of outcomes must be greater than 0.'
  if outcomes.shape.num_elements() is not None:
    if outcomes.shape.num_elements() == 0:
      raise ValueError(message)
  elif validate_args:
    assertions.append(
        tf.compat.v1.assert_greater(
            tf.size(input=outcomes), 0, message=message))

  if validate_args:
    assertions.append(
        tf.compat.v1.assert_equal(
            tf.math.is_strictly_increasing(outcomes),
            True,
            message='outcomes is not strictly increasing.'))

  return assertions
Validate `outcomes`, `logits` and `probs`'s shapes.