def append(self, transitions, rows=None):
    """Append a batch of transitions to rows of the memory.

    Args:
      transitions: Tuple of transition quantities with batch dimension.
      rows: Episodes to append to, defaults to all.

    Returns:
      Operation.
    """
    rows = tf.range(self._capacity) if rows is None else rows
    assert rows.shape.ndims == 1
    assert_capacity = tf.assert_less(
        rows, self._capacity,
        message='capacity exceeded')
    with tf.control_dependencies([assert_capacity]):
      assert_max_length = tf.assert_less(
          tf.gather(self._length, rows), self._max_length,
          message='max length exceeded')
    append_ops = []
    with tf.control_dependencies([assert_max_length]):
      for buffer_, elements in zip(self._buffers, transitions):
        timestep = tf.gather(self._length, rows)
        indices = tf.stack([rows, timestep], 1)
        append_ops.append(tf.scatter_nd_update(buffer_, indices, elements))
    with tf.control_dependencies(append_ops):
      episode_mask = tf.reduce_sum(tf.one_hot(
          rows, self._capacity, dtype=tf.int32), 0)
      return self._length.assign_add(episode_mask)

  def _variance(self):
    # We need to put the tf.where inside the outer tf.where to ensure we never
    # hit a NaN in the gradient.
    denom = tf.where(tf.greater(self.df, 2.),
                     self.df - 2.,
                     tf.ones_like(self.df))
    # Abs(scale) superfluous.
    var = (tf.ones(self.batch_shape_tensor(), dtype=self.dtype) *
           tf.square(self.scale) * self.df / denom)
    # When 1 < df <= 2, variance is infinite.
    inf = np.array(np.inf, dtype=self.dtype.as_numpy_dtype())
    result_where_defined = tf.where(
        self.df > tf.fill(self.batch_shape_tensor(), 2.),
        var,
        tf.fill(self.batch_shape_tensor(), inf, name="inf"))

    if self.allow_nan_stats:
      nan = np.array(np.nan, dtype=self.dtype.as_numpy_dtype())
      return tf.where(
          tf.greater(
              self.df,
              tf.ones(self.batch_shape_tensor(), dtype=self.dtype)),
          result_where_defined,
          tf.fill(self.batch_shape_tensor(), nan, name="nan"))
    else:
      return control_flow_ops.with_dependencies(
          [
              tf.assert_less(
                  tf.ones([], dtype=self.dtype),
                  self.df,
                  message="variance not defined for components of df <= 1"),
          ],
          result_where_defined)

  def replace(self, episodes, length, rows=None):
    """Replace full episodes.

    Args:
      episodes: Tuple of transition quantities with batch and time dimensions.
      length: Batch of sequence lengths.
      rows: Episodes to replace, defaults to all.

    Returns:
      Operation.
    """
    rows = tf.range(self._capacity) if rows is None else rows
    assert rows.shape.ndims == 1
    assert_capacity = tf.assert_less(
        rows, self._capacity, message='capacity exceeded')
    with tf.control_dependencies([assert_capacity]):
      assert_max_length = tf.assert_less_equal(
          length, self._max_length, message='max length exceeded')
    replace_ops = []
    with tf.control_dependencies([assert_max_length]):
      for buffer_, elements in zip(self._buffers, episodes):
        replace_op = tf.scatter_update(buffer_, rows, elements)
        replace_ops.append(replace_op)
    with tf.control_dependencies(replace_ops):
      return tf.scatter_update(self._length, rows, length)

  def _testSampleLogProbExact(
      self, concentrations, det_bounds, dim, means,
      num_samples=int(1e5), dtype=np.float32, target_discrepancy=0.1, seed=42):
    # For test methodology see the comment in
    # _testSampleConsistentLogProbInterval, except that this test
    # checks those parameter settings where the true volume is known
    # analytically.
    concentration = np.array(concentrations, dtype=dtype)
    det_bounds = np.array(det_bounds, dtype=dtype)
    means = np.array(means, dtype=dtype)
    # Add a tolerance to guard against some of the importance_weights exceeding
    # the theoretical maximum (importance_maxima) due to numerical inaccuracies
    # while lower bounding the determinant. See corresponding comment in
    # _testSampleConsistentLogProbInterval.
    high_tolerance = 1e-6

    testee_lkj = tfd.LKJ(
        dimension=dim, concentration=concentration, validate_args=True)
    x = testee_lkj.sample(num_samples, seed=seed)
    importance_weights = (
        tf.exp(-testee_lkj.log_prob(x)) * _det_ok_mask(x, det_bounds))
    importance_maxima = (1. / det_bounds) ** (concentration - 1) * tf.exp(
        testee_lkj._log_normalization())

    chk1 = st.assert_true_mean_equal_by_dkwm(
        importance_weights, low=0., high=importance_maxima + high_tolerance,
        expected=means, false_fail_rate=1e-6)
    chk2 = tf.assert_less(
        st.min_discrepancy_of_true_means_detectable_by_dkwm(
            num_samples, low=0., high=importance_maxima + high_tolerance,
            false_fail_rate=1e-6, false_pass_rate=1e-6),
        dtype(target_discrepancy))
    self.evaluate([chk1, chk2])

 def test_raises_when_equal(self):
   with self.test_session():
     small = tf.constant([1, 2], name="small")
     with tf.control_dependencies([tf.assert_less(small, small)]):
       out = tf.identity(small)
     with self.assertRaisesOpError("small.*small"):
       out.eval()

 def test_doesnt_raise_when_less_and_broadcastable_shapes(self):
   with self.test_session():
     small = tf.constant([1], name="small")
     big = tf.constant([3, 2], name="big")
     with tf.control_dependencies([tf.assert_less(small, big)]):
       out = tf.identity(small)
     out.eval()

 def _entropy(self):
   probs = self._probs
   if self.validate_args:
     probs = control_flow_ops.with_dependencies([
         tf.assert_less(
             probs,
             tf.constant(1., probs.dtype),
             message="Entropy is undefined when logits = inf or probs = 1.")
     ], probs)
   # Claim: entropy(p) = softplus(s)/p - s
   # where s=logits and p=probs.
   #
   # Proof:
   #
   # entropy(p)
   # := -[(1-p)log(1-p) + plog(p)]/p
   # = -[log(1-p) + plog(p/(1-p))]/p
   # = -[-softplus(s) + ps]/p
   # = softplus(s)/p - s
   #
   # since,
   # log[1-sigmoid(s)]
   # = log[1/(1+exp(s)]
   # = -log[1+exp(s)]
   # = -softplus(s)
   #
   # using the fact that,
   # 1-sigmoid(s) = sigmoid(-s) = 1/(1+exp(s))
   return tf.nn.softplus(self.logits) / probs - self.logits

 def test_doesnt_raise_when_both_empty(self):
   with self.test_session():
     larry = tf.constant([])
     curly = tf.constant([])
     with tf.control_dependencies([tf.assert_less(larry, curly)]):
       out = tf.identity(larry)
     out.eval()

 def test_raises_when_greater(self):
   with self.test_session():
     small = tf.constant([1, 2], name="small")
     big = tf.constant([3, 4], name="big")
     with tf.control_dependencies([tf.assert_less(big, small)]):
       out = tf.identity(small)
     with self.assertRaisesOpError("big.*small"):
       out.eval()

 def test_raises_when_less_but_non_broadcastable_shapes(self):
   with self.test_session():
     small = tf.constant([1, 1, 1], name="small")
     big = tf.constant([3, 2], name="big")
     with self.assertRaisesRegexp(ValueError, "broadcast"):
       with tf.control_dependencies([tf.assert_less(small, big)]):
         out = tf.identity(small)
       out.eval()

 def mixture_kl():
     with tf.control_dependencies([tf.assert_greater(consistency_trust, 0.0),
                                   tf.assert_less(consistency_trust, 1.0)]):
         uniform = tf.constant(1 / num_classes, shape=[num_classes])
         mixed_softmax1 = consistency_trust * softmax1 + (1 - consistency_trust) * uniform
         mixed_softmax2 = consistency_trust * softmax2 + (1 - consistency_trust) * uniform
         costs = tf.reduce_sum(mixed_softmax2 * tf.log(mixed_softmax2 / mixed_softmax1), axis=1)
         costs = costs * kl_cost_multiplier
         return costs

 def _maybe_assert_valid_sample(self, x):
   """Checks the validity of a sample."""
   if not self.validate_args:
     return x
   return control_flow_ops.with_dependencies([
       tf.assert_positive(x, message="sample must be positive"),
       tf.assert_less(
           x,
           tf.ones([], self.dtype),
           message="sample must be less than `1`."),
   ], x)

 def _mode(self):
   mode = (self.concentration1 - 1.) / (self.total_concentration - 2.)
   if self.allow_nan_stats:
     nan = tf.fill(
         self.batch_shape_tensor(),
         np.array(np.nan, dtype=self.dtype.as_numpy_dtype()),
         name="nan")
     is_defined = tf.logical_and(self.concentration1 > 1.,
                                 self.concentration0 > 1.)
     return tf.where(is_defined, mode, nan)
   return control_flow_ops.with_dependencies([
       tf.assert_less(
           tf.ones([], dtype=self.dtype),
           self.concentration1,
           message="Mode undefined for concentration1 <= 1."),
       tf.assert_less(
           tf.ones([], dtype=self.dtype),
           self.concentration0,
           message="Mode undefined for concentration0 <= 1.")
   ], mode)

  def _mode(self):
    a = self.concentration1
    b = self.concentration0
    mode = ((a - 1) / (a * b - 1))**(1. / a)
    if self.allow_nan_stats:
      nan = tf.fill(
          self.batch_shape_tensor(),
          np.array(np.nan, dtype=self.dtype.as_numpy_dtype),
          name="nan")
      is_defined = (self.concentration1 > 1.) & (self.concentration0 > 1.)
      return tf.where(is_defined, mode, nan)

    return control_flow_ops.with_dependencies([
        tf.assert_less(
            tf.ones([], dtype=self.concentration1.dtype),
            self.concentration1,
            message="Mode undefined for concentration1 <= 1."),
        tf.assert_less(
            tf.ones([], dtype=self.concentration0.dtype),
            self.concentration0,
            message="Mode undefined for concentration0 <= 1.")
    ], mode)

 def _maybe_assert_valid_concentration(self, concentration, validate_args):
   """Checks the validity of the concentration parameter."""
   if not validate_args:
     return concentration
   return control_flow_ops.with_dependencies([
       tf.assert_positive(
           concentration,
           message="Concentration parameter must be positive."),
       tf.assert_rank_at_least(
           concentration, 1,
           message="Concentration parameter must have >=1 dimensions."),
       tf.assert_less(
           1, tf.shape(concentration)[-1],
           message="Concentration parameter must have event_size >= 2."),
   ], concentration)

 def _mean(self):
   mean = self.rate / (self.concentration - 1.)
   if self.allow_nan_stats:
     nan = tf.fill(
         self.batch_shape_tensor(),
         np.array(np.nan, dtype=self.dtype.as_numpy_dtype()),
         name="nan")
     return tf.where(self.concentration > 1., mean, nan)
   else:
     return control_flow_ops.with_dependencies([
         tf.assert_less(
             tf.ones([], self.dtype),
             self.concentration,
             message="mean undefined when any concentration <= 1"),
     ], mean)

  def _maybe_assert_valid_y(self, y):
    if not self.validate_args:
      return y
    is_valid = [
        tf.assert_greater(
            y,
            tf.cast(-1., dtype=y.dtype.base_dtype),
            message="Inverse transformation input must be greater than -1."),
        tf.assert_less(
            y,
            tf.cast(1., dtype=y.dtype.base_dtype),
            message="Inverse transformation input must be less than 1.")
    ]

    return control_flow_ops.with_dependencies(is_valid, y)

  def __init__(self,
               low=0.,
               high=1.,
               validate_args=False,
               allow_nan_stats=True,
               name="Uniform"):
    """Initialize a batch of Uniform distributions.

    Args:
      low: Floating point tensor, lower boundary of the output interval. Must
        have `low < high`.
      high: Floating point tensor, upper boundary of the output interval. Must
        have `low < high`.
      validate_args: Python `bool`, default `False`. When `True` distribution
        parameters are checked for validity despite possibly degrading runtime
        performance. When `False` invalid inputs may silently render incorrect
        outputs.
      allow_nan_stats: Python `bool`, default `True`. When `True`, statistics
        (e.g., mean, mode, variance) use the value "`NaN`" to indicate the
        result is undefined. When `False`, an exception is raised if one or
        more of the statistic's batch members are undefined.
      name: Python `str` name prefixed to Ops created by this class.

    Raises:
      InvalidArgumentError: if `low >= high` and `validate_args=False`.
    """
    parameters = dict(locals())
    with tf.name_scope(name, values=[low, high]) as name:
      dtype = dtype_util.common_dtype([low, high], tf.float32)
      low = tf.convert_to_tensor(low, name="low", dtype=dtype)
      high = tf.convert_to_tensor(high, name="high", dtype=dtype)
      with tf.control_dependencies([
          tf.assert_less(
              low, high, message="uniform not defined when low >= high.")
      ] if validate_args else []):
        self._low = tf.identity(low)
        self._high = tf.identity(high)
        tf.assert_same_float_dtype([self._low, self._high])
    super(Uniform, self).__init__(
        dtype=self._low.dtype,
        reparameterization_type=reparameterization.FULLY_REPARAMETERIZED,
        validate_args=validate_args,
        allow_nan_stats=allow_nan_stats,
        parameters=parameters,
        graph_parents=[self._low,
                       self._high],
        name=name)

 def _variance(self):
   var = (
       tf.square(self.rate) / tf.square(self.concentration - 1.) /
       (self.concentration - 2.))
   if self.allow_nan_stats:
     nan = tf.fill(
         self.batch_shape_tensor(),
         np.array(np.nan, dtype=self.dtype.as_numpy_dtype()),
         name="nan")
     return tf.where(self.concentration > 2., var, nan)
   else:
     return control_flow_ops.with_dependencies([
         tf.assert_less(
             tf.constant(2., dtype=self.dtype),
             self.concentration,
             message="variance undefined when any concentration <= 2"),
     ], var)

 def _mode(self):
   k = tf.cast(self.event_shape_tensor()[0], self.dtype)
   mode = (self.concentration - 1.) / (
       self.total_concentration[..., tf.newaxis] - k)
   if self.allow_nan_stats:
     nan = tf.fill(
         tf.shape(mode),
         np.array(np.nan, dtype=self.dtype.as_numpy_dtype()),
         name="nan")
     return tf.where(
         tf.reduce_all(self.concentration > 1., axis=-1),
         mode, nan)
   return control_flow_ops.with_dependencies([
       tf.assert_less(
           tf.ones([], self.dtype),
           self.concentration,
           message="Mode undefined when any concentration <= 1"),
   ], mode)