def _lower_triangular_mask(shape):
  """Creates a lower-triangular boolean mask over the last 2 dimensions."""
  row_index = math_ops.cumsum(
      array_ops.ones(shape=shape, dtype=dtypes.int32), axis=-2)
  col_index = math_ops.cumsum(
      array_ops.ones(shape=shape, dtype=dtypes.int32), axis=-1)
  return math_ops.greater_equal(row_index, col_index)

 def testAxisType(self):
   for dtype in self.valid_dtypes:
     x = np.arange(1, 6).reshape([5]).astype(dtype)
     for axis_dtype in self.axis_dtypes():
       with self.cached_session(), self.test_scope():
         p = array_ops.placeholder(x.dtype)
         axis = constant_op.constant(0, axis_dtype)
         math_ops.cumsum(p, axis).eval(feed_dict={p: x})

 def testAxisType(self):
   for dtype in self.valid_dtypes:
     x = np.arange(1, 6).reshape([5]).astype(dtype)
     for axis_dtype in [dtypes.int64, dtypes.int32]:
       with self.cached_session(use_gpu=True):
         axis = constant_op.constant(0, axis_dtype)
         tf_out = math_ops.cumsum(x, axis).eval()

def power_sums_tensor(array_size, power_matrix, multiplier):
  r"""Computes \sum_{i=0}^{N-1} A^i B (A^i)^T for N=0..(array_size + 1).

  Args:
    array_size: The number of non-trivial sums to pre-compute.
    power_matrix: The "A" matrix above.
    multiplier: The "B" matrix above
  Returns:
    A Tensor with S[N] = \sum_{i=0}^{N-1} A^i B (A^i)^T
      S[0] is the zero matrix
      S[1] is B
      S[2] is A B A^T + B
      ...and so on
  """
  array_size = math_ops.cast(array_size, dtypes.int32)
  power_matrix = ops.convert_to_tensor(power_matrix)
  identity_like_power_matrix = linalg_ops.eye(
      array_ops.shape(power_matrix)[0], dtype=power_matrix.dtype)
  identity_like_power_matrix.set_shape(
      ops.convert_to_tensor(power_matrix).get_shape())
  transition_powers = functional_ops.scan(
      lambda previous_power, _: math_ops.matmul(previous_power, power_matrix),
      math_ops.range(array_size - 1),
      initializer=identity_like_power_matrix)
  summed = math_ops.cumsum(
      array_ops.concat([
          array_ops.expand_dims(multiplier, 0), math_ops.matmul(
              batch_times_matrix(transition_powers, multiplier),
              transition_powers,
              adjoint_b=True)
      ], 0))
  return array_ops.concat(
      [array_ops.expand_dims(array_ops.zeros_like(multiplier), 0), summed], 0)

  def _compare(self, x, axis, exclusive, reverse):
    np_out = handle_options(np.cumsum, x, axis, exclusive, reverse)
    with self.cached_session(), self.test_scope():
      p = array_ops.placeholder(x.dtype)
      tf_out = math_ops.cumsum(p, axis, exclusive, reverse).eval(
          feed_dict={p: x})

    self.assertAllClose(np_out, tf_out)

 def _compareGradient(self, shape, axis, exclusive, reverse):
   x = np.arange(0, 50).reshape(shape).astype(np.float64)
   with self.cached_session(use_gpu=True):
     t = ops.convert_to_tensor(x)
     result = math_ops.cumsum(t, axis, exclusive, reverse)
     jacob_t, jacob_n = gradient_checker.compute_gradient(
         t, shape, result, shape, x_init_value=x, delta=1)
   self.assertAllClose(jacob_t, jacob_n, rtol=1e-8, atol=1e-8)

def _CumsumGrad(op, grad):
  axis = op.inputs[1]
  exclusive = op.get_attr("exclusive")
  reverse = op.get_attr("reverse")
  return [
      math_ops.cumsum(grad, axis, exclusive=exclusive, reverse=not reverse),
      None
  ]

def segment_ids_to_row_splits(segment_ids, num_segments=None,
                              out_type=None, name=None):
  """Generates the RaggedTensor `row_splits` corresponding to a segmentation.

  Returns an integer vector `splits`, where `splits[0] = 0` and
  `splits[i] = splits[i-1] + count(segment_ids==i)`.  Example:

  ```python
  >>> ragged.segment_ids_to_row_splits([0, 0, 0, 2, 2, 3, 4, 4, 4]).eval()
  [ 0 3 3 5 6 9 ]
  ```

  Args:
    segment_ids: A 1-D integer Tensor.
    num_segments: A scalar integer indicating the number of segments.  Defaults
      to `max(segment_ids) + 1` (or zero if `segment_ids` is empty).
    out_type: The dtype for the return value.  Defaults to `segment_ids.dtype`,
      or `tf.int64` if `segment_ids` does not have a dtype.
    name: A name prefix for the returned tensor (optional).

  Returns:
    A sorted 1-D integer Tensor, with `shape=[num_segments + 1]`.
  """
  if out_type is None:
    if isinstance(segment_ids, ops.Tensor):
      out_type = segment_ids.dtype
    elif isinstance(num_segments, ops.Tensor):
      out_type = num_segments.dtype
    else:
      out_type = dtypes.int64
  else:
    out_type = dtypes.as_dtype(out_type)
  with ops.name_scope(name, "SegmentIdsToRaggedSplits", [segment_ids]) as name:
    # Note: we cast int64 tensors to int32, since bincount currently only
    # supports int32 inputs.
    segment_ids = ragged_util.convert_to_int_tensor(segment_ids, "segment_ids",
                                                    dtype=dtypes.int32)
    segment_ids.shape.assert_has_rank(1)
    if num_segments is not None:
      num_segments = ragged_util.convert_to_int_tensor(num_segments,
                                                       "num_segments",
                                                       dtype=dtypes.int32)
      num_segments.shape.assert_has_rank(0)

    row_lengths = math_ops.bincount(
        segment_ids,
        minlength=num_segments,
        maxlength=num_segments,
        dtype=out_type)
    splits = array_ops.concat([[0], math_ops.cumsum(row_lengths)], axis=0)

    # Update shape information, if possible.
    if num_segments is not None:
      const_num_segments = tensor_util.constant_value(num_segments)
      if const_num_segments is not None:
        splits.set_shape(tensor_shape.TensorShape([const_num_segments + 1]))

    return splits

def _CumprodGrad(op, grad):
  x = op.inputs[0]
  axis = op.inputs[1]
  reverse = op.get_attr("reverse")

  # TODO This fails when x contains 0 and should be fixed
  prod = math_ops.cumprod(x, axis=axis, reverse=reverse)
  out = math_ops.cumsum(prod * grad, axis=axis, reverse=(not reverse))
  return [out / x, None]

def _effective_sample_size_single_state(states, filter_beyond_lag,
                                        filter_threshold):
  """ESS computation for one single Tensor argument."""

  with ops.name_scope(
      "effective_sample_size_single_state",
      values=[states, filter_beyond_lag, filter_threshold]):

    states = ops.convert_to_tensor(states, name="states")
    dt = states.dtype

    # filter_beyond_lag == None ==> auto_corr is the full sequence.
    auto_corr = sample_stats.auto_correlation(
        states, axis=0, max_lags=filter_beyond_lag)
    if filter_threshold is not None:
      filter_threshold = ops.convert_to_tensor(
          filter_threshold, dtype=dt, name="filter_threshold")
      # Get a binary mask to zero out values of auto_corr below the threshold.
      #   mask[i, ...] = 1 if auto_corr[j, ...] > threshold for all j <= i,
      #   mask[i, ...] = 0, otherwise.
      # So, along dimension zero, the mask will look like [1, 1, ..., 0, 0,...]
      # Building step by step,
      #   Assume auto_corr = [1, 0.5, 0.0, 0.3], and filter_threshold = 0.2.
      # Step 1:  mask = [False, False, True, False]
      mask = auto_corr < filter_threshold
      # Step 2:  mask = [0, 0, 1, 1]
      mask = math_ops.cast(mask, dtype=dt)
      # Step 3:  mask = [0, 0, 1, 2]
      mask = math_ops.cumsum(mask, axis=0)
      # Step 4:  mask = [1, 1, 0, 0]
      mask = math_ops.maximum(1. - mask, 0.)
      auto_corr *= mask

    # With R[k] := auto_corr[k, ...],
    # ESS = N / {1 + 2 * Sum_{k=1}^N (N - k) / N * R[k]}
    #     = N / {-1 + 2 * Sum_{k=0}^N (N - k) / N * R[k]} (since R[0] = 1)
    #     approx N / {-1 + 2 * Sum_{k=0}^M (N - k) / N * R[k]}
    # where M is the filter_beyond_lag truncation point chosen above.

    # Get the factor (N - k) / N, and give it shape [M, 1,...,1], having total
    # ndims the same as auto_corr
    n = _axis_size(states, axis=0)
    k = math_ops.range(0., _axis_size(auto_corr, axis=0))
    nk_factor = (n - k) / n
    if auto_corr.shape.ndims is not None:
      new_shape = [-1] + [1] * (auto_corr.shape.ndims - 1)
    else:
      new_shape = array_ops.concat(
          ([-1],
           array_ops.ones([array_ops.rank(auto_corr) - 1], dtype=dtypes.int32)),
          axis=0)
    nk_factor = array_ops.reshape(nk_factor, new_shape)

    return n / (-1 + 2 * math_ops.reduce_sum(nk_factor * auto_corr, axis=0))

  def test_searchsorted(self):
    sorted_inputs = math_ops.cumsum(random_ops.random_uniform([3, 2, 4]),
                                    axis=-1)
    values = random_ops.random_uniform([2, 3], minval=-1, maxval=4.5)

    def loop_fn(i):
      inputs_i = array_ops.gather(sorted_inputs, i)
      return [array_ops.searchsorted(inputs_i, values, out_type=dtypes.int32,
                                     side="left"),  # creates LowerBound op.
              array_ops.searchsorted(inputs_i, values, out_type=dtypes.int64,
                                     side="right")]  # creates UpperBound op.

    self._test_loop_fn(loop_fn, 3, loop_fn_dtypes=[dtypes.int32, dtypes.int64])

 def make_tril_ids(n):
   """Internal helper to create vector of linear indices into y."""
   cols = array_ops.reshape(array_ops.tile(math_ops.range(n), [n]), [n, n])
   rows = array_ops.tile(
       array_ops.expand_dims(math_ops.range(n), -1), [1, n])
   pred = math_ops.greater(cols, rows)
   tril_ids = array_ops.tile(array_ops.reshape(
       math_ops.cumsum(math_ops.range(n)), [n, 1]), [1, n]) + cols
   tril_ids = math_ops.select(pred,
                              array_ops.zeros([n, n], dtype=dtypes.int32),
                              tril_ids + 1)
   tril_ids = array_ops.reshape(tril_ids, [-1])
   return tril_ids

  def _update_mask(self, weights, threshold):
    """Updates the mask for a given weight tensor.

    This functions first computes the cdf of the weight tensor, and estimates
    the threshold value such that 'desired_sparsity' fraction of weights
    have magnitude less than the threshold.

    Args:
      weights: The weight tensor that needs to be masked.
      threshold: The current threshold value. The function will compute a new
        threshold and return the exponential moving average using the current
        value of threshold

    Returns:
      new_threshold: The new value of the threshold based on weights, and
        sparsity at the current global_step
      new_mask: A numpy array of the same size and shape as weights containing
        0 or 1 to indicate which of the values in weights falls below
        the threshold

    Raises:
      ValueError: if sparsity is not defined
    """
    if self._sparsity is None:
      raise ValueError('Sparsity variable undefined')

    with ops.name_scope(weights.op.name + '_pruning_ops'):
      abs_weights = math_ops.abs(weights)
      max_value = math_ops.reduce_max(abs_weights)
      histogram = _histogram(
          abs_weights, [0.0, max_value],
          nbins=self._spec.nbins,
          dtype=np.float32)

      cdf = math_ops.cumsum(histogram)
      norm_cdf = math_ops.div(cdf, math_ops.reduce_sum(histogram))
      current_threshold = math_ops.multiply(
          math_ops.div(
              math_ops.reduce_sum(
                  math_ops.cast(
                      math_ops.less(norm_cdf, self._sparsity), np.float32)),
              float(self._spec.nbins)), max_value)

      smoothed_threshold = math_ops.add_n([
          math_ops.multiply(current_threshold, 1 - self._spec.threshold_decay),
          math_ops.multiply(threshold, self._spec.threshold_decay)
      ])
      new_mask = math_ops.cast(
          math_ops.greater(abs_weights, smoothed_threshold), np.float32)
    return smoothed_threshold, new_mask

def _randomize(coeffs, radixes, seed=None):
  """Applies the Owen randomization to the coefficients."""
  given_dtype = coeffs.dtype
  coeffs = math_ops.to_int32(coeffs)
  num_coeffs = array_ops.shape(coeffs)[-1]
  radixes = array_ops.reshape(math_ops.to_int32(radixes), [-1])
  perms = _get_permutations(num_coeffs, radixes, seed=seed)
  perms = array_ops.reshape(perms, [-1])
  radix_sum = math_ops.reduce_sum(radixes)
  radix_offsets = array_ops.reshape(math_ops.cumsum(radixes, exclusive=True),
                                    [-1, 1])
  offsets = radix_offsets + math_ops.range(num_coeffs) * radix_sum
  permuted_coeffs = array_ops.gather(perms, coeffs + offsets)
  return math_ops.cast(permuted_coeffs, dtype=given_dtype)

  def testReadmeExample(self):
    data = random_ops.random_uniform((128, 128), 0, 10, dtype=dtypes.int32)
    histogram = math_ops.bincount(data, minlength=10, maxlength=10)
    cdf = math_ops.cumsum(histogram, exclusive=False)
    cdf = array_ops.pad(cdf, [[1, 0]])
    cdf = array_ops.reshape(cdf, [1, 1, -1])

    data = math_ops.cast(data, dtypes.int16)
    encoded = coder_ops.range_encode(data, cdf, precision=14)
    decoded = coder_ops.range_decode(
        encoded, array_ops.shape(data), cdf, precision=14)

    with self.test_session() as sess:
      self.assertAllEqual(*sess.run((data, decoded)))

def sparsemax(logits, name=None):
  """Computes sparsemax activations [1].

  For each batch `i` and class `j` we have
    sparsemax[i, j] = max(logits[i, j] - tau(logits[i, :]), 0)

  [1]: https://arxiv.org/abs/1602.02068

  Args:
    logits: A `Tensor`. Must be one of the following types: `half`, `float32`,
      `float64`.
    name: A name for the operation (optional).

  Returns:
    A `Tensor`. Has the same type as `logits`.
  """

  with ops.name_scope(name, "sparsemax", [logits]) as name:
    logits = ops.convert_to_tensor(logits, name="logits")
    obs = array_ops.shape(logits)[0]
    dims = array_ops.shape(logits)[1]

    z = logits - math_ops.reduce_mean(logits, axis=1)[:, array_ops.newaxis]

    # sort z
    z_sorted, _ = nn.top_k(z, k=dims)

    # calculate k(z)
    z_cumsum = math_ops.cumsum(z_sorted, axis=1)
    k = math_ops.range(
        1, math_ops.cast(dims, logits.dtype) + 1, dtype=logits.dtype
    )
    z_check = 1 + k * z_sorted > z_cumsum
    # because the z_check vector is always [1,1,...1,0,0,...0] finding the
    # (index + 1) of the last `1` is the same as just summing the number of 1.
    k_z = math_ops.reduce_sum(math_ops.cast(z_check, dtypes.int32), axis=1)

    # calculate tau(z)
    indices = array_ops.stack([math_ops.range(0, obs), k_z - 1], axis=1)
    tau_sum = array_ops.gather_nd(z_cumsum, indices)
    tau_z = (tau_sum - 1) / math_ops.cast(k_z, logits.dtype)

    # calculate p
    return math_ops.maximum(
        math_ops.cast(0, logits.dtype),
        z - tau_z[:, array_ops.newaxis]
    )

def segment_ids_to_row_splits(segment_ids, num_segments=None, name=None):
  """Generates the RaggedTensor `row_splits` corresponding to a segmentation.

  Returns an integer vector `splits`, where `splits[0] = 0` and
  `splits[i] = splits[i-1] + count(segment_ids==i)`.  Example:

  ```python
  >>> ragged.segment_ids_to_row_splits([0, 0, 0, 2, 2, 3, 4, 4, 4]).eval()
  [ 0 3 3 5 6 9 ]
  ```

  Args:
    segment_ids: A 1-D integer Tensor.
    num_segments: A scalar integer indicating the number of segments.  Defaults
      to `max(segment_ids) + 1` (or zero if `segment_ids` is empty).
    name: A name prefix for the returned tensor (optional).

  Returns:
    A sorted 1-D int64 Tensor, with `shape=[num_segments + 1]`.
  """
  with ops.name_scope(name, "SegmentIdsToRaggedSplits", [segment_ids]) as name:
    segment_ids = ragged_util.convert_to_int_tensor(segment_ids, "segment_ids")
    segment_ids.shape.assert_has_rank(1)
    if num_segments is not None:
      num_segments = ragged_util.convert_to_int_tensor(num_segments,
                                                       "num_segments")
      num_segments.shape.assert_has_rank(0)

    row_lengths = math_ops.bincount(
        segment_ids,
        minlength=num_segments,
        maxlength=num_segments,
        dtype=dtypes.int64)
    splits = array_ops.concat([[0], math_ops.cumsum(row_lengths)], axis=0)

    # Update shape information, if possible.
    if num_segments is not None:
      const_num_segments = tensor_util.constant_value(num_segments)
      if const_num_segments is not None:
        splits.set_shape(tensor_shape.TensorShape([const_num_segments + 1]))

    return splits

def safe_cumprod(x, *args, **kwargs):
  """Computes cumprod of x in logspace using cumsum to avoid underflow.

  The cumprod function and its gradient can result in numerical instabilities
  when its argument has very small and/or zero values.  As long as the argument
  is all positive, we can instead compute the cumulative product as
  exp(cumsum(log(x))).  This function can be called identically to tf.cumprod.

  Args:
    x: Tensor to take the cumulative product of.
    *args: Passed on to cumsum; these are identical to those in cumprod.
    **kwargs: Passed on to cumsum; these are identical to those in cumprod.
  Returns:
    Cumulative product of x.
  """
  with ops.name_scope(None, "SafeCumprod", [x]):
    x = ops.convert_to_tensor(x, name="x")
    tiny = np.finfo(x.dtype.as_numpy_dtype).tiny
    return math_ops.exp(math_ops.cumsum(
        math_ops.log(clip_ops.clip_by_value(x, tiny, 1)), *args, **kwargs))

def from_row_lengths(values, row_lengths, name=None):
  """Creates a `RaggedTensor` with rows partitioned by `row_lengths`.

  The returned `RaggedTensor` corresponds with the python list defined by:

  ```python
  result = [[values.pop(0) for i in range(length)]
            for length in row_lengths]
  ```

  Args:
    values: A potentially ragged tensor with shape `[nvals, ...]`.
    row_lengths: A 1-D int64 tensor with shape `[nrows]`.  Must be nonnegative.
      `sum(row_lengths)` must be `nvals`.
    name: A name prefix for the RaggedTensor (optional).

  Returns:
    A `RaggedTensor`.  `result.rank = values.rank + 1`.
    `result.ragged_rank = values.ragged_rank + 1`.

  #### Example:
    ```python
    >>> rt = ragged.from_row_lengths(
    ...     values=[3, 1, 4, 1, 5, 9, 2, 6],
    ...     row_lengths=[4, 0, 3, 1, 0])
    >>> rt.eval().tolist()
    [[3, 1, 4, 1], [], [5, 9, 2], [6], []]
    ```
  """
  with ops.name_scope(name, 'RaggedFromRowLengths', [values, row_lengths]):
    values = convert_to_tensor_or_ragged_tensor(values, name='values')
    row_lengths = ops.convert_to_tensor(row_lengths, dtypes.int64,
                                        'row_lengths')
    row_lengths.shape.assert_has_rank(1)
    row_limits = math_ops.cumsum(row_lengths)
    row_splits = array_ops.concat([[0], row_limits], axis=0)
    return ragged_tensor.RaggedTensor(
        values=values,
        row_splits=row_splits,
        cached_row_lengths=row_lengths,
        internal=True)

def compute_cdf_from_histogram(values, value_range, **kwargs):
  """Returns the normalized cumulative distribution of the given values tensor.

  Computes the histogram and uses tf.cumsum to arrive at cdf

  Args:
    values:  Numeric `Tensor`.
    value_range:  Shape [2] `Tensor` of same `dtype` as `values`.
    **kwargs: keyword arguments: nbins, name

  Returns:
    A 1-D `Tensor` holding normalized cdf of values.

  """
  nbins = kwargs.get('nbins', _NBINS)
  name = kwargs.get('name', None)
  with ops.name_scope(name, 'cdf', [values, value_range, nbins]):
    histogram = _histogram(
        values, value_range, dtype=dtypes.float32, nbins=nbins)
    cdf = math_ops.cumsum(histogram)
    return math_ops.div(cdf, math_ops.reduce_max(cdf))