def _verify_input(tensor_list, labels, probs_list):
  """Verify that batched inputs are well-formed."""
  checked_probs_list = []
  for probs in probs_list:
    # Since number of classes shouldn't change at runtime, probabilities shape
    # should be fully defined.
    probs.get_shape().assert_is_fully_defined()

    # Probabilities must be 1D.
    probs.get_shape().assert_has_rank(1)

    # Probabilities must be nonnegative and sum to one.
    tol = 1e-6
    prob_sum = math_ops.reduce_sum(probs)
    checked_probs = control_flow_ops.with_dependencies([
        check_ops.assert_non_negative(probs),
        check_ops.assert_less(prob_sum, 1.0 + tol),
        check_ops.assert_less(1.0 - tol, prob_sum)
    ], probs)
    checked_probs_list.append(checked_probs)

  # All probabilities should be the same length.
  prob_length = checked_probs_list[0].get_shape().num_elements()
  for checked_prob in checked_probs_list:
    if checked_prob.get_shape().num_elements() != prob_length:
      raise ValueError('Probability parameters must have the same length.')

  # Labels tensor should only have batch dimension.
  labels.get_shape().assert_has_rank(1)

  for tensor in tensor_list:
    # Data tensor should have a batch dimension.
    shape = tensor.get_shape().with_rank_at_least(1)

    # Data and label batch dimensions must be compatible.
    tensor_shape.dimension_at_index(shape, 0).assert_is_compatible_with(
        labels.get_shape()[0])

  # Data and labels must have the same, strictly positive batch size. Since we
  # can't assume we know the batch size at graph creation, add runtime checks.
  labels_batch_size = array_ops.shape(labels)[0]
  lbl_assert = check_ops.assert_positive(labels_batch_size)

  # Make each tensor depend on its own checks.
  labels = control_flow_ops.with_dependencies([lbl_assert], labels)
  tensor_list = [
      control_flow_ops.with_dependencies([
          lbl_assert,
          check_ops.assert_equal(array_ops.shape(x)[0], labels_batch_size)
      ], x) for x in tensor_list
  ]

  # Label's classes must be integers 0 <= x < num_classes.
  labels = control_flow_ops.with_dependencies([
      check_ops.assert_integer(labels), check_ops.assert_non_negative(labels),
      check_ops.assert_less(labels, math_ops.cast(prob_length, labels.dtype))
  ], labels)

  return tensor_list, labels, checked_probs_list

  def _check_domain_range_possibly_add_asserts(self):
    """Static check of init arg `num_rows`, possibly add asserts."""
    # Possibly add asserts.
    if self._assert_proper_shapes:
      self._num_rows = control_flow_ops.with_dependencies([
          check_ops.assert_rank(
              self._num_rows,
              0,
              message="Argument num_rows must be a 0-D Tensor."),
          check_ops.assert_non_negative(
              self._num_rows,
              message="Argument num_rows must be non-negative."),
      ], self._num_rows)
      self._num_columns = control_flow_ops.with_dependencies([
          check_ops.assert_rank(
              self._num_columns,
              0,
              message="Argument num_columns must be a 0-D Tensor."),
          check_ops.assert_non_negative(
              self._num_columns,
              message="Argument num_columns must be non-negative."),
      ], self._num_columns)

    # Static checks.
    if not self._num_rows.dtype.is_integer:
      raise TypeError("Argument num_rows must be integer type.  Found:"
                      " %s" % self._num_rows)

    if not self._num_columns.dtype.is_integer:
      raise TypeError("Argument num_columns must be integer type.  Found:"
                      " %s" % self._num_columns)

    num_rows_static = self._num_rows_static
    num_columns_static = self._num_columns_static

    if num_rows_static is not None:
      if num_rows_static.ndim != 0:
        raise ValueError("Argument num_rows must be a 0-D Tensor.  Found:"
                         " %s" % num_rows_static)

      if num_rows_static < 0:
        raise ValueError("Argument num_rows must be non-negative.  Found:"
                         " %s" % num_rows_static)
    if num_columns_static is not None:
      if num_columns_static.ndim != 0:
        raise ValueError("Argument num_columns must be a 0-D Tensor.  Found:"
                         " %s" % num_columns_static)

      if num_columns_static < 0:
        raise ValueError("Argument num_columns must be non-negative.  Found:"
                         " %s" % num_columns_static)

 def test_raises_when_negative(self):
   with self.test_session():
     zoe = constant_op.constant([-1, -2], name="zoe")
     with ops.control_dependencies([check_ops.assert_non_negative(zoe)]):
       out = array_ops.identity(zoe)
     with self.assertRaisesOpError("zoe"):
       out.eval()

  def _check_batch_shape_possibly_add_asserts(self):
    """Static check of init arg `batch_shape`, possibly add asserts."""
    if self._batch_shape_arg is None:
      return

    # Possibly add asserts
    if self._assert_proper_shapes:
      self._batch_shape_arg = control_flow_ops.with_dependencies(
          [
              check_ops.assert_rank(
                  self._batch_shape_arg,
                  1,
                  message="Argument batch_shape must be a 1-D Tensor."),
              check_ops.assert_non_negative(
                  self._batch_shape_arg,
                  message="Argument batch_shape must be non-negative."),
          ],
          self._batch_shape_arg)

    # Static checks
    if not self._batch_shape_arg.dtype.is_integer:
      raise TypeError("Argument batch_shape must be integer type.  Found:"
                      " %s" % self._batch_shape_arg)

    if self._batch_shape_static is None:
      return  # Cannot do any other static checks.

    if self._batch_shape_static.ndim != 1:
      raise ValueError("Argument batch_shape must be a 1-D Tensor.  Found:"
                       " %s" % self._batch_shape_static)

    if np.any(self._batch_shape_static < 0):
      raise ValueError("Argument batch_shape must be non-negative.  Found:"
                       "%s" % self._batch_shape_static)

 def _assert_valid_n(self, n, validate_args):
   n = ops.convert_to_tensor(n, name="n")
   if not validate_args:
     return n
   return control_flow_ops.with_dependencies(
       [check_ops.assert_non_negative(n),
        distribution_util.assert_integer_form(n)], n)

  def get_sample_ndims(self, x, name="get_sample_ndims"):
    """Returns number of dimensions corresponding to iid draws ("sample").

    Args:
      x: `Tensor`.
      name: `String`. The name to give this op.

    Returns:
      sample_ndims: `Tensor` (0D, `int32`).

    Raises:
      ValueError: if `sample_ndims` is calculated to be negative.
    """
    with self._name_scope(name, values=[x]):
      ndims = self.get_ndims(x, name=name)
      if self._is_all_constant_helper(ndims, self.batch_ndims,
                                      self.event_ndims):
        ndims = tensor_util.constant_value(ndims)
        sample_ndims = (ndims - self._batch_ndims_static -
                        self._event_ndims_static)
        if sample_ndims < 0:
          raise ValueError(
              "expected batch_ndims(%d) + event_ndims(%d) <= ndims(%d)" %
              (self._batch_ndims_static, self._event_ndims_static, ndims))
        return ops.convert_to_tensor(sample_ndims, name="sample_ndims")
      else:
        with ops.name_scope(name="sample_ndims"):
          sample_ndims = ndims - self.batch_ndims - self.event_ndims
          if self.validate_args:
            sample_ndims = control_flow_ops.with_dependencies(
                [check_ops.assert_non_negative(sample_ndims)], sample_ndims)
        return sample_ndims

 def _maybe_assert_valid_x(self, x):
   if not self.validate_args:
     return x
   is_valid = check_ops.assert_non_negative(
       x,
       message="Forward transformation input must be at least {}.".format(0))
   return control_flow_ops.with_dependencies([is_valid], x)

  def __init__(self,
               n,
               logits=None,
               p=None,
               validate_args=False,
               allow_nan_stats=True,
               name="Binomial"):
    """Initialize a batch of Binomial distributions.

    Args:
      n:  Non-negative floating point tensor with shape broadcastable to
        `[N1,..., Nm]` with `m >= 0` and the same dtype as `p` or `logits`.
        Defines this as a batch of `N1 x ... x Nm` different Binomial
        distributions. Its components should be equal to integer values.
      logits: Floating point tensor representing the log-odds of a
        positive event with shape broadcastable to `[N1,..., Nm]` `m >= 0`, and
        the same dtype as `n`. Each entry represents logits for the probability
        of success for independent Binomial distributions.
      p:  Positive floating point tensor with shape broadcastable to
        `[N1,..., Nm]` `m >= 0`, `p in [0, 1]`. Each entry represents the
        probability of success for independent Binomial distributions.
      validate_args: `Boolean`, default `False`.  Whether to assert valid values
        for parameters `n`, `p`, and `x` in `prob` and `log_prob`.
        If `False` and inputs are invalid, correct behavior is not guaranteed.
      allow_nan_stats: `Boolean`, default `True`.  If `False`, raise an
        exception if a statistic (e.g. mean/mode/etc...) is undefined for any
        batch member.  If `True`, batch members with valid parameters leading to
        undefined statistics will return NaN for this statistic.
      name: The name to prefix Ops created by this distribution class.

    Examples:

    ```python
    # Define 1-batch of a binomial distribution.
    dist = Binomial(n=2., p=.9)

    # Define a 2-batch.
    dist = Binomial(n=[4., 5], p=[.1, .3])
    ```

    """
    self._logits, self._p = distribution_util.get_logits_and_prob(
        name=name, logits=logits, p=p, validate_args=validate_args)
    with ops.name_scope(name, values=[n]) as ns:
      with ops.control_dependencies([
          check_ops.assert_non_negative(
              n, message="n has negative components."),
          distribution_util.assert_integer_form(
              n, message="n has non-integer components."),
      ] if validate_args else []):
        self._n = array_ops.identity(n, name="n")
        super(Binomial, self).__init__(
            dtype=self._p.dtype,
            parameters={"n": self._n, "p": self._p, "logits": self._logits},
            is_continuous=False,
            is_reparameterized=False,
            validate_args=validate_args,
            allow_nan_stats=allow_nan_stats,
            name=ns)

 def _assert_valid_sample(self, x, check_integer=True):
   if not self.validate_args: return x
   with ops.name_scope("check_x", values=[x]):
     dependencies = [check_ops.assert_non_negative(x)]
     if check_integer:
       dependencies += [distribution_util.assert_integer_form(
           x, message="x has non-integer components.")]
     return control_flow_ops.with_dependencies(dependencies, x)

def embed_check_nonnegative_discrete(x, check_integer=True):
  """Assert x is a non-negative tensor, and optionally of integers."""
  assertions = [check_ops.assert_non_negative(
      x, message="x must be non-negative.")]
  if check_integer:
    assertions += [assert_integer_form(
        x, message="x cannot contain fractional components.")]
  return control_flow_ops.with_dependencies(assertions, x)

 def test_empty_tensor_doesnt_raise(self):
   # A tensor is non-negative when it satisfies:
   #   For every element x_i in x, x_i >= 0
   # and an empty tensor has no elements, so this is trivially satisfied.
   # This is standard set theory.
   empty = constant_op.constant([], name="empty")
   with ops.control_dependencies([check_ops.assert_non_negative(empty)]):
     out = array_ops.identity(empty)
   self.evaluate(out)

 def _maybe_assert_valid_y(self, y):
   if not self.validate_args:
     return y
   is_positive = check_ops.assert_non_negative(
       y, message="Inverse transformation input must be greater than 0.")
   less_than_one = check_ops.assert_less_equal(
       y, constant_op.constant(1., y.dtype),
       message="Inverse transformation input must be less than or equal to 1.")
   return control_flow_ops.with_dependencies([is_positive, less_than_one], y)

def _assert_range(labels, n_classes):
  assert_less = check_ops.assert_less(
      labels,
      ops.convert_to_tensor(n_classes, dtype=labels.dtype),
      message='Label IDs must < n_classes')
  assert_greater = check_ops.assert_non_negative(
      labels, message='Label IDs must >= 0')
  with ops.control_dependencies((assert_less, assert_greater)):
    return array_ops.identity(labels)

 def _inverse_log_det_jacobian(self, y):
   # If event_ndims = 2,
   # F^{-1}(y) = (-y, y), so DF^{-1}(y) = (-1, 1),
   # so Log|DF^{-1}(y)| = Log[1, 1] = [0, 0].
   zeros = constant_op.constant(0., dtype=y.dtype)
   if self.validate_args:
     zeros = control_flow_ops.with_dependencies(
         [check_ops.assert_non_negative(y, message="Argument y was negative")],
         zeros)
   return zeros, zeros

 def _check_x(self, x, check_integer=True):
   with ops.name_scope('check_x', values=[x]):
     x = ops.convert_to_tensor(x, name="x")
     if not self.validate_args:
       return x
     dependencies = [check_ops.assert_non_negative(x)]
     if check_integer:
       dependencies += [distribution_util.assert_integer_form(
           x, message="x has non-integer components.")]
     return control_flow_ops.with_dependencies(dependencies, x)

  def _prepare(self):
    # We need to put the conversion and check here because a user will likely
    # want to decay the learning rate dynamically.
    self._learning_rate_tensor = control_flow_ops.with_dependencies([
        check_ops.assert_non_negative(
            self._learning_rate, message='`learning_rate` must be non-negative')
    ], ops.convert_to_tensor(self._learning_rate, name='learning_rate_tensor'))
    self._decay_tensor = ops.convert_to_tensor(
        self._preconditioner_decay_rate, name='preconditioner_decay_rate')

    super(SGLDOptimizer, self)._prepare()

 def _check_counts(self, counts):
   counts = ops.convert_to_tensor(counts, name="counts_before_deps")
   if not self.validate_args:
     return counts
   return control_flow_ops.with_dependencies([
       check_ops.assert_non_negative(
           counts, message="counts has negative components."),
       check_ops.assert_less_equal(
           counts, self._n, message="counts are not less than or equal to n."),
       distribution_util.assert_integer_form(
           counts, message="counts have non-integer components.")], counts)

 def _inverse_log_det_jacobian(self, y):
   # If event_ndims = 2,
   # F^{-1}(y) = (-y, y), so DF^{-1}(y) = (-1, 1),
   # so Log|DF^{-1}(y)| = Log[1, 1] = [0, 0].
   batch_shape = array_ops.shape(y)[:array_ops.rank(y) - self.event_ndims]
   zeros = array_ops.zeros(batch_shape, dtype=y.dtype)
   if self.validate_args:
     zeros = control_flow_ops.with_dependencies(
         [check_ops.assert_non_negative(y, message="Argument y was negative")],
         zeros)
   return zeros, zeros

  def _check_counts(self, counts):
    """Check counts for proper shape, values, then return tensor version."""
    counts = ops.convert_to_tensor(counts, name='counts')
    if not self.validate_args:
      return counts
    candidate_n = math_ops.reduce_sum(counts, reduction_indices=[-1])

    return control_flow_ops.with_dependencies([
        check_ops.assert_non_negative(counts),
        check_ops.assert_equal(self._n, candidate_n),
        _assert_integer_form(counts)], counts)

 def _maybe_assert_valid(self, x):
   if not self.validate_args:
     return x
   return control_flow_ops.with_dependencies([
       check_ops.assert_non_negative(
           x,
           message="sample must be non-negative"),
       check_ops.assert_less_equal(
           x, array_ops.ones([], self.concentration0.dtype),
           message="sample must be no larger than `1`."),
   ], x)