code: # Copyright 2018 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """Support for ragged tensors."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import numpy as np
  
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import check_ops
  from tensorflow.python.ops import gen_ragged_math_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops.ragged import ragged_functional_ops
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow.python.ops.ragged import ragged_util
  from tensorflow.python.ops.ragged import segment_id_ops
  
  
  #===============================================================================
  # ragged.range
  #===============================================================================
  # pylint: disable=redefined-builtin
  def range(starts, limits=None, deltas=1, dtype=None, name=None):
    """Returns a `RaggedTensor` containing the specified sequences of numbers.
  
    Each row of the returned `RaggedTensor` contains a single sequence:
  
    ```python
    ragged.range(starts, limits, deltas)[i] ==
        tf.range(starts[i], limits[i], deltas[i])
    ```
  
    If `start[i] < limits[i] and deltas[i] > 0`, then `output[i]` will be an
    empty list.  Similarly, if `start[i] > limits[i] and deltas[i] < 0`, then
    `output[i]` will be an empty list.  This behavior is consistent with the
    Python `range` function, but differs from the `tf.range` op, which returns
    an error for these cases.
  
    Examples:
  
    ```python
    >>> ragged.range([3, 5, 2]).eval().tolist()
    [[0, 1, 2], [0, 1, 2, 3, 4], [0, 1]]
    >>> ragged.range([0, 5, 8], [3, 3, 12]).eval().tolist()
    [[0, 1, 2], [], [8, 9, 10, 11]]
    >>> ragged.range([0, 5, 8], [3, 3, 12], 2).eval().tolist()
    [[0, 2], [], [8, 10]]
    ```
  
    The input tensors `starts`, `limits`, and `deltas` may be scalars or vectors.
    The vector inputs must all have the same size.  Scalar inputs are broadcast
    to match the size of the vector inputs.
  
    Args:
      starts: Vector or scalar `Tensor`.  Specifies the first entry for each range
        if `limits` is not `None`; otherwise, specifies the range limits, and the
        first entries default to `0`.
      limits: Vector or scalar `Tensor`.  Specifies the exclusive upper limits for
        each range.
      deltas: Vector or scalar `Tensor`.  Specifies the increment for each range.
        Defaults to `1`.
      dtype: The type of the elements of the resulting tensor.  If not specified,
        then a value is chosen based on the other args.
      name: A name for the operation.
  
    Returns:
      A `RaggedTensor` of type `dtype` with `ragged_rank=1`.
    """
    if limits is None:
      starts, limits = 0, starts
  
    with ops.name_scope(name, 'RaggedRange', [starts, limits, deltas]) as name:
      starts = ops.convert_to_tensor(starts, dtype=dtype, name='starts')
      limits = ops.convert_to_tensor(limits, dtype=dtype, name='limits')
      deltas = ops.convert_to_tensor(deltas, dtype=dtype, name='deltas')
  
      # infer dtype if not explicitly provided
      if dtype is None:
        starts, limits, deltas = _infer_matching_dtype(
            [starts, limits, deltas],
            [dtypes.int32, dtypes.int64, dtypes.float32, dtypes.float64])
  
      result = gen_ragged_math_ops.ragged_range(starts, limits, deltas, name=name)
      return ragged_tensor.RaggedTensor.from_row_splits(result.rt_dense_values,
                                                        result.rt_nested_splits)
  
  
  def _infer_matching_dtype(tensors, dtype_hierarchy):
    """Infers a matching dtype for tensors, and casts them to that dtype."""
    assert all(t.dtype in dtype_hierarchy for t in tensors)
    inferred_dtype = max([t.dtype for t in tensors], key=dtype_hierarchy.index)
    return [math_ops.cast(t, inferred_dtype) for t in tensors]
  
  
  #===============================================================================
  # ragged_segment_<AGGREGATE>
  #===============================================================================
  
  # Docstring template used for the raggged_segment_<AGGREGATE> ops.
  _RAGGED_SEGMENT_DOCSTRING = """\
  Computes the %(combination)s along segments of a RaggedTensor.
  
    Returns a RaggedTensor `output` with `num_segments` rows, where the row
    `output[i]` is formed by taking the %(combination)s of all rows of `data`
    whose corresponding `segment_id` is `i`.
  
    The length of the row `output[i]` will be the maximum of the lengths of
    all rows of `data` whose corresponding `segment_id` is `i`.  If no `data`
    rows correspond to a given segment ID, then the output row for that segment
    ID will be empty.
  
    Args:
      data: A `RaggedTensor` containing the values to combine.
      segment_ids: A `Tensor` or `RaggedTensor`.  Must have type `int64` or
        `int32`.  `segment_ids.shape` must be a prefix of `data.shape`.
        Must be greater than or equal to zero, and less than `num_segments`.
        `segment_ids` is not required to be sorted.
      num_segments: An `int32` or `int64` scalar specifying the number of
        distinct segment ids.
      name: A name prefix for the returned tensor (optional).
    Returns:
      A `RaggedTensor` containing the %(combined)s values.  The returned tensor
      has the same dtype as `data`, and its shape is
      `[num_segments] + data.shape[segment_ids.rank:]`.
    Raises:
      ValueError: If `segment_ids.shape` is not a prefix of `data.shape`.
  """
  
  
  def _ragged_segment_aggregate(unsorted_segment_op,
                                data,
                                segment_ids,
                                num_segments,
                                name=None):
    """Aggregates along segments of a RaggedTensor using `unsorted_segment_op`.
  
    Returns a RaggedTensor `output` with `num_segments` rows, where the row
    `output[i]` is formed by combining all rows of `data` whose corresponding
    `segment_id` is `i`.  The values in each row are combined using
    `unsorted_segment_op`.
  
    The length of the row `output[i]` will be the maximum of the lengths of
    all rows of `data` whose corresponding `segment_id` is `i`.  If no `data`
    rows correspond to a given segment ID, then the output row for that segment
    ID will be empty.
  
    Args:
      unsorted_segment_op: The tensorflow `op` that should be used to combine
        values in each row.  Must have the same signature and basic behavior as
        `unsorted_segment_sum`, `unsorted_segment_max`, etc.
      data: A `RaggedTensor` containing the values to be combined.
      segment_ids: A `Tensor` or `RaggedTensor`.  Must have type `int64` or
        `int32`.  `segment_ids.shape` must be a prefix of `data.shape`.
        `segment_ids` is not required to be sorted.
      num_segments: An `int32` or `int64` scalar.
      name: A name prefix for the returned tensor (optional).
  
    Returns:
      A `RaggedTensor` containing the aggregated values.  The returned tensor
      has the same dtype as `data`, and its shape is
      `[num_segments] + data.shape[segment_ids.rank:]`.
    Raises:
      ValueError: If segment_ids.shape is not a prefix of data.shape.
    """
    if not (ragged_tensor.is_ragged(data) or
            ragged_tensor.is_ragged(segment_ids)):
      return unsorted_segment_op(data, segment_ids, num_segments, name)
  
    with ops.name_scope(name, 'RaggedSegment',
                        [data, segment_ids, num_segments]) as name:
      data = ragged_tensor.convert_to_tensor_or_ragged_tensor(data, name='data')
      segment_ids = ragged_tensor.convert_to_tensor_or_ragged_tensor(
          segment_ids, name='segment_ids')
  
      if ragged_tensor.is_ragged(segment_ids):
        if not ragged_tensor.is_ragged(data):
          raise ValueError('segment_ids.shape must be a prefix of data.shape, '
                           'but segment_ids is ragged and data is not.')
        check_splits = check_ops.assert_equal(
            segment_ids.row_splits,
            data.row_splits,
            message='segment_ids.shape must be a prefix of data.shape')
        with ops.control_dependencies([check_splits]):
          return _ragged_segment_aggregate(unsorted_segment_op, data.values,
                                           segment_ids.values, num_segments, name)
  
      segment_ids = math_ops.cast(segment_ids, dtypes.int64)
  
      # Find the length of each row in data.  (dtype=int64, shape=[data_nrows])
      data_row_lengths = data.row_splits[1:] - data.row_splits[:-1]
  
      # Find the length that each output row will have.  The length of the row
      # corresponding to segment `id` is `max(data_row_lengths[i])` where
      # `segment_ids[i]=id`.  (dtype=int64, shape=[output_nrows])
      output_row_lengths = math_ops.maximum(
          math_ops.unsorted_segment_max(data_row_lengths, segment_ids,
                                        num_segments), 0)
      assert output_row_lengths.dtype == dtypes.int64
  
      # Build the splits tensor for the output RaggedTensor.
      output_splits = array_ops.concat([
          array_ops.zeros([1], dtypes.int64),
          math_ops.cumsum(output_row_lengths)
      ],
                                       axis=0)
  
      # For each row in `data`, find the start & limit position where that row's
      # values will be aggregated in output.values.
      data_row_to_out_row_start = array_ops.gather(output_splits, segment_ids)
      data_row_to_out_row_limit = data_row_to_out_row_start + data_row_lengths
  
      # For each value in `data.values`, find the position where it will
      # aggregated in `output.values`.
      # Get the target output values index for each data values index.
      data_val_to_out_val_index = range(data_row_to_out_row_start,
                                        data_row_to_out_row_limit).values
  
      # Recursively aggregate the values.
      output_values = _ragged_segment_aggregate(unsorted_segment_op, data.values,
                                                data_val_to_out_val_index,
                                                output_splits[-1])
      return ragged_tensor.RaggedTensor.from_row_splits(output_values,
                                                        output_splits)
  
  
  def segment_sum(data, segment_ids, num_segments, name=None):
    # For docs, see: _RAGGED_SEGMENT_DOCSTRING
    return _ragged_segment_aggregate(math_ops.unsorted_segment_sum, data,
                                     segment_ids, num_segments, name or
                                     'RaggedSegmentSum')
  
  
  def segment_prod(data, segment_ids, num_segments, name=None):
    # For docs, see: _RAGGED_SEGMENT_DOCSTRING
    return _ragged_segment_aggregate(math_ops.unsorted_segment_prod, data,
                                     segment_ids, num_segments, name or
                                     'RaggedSegmentProd')
  
  
  def segment_min(data, segment_ids, num_segments, name=None):
    # For docs, see: _RAGGED_SEGMENT_DOCSTRING
    return _ragged_segment_aggregate(math_ops.unsorted_segment_min, data,
                                     segment_ids, num_segments, name or
                                     'RaggedSegmentMin')
  
  
  def segment_max(data, segment_ids, num_segments, name=None):
    # For docs, see: _RAGGED_SEGMENT_DOCSTRING
    return _ragged_segment_aggregate(math_ops.unsorted_segment_max, data,
                                     segment_ids, num_segments, name or
                                     'RaggedSegmentMax')
  
  
  def segment_mean(data, segment_ids, num_segments, name=None):
    """For docs, see: _RAGGED_SEGMENT_DOCSTRING."""
    with ops.name_scope(name, 'RaggedSegmentMean',
                        [data, segment_ids, num_segments]):
      total = segment_sum(data, segment_ids, num_segments)
      ones = ragged_tensor.RaggedTensor.from_nested_row_splits(
          array_ops.ones_like(data.flat_values), data.nested_row_splits)
      count = segment_sum(ones, segment_ids, num_segments)
      if ragged_tensor.is_ragged(total):
        return total.with_flat_values(total.flat_values / count.flat_values)
      else:
        return total / count
  
  
  def segment_sqrt_n(data, segment_ids, num_segments, name=None):
    """For docs, see: _RAGGED_SEGMENT_DOCSTRING."""
    with ops.name_scope(name, 'RaggedSegmentSqrtN',
                        [data, segment_ids, num_segments]):
      total = segment_sum(data, segment_ids, num_segments)
      ones = ragged_tensor.RaggedTensor.from_nested_row_splits(
          array_ops.ones_like(data.flat_values), data.nested_row_splits)
      count = segment_sum(ones, segment_ids, num_segments)
      if ragged_tensor.is_ragged(total):
        return total.with_flat_values(
            total.flat_values / math_ops.sqrt(count.flat_values))
      else:
        return total / math_ops.sqrt(count)
  
  
  def _set_ragged_segment_docstring(func, combination, combined):
    func.__doc__ = _RAGGED_SEGMENT_DOCSTRING % dict(
        combination=combination, combined=combined)
  
  
  _set_ragged_segment_docstring(segment_sum, 'sum', 'summed')
  _set_ragged_segment_docstring(segment_prod, 'product', 'multiplied')
  _set_ragged_segment_docstring(segment_min, 'minimum', 'minimized')
  _set_ragged_segment_docstring(segment_max, 'maximum', 'maximized')
  _set_ragged_segment_docstring(segment_mean, 'mean', 'averaged')
  _set_ragged_segment_docstring(segment_sqrt_n, 'sum divided by sqrt(N)',
                                'summed')
  
  #===============================================================================
  # ragged_reduce_<AGGREGATE>
  #===============================================================================
  
  # Docstring template used for ragged_reduce_<AGGREGATE> ops.
  _RAGGED_REDUCE_DOCSTRING = """\
  Computes the %(combination)s of elements across dimensions of a `RaggedTensor`.
  
    Reduces `input_tensor` along the dimensions given in `axis` by taking the
    %(combination)s of values.  If a reduced dimension has no elements for
    some index, then the value for that index will be %(default)s.
  
    The rank of the tensor is reduced by `1` for each entry in `axis`.  If
    `axis` is not specified, then all dimensions are reduced, and a scalar
    value is returned.
    Args:
      input_tensor: A `RaggedTensor` containing the values to be %(combined)s.
      axis: The dimensions to reduce.  May be `None` (to reduce all axes), an
        `int` (to reduce a single axis), a `list` or `tuple` of `int` (to reduce
        a given set of axes), or a `Tensor` with a constant value.  Must be in
        the range `[0, input_tensor.rank]`.
      name: A name prefix for the returned tensor (optional).
    Returns:
      A `RaggedTensor` containing the %(combined)s values.  The returned tensor
      has the same dtype as `data`, and its shape is given by removing the
      dimensions specified in `axis` from `input_tensor.shape`.  The `ragged_rank`
      of the returned tensor is given by substracting any ragged dimensions
      specified in `axis` from `input_tensor.ragged_rank`.
    Raises:
      ValueError: If `axis` contains a `Tensor` whose value is not constant.
    ####Example:
      ```python%(example)s    ```
  """
  _RAGGED_REDUCE_SUM_EXAMPLE = """
      >>> rt = ragged.constant([[3, 1, 4], [1, 5], [9], [2, 6]])
      >>> ragged.reduce_sum(rt, axis=0).eval().tolist()
      [15, 12, 4]  # = [3+1+9+2, 1+5+6, 4]
      >>> ragged.reduce_sum(rt, axis=1).eval().tolist()
      [8, 6, 9, 8]  # = [3+1+4, 1+5, 9, 2+6]
  """
  _RAGGED_REDUCE_PROD_EXAMPLE = """
      >>> rt = ragged.constant([[3, 1, 4], [1, 5], [9], [2, 6]])
      >>> ragged.reduce_prod(rt, axis=0).eval().tolist()
      [54, 30, 4]  # = [3*1*9*2, 1*5*6, 4]
      >>> ragged.reduce_prod(rt, axis=1).eval().tolist()
      [12, 5, 9, 12]  # = [3*1*4, 1*5, 9, 2*6]
  """
  _RAGGED_REDUCE_MIN_EXAMPLE = """
      >>> rt = ragged.constant([[3, 1, 4], [1, 5], [9], [2, 6]])
      >>> ragged.reduce_min(rt, axis=0).eval().tolist()
      [1, 1, 4]  # = [min(3, 1, 9, 2), min(1, 5, 6), 4]
      >>> ragged.reduce_min(rt, axis=1).eval().tolist()
      [1, 1, 9, 2]  # = [min(3, 1, 4), min(1, 5), 9, min(2, 6)]
  """
  _RAGGED_REDUCE_MAX_EXAMPLE = """
      >>> rt = ragged.constant([[3, 1, 4], [1, 5], [9], [2, 6]])
      >>> ragged.reduce_max(rt, axis=0).eval().tolist()
      [9, 6, 4]  # = [max(3, 1, 9, 2), max(1, 5, 6), 4]
      >>> ragged.reduce_max(rt, axis=1).eval().tolist()
      [4, 5, 9, 6]  # = [max(3, 1, 4), max(1, 5), 9, max(2, 6)]
  """
  _RAGGED_REDUCE_MEAN_EXAMPLE = """
      >>> rt = ragged.constant([[3, 1, 4], [1, 5], [9], [2, 6]])
      >>> ragged.reduce_mean(rt, axis=0).eval().tolist()
      [3.75, 4, 4]  # = [mean(3, 1, 9, 2), mean(1, 5, 6), 4]
      >>> ragged.reduce_mean(rt, axis=1).eval().tolist()
      [2.66666, 3, 9, 4]  # = [mean(3, 1, 4), mean(1, 5), 9, mean(2, 6)]
  """
  _RAGGED_REDUCE_ALL_EXAMPLE = """
      >>> rt = ragged.constant([[True, True], [True, True, False, True], [False, True]])
      >>> ragged.reduce_all(rt, axis=0).eval().tolist()
      [False, True, False, True]
      >>> ragged.reduce_all(rt, axis=1).eval().tolist()
      [True, False, False]
  """
  _RAGGED_REDUCE_ANY_EXAMPLE = """
      >>> rt = ragged.constant([[True, True], [True, True, False, True], [False, True]])
      >>> ragged.reduce_any(rt, axis=0).eval().tolist()
      [True, True, False, True]
      >>> ragged.reduce_any(rt, axis=1).eval().tolist()
      [True, True, True]
  """
  
  
  def _ragged_reduce_aggregate(reduce_op,
                               unsorted_segment_op,
                               rt_input,
                               axis,
                               keepdims,
                               name=None):
    """Aggregates across axes of a RaggedTensor using the given `Tensor` ops.
  
    Reduces `rt_input` along the dimensions given in `axis`.  The rank of the
    tensor is reduced by 1 for each entry in `axis`.  If `axis` is not specified,
    then all dimensions are reduced, and a scalar value is returned.
  
    This op assumes that `reduce_op` and `unsorted_segment_op` are associative;
    if not, then reducing multiple axes will return incorrect results.  (In
    particular, reducing multiple axes is currently implemented by reducing the
    axes one at a time.)
  
    Args:
      reduce_op: The tensorflow `op` that should be used to reduce values in
        uniform dimensions.  Must have the same signature and basic behavior as
        `reduce_sum`, `reduce_max`, etc.
      unsorted_segment_op: The tensorflow `op` that should be used to combine
        values in ragged dimensions.  Must have the same signature and basic
        behavior as `unsorted_segment_sum`, `unsorted_segment_max`, etc.
      rt_input: A `Tensor` or `RaggedTensor` containing the values to be reduced.
      axis: The axis or axes to reduce.  May be `None` (to reduce all axes), an
        `int` (to reduce a single axis), a `list` or `tuple` of `int` (to reduce a
        given set of axes), or a `Tensor` with a constant value.  Must be in the
        range `[0, rt_input.rank)`.
      keepdims: If true, retains reduced dimensions with length 1.
      name: A name prefix for the returned tensor (optional).
  
    Returns:
      A `RaggedTensor` containing the reduced values.  The returned tensor
      has the same dtype as `data`, and its shape is given by removing the
      dimensions specified in `axis` from `rt_input.shape`.  The `ragged_rank`
      of the returned tensor is given by substracting any ragged dimensions
      specified in `axis` from `rt_input.ragged_rank`.
    Raises:
      ValueError: If `axis` contains a `Tensor` whose value is not constant.
    """
    if not ragged_tensor.is_ragged(rt_input):
      return reduce_op(rt_input, axis, name=name)
  
    if keepdims:
      raise ValueError('keepdims=True is not supported for RaggedTensors.')
  
    if isinstance(axis, ops.Tensor):
      axis = tensor_util.constant_value(axis)
      if axis is None:
        raise ValueError('axis must be known at graph construction time.')
      if isinstance(axis, np.ndarray):
        axis = axis.tolist()
  
    # When reducing all axes, just ignore splits & reduce the inner values.
    if axis is None:
      return reduce_op(rt_input.flat_values, None, name=name)
  
    with ops.name_scope(name, 'RaggedReduce', [rt_input, axis]):
      if isinstance(axis, (tuple, list)):
        if not axis:
          return rt_input
        elif len(axis) == 1:
          axis = axis[0]
        else:
          # When reducing multiple axes, just reduce one at a time.  This is less
          # efficient, and only works for associative ops.  (In particular, it
          # does not work for reduce_mean.)  However, reducing multiple axes at
          # once will probably require a nontrivial c++ op.
          axis = sorted(axis)
          inner_reduced = _ragged_reduce_aggregate(reduce_op, unsorted_segment_op,
                                                   rt_input, axis[-1], keepdims)
          return _ragged_reduce_aggregate(reduce_op, unsorted_segment_op,
                                          inner_reduced, axis[:-1], keepdims)
  
      rt_input = ragged_tensor.convert_to_tensor_or_ragged_tensor(
          rt_input, name='rt_input')
  
      axis = ragged_util.get_positive_axis(axis, rt_input.shape.ndims)
  
      if axis == 0:
        # out[i_1, i_2, ..., i_N] = sum_{j} rt_input[j, i_1, i_2, ..., i_N]
        row_lengths = rt_input.row_splits[1:] - rt_input.row_splits[:-1]
        num_segments = math_ops.maximum(math_ops.reduce_max(row_lengths), 0)
        segment_ids = range(row_lengths).values
        return _ragged_segment_aggregate(unsorted_segment_op, rt_input.values,
                                         segment_ids, num_segments)
      elif axis == 1:
        # out[i_0, i_1, i_2, ..., i_N] = sum_{j} rt_input[i_0, j, i_2, ..., i_N]
        num_segments = array_ops.shape(rt_input.row_splits)[0] - 1
        segment_ids = segment_id_ops.row_splits_to_segment_ids(
            rt_input.row_splits)
        return _ragged_segment_aggregate(unsorted_segment_op, rt_input.values,
                                         segment_ids, num_segments)
      else:
        # out[i_0, ..., i_[axis-1], i_axis+1], ..., i_N] =
        #     sum_{j} rt_input [i_0, ..., i_[axis-1], j, i_axis+1], ..., i_N]
        return rt_input.with_values(
            _ragged_reduce_aggregate(reduce_op, unsorted_segment_op,
                                     rt_input.values, axis - 1, keepdims))
  
  
  def reduce_sum(input_tensor, axis=None, keepdims=None, name=None):
    """For docs, see: _RAGGED_REDUCE_DOCSTRING."""
    return _ragged_reduce_aggregate(math_ops.reduce_sum,
                                    math_ops.unsorted_segment_sum, input_tensor,
                                    axis, keepdims, name or 'RaggedReduceSum')
  
  
  def reduce_prod(input_tensor, axis=None, keepdims=None, name=None):
    """For docs, see: _RAGGED_REDUCE_DOCSTRING."""
    return _ragged_reduce_aggregate(math_ops.reduce_prod,
                                    math_ops.unsorted_segment_prod, input_tensor,
                                    axis, keepdims, name or 'RaggedReduceProd')
  
  
  def reduce_min(input_tensor, axis=None, keepdims=None, name=None):
    """For docs, see: _RAGGED_REDUCE_DOCSTRING."""
    return _ragged_reduce_aggregate(math_ops.reduce_min,
                                    math_ops.unsorted_segment_min, input_tensor,
                                    axis, keepdims, name or 'RaggedReduceMin')
  
  
  def reduce_max(input_tensor, axis=None, keepdims=None, name=None):
    """For docs, see: _RAGGED_REDUCE_DOCSTRING."""
    return _ragged_reduce_aggregate(math_ops.reduce_max,
                                    math_ops.unsorted_segment_max, input_tensor,
                                    axis, keepdims, name or 'RaggedReduceMax')
  
  
  def reduce_mean(input_tensor, axis=None, keepdims=None, name=None):
    """For docs, see: _RAGGED_REDUCE_DOCSTRING."""
    with ops.name_scope(name, 'RaggedReduceMean', [input_tensor, axis]):
      total = reduce_sum(input_tensor, axis, keepdims)
      if ragged_tensor.is_ragged(input_tensor):
        ones = ragged_tensor.RaggedTensor.from_nested_row_splits(
            array_ops.ones_like(input_tensor.flat_values),
            input_tensor.nested_row_splits)
      else:
        ones = array_ops.ones_like(input_tensor)
      count = reduce_sum(ones, axis, keepdims)
      if ragged_tensor.is_ragged(total):
        return ragged_tensor.RaggedTensor.from_nested_row_splits(
            total.flat_values / count.flat_values, total.nested_row_splits)
      else:
        return total / count
  
  
  def _cast(input_tensor, dtype):
    return ragged_functional_ops.map_flat_values(math_ops.cast, input_tensor,
                                                 dtype)
  
  
  def reduce_all(input_tensor, axis=None, keepdims=None, name=None):
    """For docs, see: _RAGGED_REDUCE_DOCSTRING."""
    with ops.name_scope(name, 'RaggedReduceAll', [input_tensor, axis]):
      return _cast(
          reduce_prod(_cast(input_tensor, dtypes.int32), axis, keepdims),
          dtypes.bool)
  
  
  def reduce_any(input_tensor, axis=None, keepdims=None, name=None):
    """For docs, see: _RAGGED_REDUCE_DOCSTRING."""
    with ops.name_scope(name, 'RaggedReduceAny', [input_tensor, axis]):
      return _cast(
          reduce_sum(_cast(input_tensor, dtypes.int32), axis, keepdims),
          dtypes.bool)
  
  
  def _set_ragged_reduce_docstring(func, combination, combined, default, example):
    func.__doc__ = _RAGGED_REDUCE_DOCSTRING % dict(
        combination=combination,
        combined=combined,
        default=default,
        example=example)
  
  
  _set_ragged_reduce_docstring(reduce_sum, 'sum', 'summed', '0',
                               _RAGGED_REDUCE_SUM_EXAMPLE)
  _set_ragged_reduce_docstring(reduce_prod, 'product', 'multiplied', '1',
                               _RAGGED_REDUCE_PROD_EXAMPLE)
  _set_ragged_reduce_docstring(reduce_min, 'minimum', 'minimized',
                               '`input_tensor.dtype.min`',
                               _RAGGED_REDUCE_MIN_EXAMPLE)
  _set_ragged_reduce_docstring(reduce_max, 'maximum', 'maximized',
                               '`input_tensor.dtype.max`',
                               _RAGGED_REDUCE_MAX_EXAMPLE)
  _set_ragged_reduce_docstring(reduce_mean, 'mean', 'averaged', 'NaN',
                               _RAGGED_REDUCE_MEAN_EXAMPLE)
  
  _set_ragged_reduce_docstring(reduce_all, 'logical and', 'and-ed', 'True',
                               _RAGGED_REDUCE_ALL_EXAMPLE)
  _set_ragged_reduce_docstring(reduce_any, 'logical or', 'or-ed', 'False',
                               _RAGGED_REDUCE_ANY_EXAMPLE)
   # Copyright 2016 The TensorFlow Authors. All Rights Reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # ==============================================================================
  """This module contains the user-facing API for AutoGraph."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import functools
  import sys
  
  from enum import Enum
  
  # pylint:disable=g-bad-import-order
  import numpy as np
  # pylint:enable=g-bad-import-order
  
  
  from tensorflow.python.autograph.core import config
  from tensorflow.python.autograph.core import converter
  from tensorflow.python.autograph.impl import conversion
  from tensorflow.python.autograph.operators import py_builtins
  from tensorflow.python.autograph.pyct import compiler
  from tensorflow.python.autograph.pyct import inspect_utils
  from tensorflow.python.autograph.utils import py_func
  from tensorflow.python.data.util import nest
  from tensorflow.python.framework import tensor_util
  from tensorflow.python.platform import tf_logging as logging
  from tensorflow.python.util import tf_decorator
  from tensorflow.python.util import tf_inspect
  
  # TODO(mdan): Properly document the type hints.
  # TODO(mdan): Reduce the type hint information to (module, type).
  # (currently we require (module + class name, type))
  
  
  # TODO(mdan): This should behave like to_graph (e.g. convert statically).
  # TODO(znado): Make an alias so can write Verbosity directly without needing
  # to write converter.
  def convert(
      recursive=False,
      verbose=converter.Verbosity.BRIEF,
      optional_features=converter.Feature.ALL):
    """Decorator that compiles a function to use TensorFlow ops.
  
    The decorator is dynamic - it recompiles the target whenever the decorated
    function is called. This means the parameter values are known at conversion.
    It also means that repeated calls with different types of parameters will be
    correctly processed.
  
    Args:
      recursive: bool, whether to recursively convert any functions or classes
        that the converted function may use.
      verbose: converter.Verbosity, the level of verbosity.
      optional_features: converted.Feature, allows toggling optional or
        experimental features. When set to None, only the core features are
        enabled.
  
    Returns:
      Callable, a decorator that converts the given function into an equivalent
      function that uses TensorFlow ops.
    """
  
    def decorator(f):
      """Decorator implementation."""
  
      @functools.wraps(f)
      def wrapper(*args, **kwargs):
        return converted_call(
            f, None,
            converter.ConversionOptions(
                recursive=recursive,
                verbose=verbose,
                force_conversion=True,
                optional_features=optional_features,
            ), *args, **kwargs)
  
      wrapper = tf_decorator.make_decorator(f, wrapper)
  
      # Sometimes the decorator is just desugared, making it impossible to detect.
      # This attribute makes detection easier.
      setattr(wrapper, '__ag_compiled', True)
      return wrapper
  
    return decorator
  
  
  class RunMode(Enum):
    """Specifies the way a converted function or method should be executed in TF.
  
    Attributes:
     * GRAPH: Call this function directly, as-is. This is suitable for functions
         that were already designed for TF graphs and contain ops.
     * PY_FUNC: Wrap this function into a py_func op. This is suitable for code
         that will only run correctly in Python, for example code that renders
         to the display, reads keyboard input, etc.
    """
    GRAPH = 1
    PY_FUNC = 2
  
  
  def do_not_convert(run_as=RunMode.GRAPH, return_dtypes=None):
    """Decorator that suppresses the conversion of a function.
  
    See also: docs/pyfunc_dtypes.md
  
    Args:
      run_as: RunMode, specifies how to use the function in TensorFlow.
      return_dtypes: Optional[Iterable[ Union[tf.DType,
        utils.py_func.MatchDType]]], the return data types of the converted
        function, if run_as is RunMode.PY_FUNC. Ignored otherwise. May be set to
        None if the function has no return values.
  
    Returns:
      Callable, a decorator that wraps the original function.
    """
  
    def decorator(f):
      """Decorator implementation."""
  
      @functools.wraps(f)
      def graph_wrapper(*args, **kwargs):
        return f(*args, **kwargs)
  
      @functools.wraps(f)
      def py_func_wrapper(*args, **kwargs):
        if kwargs:
          raise NotImplementedError('RunMode.PY_FUNC does not yet support kwargs')
        # TODO(mdan): Add support for kwargs.
        return py_func.wrap_py_func(
            f, return_dtypes, args, kwargs, use_dummy_return=not return_dtypes)
  
      if run_as == RunMode.GRAPH:
        wrapper = graph_wrapper
      elif run_as == RunMode.PY_FUNC:
        wrapper = py_func_wrapper
      else:
        raise ValueError('unknown value for run_as: %s' % run_as)
  
      # Sometimes the decorator is just desugared, making it impossible to detect.
      # This attribute makes detection easier.
      setattr(wrapper, '__ag_compiled', True)
      return wrapper
  
    return decorator
  
  
  # TODO(mdan): Move to a private, undocumented module.
  def converted_call(f, owner, options, *args, **kwargs):
    """Compiles a function call inline. For internal use only."""
    if options.verbose >= converter.Verbosity.VERBOSE:
      logging.info('Converted call: {}; owner: {}'.format(f, owner))
  
    if owner is not None:
      if not isinstance(f, str):
        raise ValueError(
            'When owner is specified, the function name must be specified as'
            ' a string: {}'.format(f))
  
      # Special case when the owner is a 'super' object. In that case lookups of
      # dynamic attributes won't work. See
      # inspect_utils.SuperWrapperForDynamicAttrs.
      if isinstance(owner, super):
        owner = inspect_utils.SuperWrapperForDynamicAttrs(owner)
  
      f = getattr(owner, f)
  
    if inspect_utils.isbuiltin(f):
      return py_builtins.overload_of(f)(*args, **kwargs)
  
    # TODO(mdan): This needs cleanup.
    # In particular, we may want to avoid renaming functions altogether.
    if not options.force_conversion and conversion.is_whitelisted_for_graph(f):
  
      # Args typically include `self`, as required by the conversion process.
      # When conversion is skipped, `self` is not necessary, because the
      # original bound method is being executed. This code removes it.
      if tf_inspect.ismethod(f) and args:
        f_class = inspect_utils.getmethodclass(f)
        if args[0] is f_class:
          args = args[1:]
  
      return f(*args, **kwargs)
  
    # internal_convert_user_code is for example turned off when issuing a dynamic
    # call conversion from generated code while in nonrecursive mode. In that
    # case we evidently don't want to recurse, but we still have to convert
    # things like builtins.
    if not options.internal_convert_user_code:
      return f(*args, **kwargs)
  
    # Unwrap functools.partial objects
    # TODO(allenl, mdan): Consider sharing unwrapping logic with tf_inspect.
    while isinstance(f, functools.partial):
      args = f.args + args
      new_kwargs = {}
      if f.keywords is not None:
        new_kwargs.update(f.keywords)
      new_kwargs.update(kwargs)
      kwargs = new_kwargs
      f = f.func
  
    if tf_inspect.isfunction(f) or tf_inspect.ismethod(f):
      # Regular functions
      target_entity = f
      arg_map_target = f
      f_class = inspect_utils.getmethodclass(f)
  
      # TODO(b/119246461): This may be more elegantly handled using __get__?
      if f_class is not None:
        # If this is a method call, it may or may not include self.
        #
        # Example when self is included:
        #   converted_call(to_graph(foo.bar), foo)
        #
        # Example when self is not included:
        #   super(...).foo(args)
        #
        if owner is not None and (not args or args[0] is not owner):
          effective_args = (owner,) + args
        else:
          # When the owner is not specified, use the result of
          # inspect_utils.getmethodclass.
          # TODO(b/119246461): Make sure an owner is always specified.
          if not args or args[0] is not f_class:
            effective_args = (f_class,) + args
          else:
            effective_args = (f_class,) + args[1:]
        partial_types = (f_class,)
      else:
        effective_args = args
        partial_types = ()
  
    elif tf_inspect.isclass(f):
      # Constructors
      target_entity = f
      arg_map_target = f.__init__
      effective_args = args
      partial_types = ()
  
    elif hasattr(f, '__call__') and hasattr(f, '__class__'):
      # Callable objects
      target_entity = f.__call__
      arg_map_target = f.__call__
      effective_args = (f,) + args
      partial_types = (f.__class__,)
  
    else:
      NotImplementedError('unknown callable type "%s"' % type(f))
  
    arg_values = tf_inspect.getcallargs(arg_map_target, *args, **kwargs)
    arg_types = {}
    for name, arg in arg_values.items():
      arg_class = arg.__class__
      arg_types[name] = (arg_class.__name__, arg_class)
  
    # When called from within a decorator, this is the only indication that
    # the function is a method - it appears that the decorator is applied
    # before the method is bound.
    if not partial_types:
      if 'self' in arg_values:
        if tf_inspect.isclass(arg_values['self'].__class__):
          partial_types = (arg_values['self'].__class__,)
      elif 'cls' in arg_values:
        if tf_inspect.isclass(arg_values['cls']):
          partial_types = (arg_values['cls'],)
  
    converted_f = to_graph(
        target_entity,
        recursive=options.recursive,
        verbose=options.verbose,
        arg_values=arg_values,
        arg_types=arg_types,
        partial_types=partial_types,
        strip_decorators=options.strip_decorators,
        optional_features=options.optional_features)
  
    result = converted_f(*effective_args, **kwargs)
  
    # The converted function's closure is simply inserted into the function's
    # module __dict__. Since modules are permanently cached, that results in
    # leaking the entire closure.
    # Normally, it's not safe to delete the module because that may release said
    # closure as well. However, in the case of converted_call we are certain the
    # function will not be executed again, so the closure should no longer be
    # needed so long as the function doesn't return any executable code.
    # TODO(mdan): Attach the closure properly, using cells.
    if all(map(_is_not_callable, nest.flatten(result))):
      del sys.modules[converted_f.__module__]
  
    return result
  
  
  def _is_not_callable(obj):
    # TODO(brianklee): Handle case when obj is a tensor dependent on a py_func.
    if isinstance(obj, (int, float, complex, str, bool)):
      return True
    if isinstance(obj, (np.ndarray, np.generic)):
      return True
    if tensor_util.is_tensor(obj):
      return True
    return False
  
  
  # TODO(mdan): Rename: to_ops?
  # TODO(mdan): Look into overloading as function and decorator, like tfe.defun?
  # TODO(mdan): Remove partial_types.
  def to_graph(e,
               recursive=True,
               verbose=converter.Verbosity.VERBOSE,
               arg_values=None,
               arg_types=None,
               partial_types=None,
               strip_decorators=None,
               optional_features=converter.Feature.ALL):
    """Converts a Python entity into equivalent code that uses TensorFlow ops.
  
    Supported Python entities include:
      * functions
      * classes
  
    Classes are converted by converting all their methods into a new class.
  
    Args:
      e: Union[Callable, Type], the Python entity to convert.
      recursive: bool, whether to recursively convert any functions that the
        converted function may call.
      verbose: converter.Verbosity, the level of printing verbosity to use.
      arg_values: Optional[Dict[Text, Any]], value hints for symbols including
        function arguments.
      arg_types: Optional[Dict[Text, Type]], type hints for symbols including
        function arguments.
      partial_types: Set[Type], reserved for internal use.
      strip_decorators: Tuple[Callable], same as
        ConversionOptions.strip_decorators.
      optional_features: Union[Feature, Set[Feature]], same as
        ConversionOptions.optional_features.
  
    Returns:
      Union[Callable, Type], the converted entity, which is the same kind as e
      (that is, a function is e is a function, a class if e is a class, etc.) but
      its code has been converted to use TF ops.
  
    Raises:
      ValueError: If the entity could not be converted.
    """
    if strip_decorators is None:
      strip_decorators = ()
    strip_decorators += (convert, do_not_convert, converted_call)
  
    program_ctx = converter.ProgramContext(
        options=converter.ConversionOptions(
            recursive=recursive,
            verbose=verbose,
            strip_decorators=strip_decorators,
            optional_features=optional_features),
        partial_types=partial_types,
        autograph_module=tf_inspect.getmodule(to_graph),
        uncompiled_modules=config.DEFAULT_UNCOMPILED_MODULES)
    _, name, namespace = conversion.entity_to_graph(e, program_ctx, arg_values,
                                                    arg_types)
  
    nodes = []
    for dep in reversed(program_ctx.conversion_order):
      nodes.extend(program_ctx.dependency_cache[dep])
  
    compiled_module, _ = compiler.ast_to_object(
        nodes,
        source_prefix=program_ctx.required_imports,
        include_source_map=True)
  
    # The compiled code should see everything the entry entity saw.
    # TODO(mdan): This might not work well if the call tree spans modules?
    for key, val in namespace.items():
      # Avoid overwriting entities that have been transformed.
      if key not in compiled_module.__dict__:
        compiled_module.__dict__[key] = val
    compiled = getattr(compiled_module, name)
  
    if tf_inspect.isfunction(e):
      compiled.__defaults__ = e.__defaults__
  
    if hasattr(compiled, '__globals__'):
      # Remove self to avoid circular references. This will probably only work
      # so long as the function is not reentrant.
      del compiled.__globals__[name]
  
    # Need this so the source_mapping attribute is available for the context
    # manager to access for runtime errors.
    #
    # Note that compiler.ast_to_object attaches the source map 'ag_source_map__'
    # symbol to the compiled module.
    # TODO(mdan): Record this statically in the generated code.
    # TODO(mdan): Rename this attribute to 'autograph_info__'
    source_map_attribute_name = 'ag_source_map'
    if getattr(compiled, source_map_attribute_name, None) is not None:
      raise ValueError('cannot convert %s because is has an attribute '
                       '"%s", which is reserved for AutoGraph.' %
                       (compiled, source_map_attribute_name))
    setattr(compiled, source_map_attribute_name,
            compiled_module.__dict__['ag_source_map__'])
  
    return compiled
  
  
  def to_code(e,
              recursive=True,
              arg_values=None,
              arg_types=None,
              partial_types=None,
              indentation='  '):
    """Returns the equivalent code that uses TensorFlow ops.
  
    Also see: `to_graph`, `convert`
  
    Args:
      e: Union[Callable, Type], the Python entity to convert.
      recursive: bool, whether to recursively convert any functions that the
        converted function may call.
      arg_values: Optional[Dict[Text, Any]], value hints for symbols including
        function arguments.
      arg_types: Optional[Dict[Text, Type]], type hints for symbols including
        function arguments.
      partial_types: Set[Type], reserved for internal use.
      indentation: Text, when to use for each level of indentation.
  
    Returns:
      Text, the converted code.
    """
    program_ctx = converter.ProgramContext(
        options=converter.ConversionOptions(
            recursive=recursive,
            strip_decorators=(convert, do_not_convert, converted_call)),
        partial_types=partial_types,
        autograph_module=tf_inspect.getmodule(to_graph),
        uncompiled_modules=config.DEFAULT_UNCOMPILED_MODULES)
    conversion.entity_to_graph(e, program_ctx, arg_values, arg_types)
  
    code = '\n'.join(
        compiler.ast_to_source(program_ctx.dependency_cache[dep], indentation)
        for dep in reversed(program_ctx.conversion_order))
  
    return program_ctx.required_imports + '\n\n' + code
  
apis: tensorflow.python.ops.ragged.ragged_util.get_positive_axis tensorflow.python.ops.ragged.ragged_tensor.RaggedTensor.from_row_splits tensorflow.python.ops.ragged.ragged_tensor.RaggedTensor.from_nested_row_splits tensorflow.python.ops.ragged.segment_id_ops.row_splits_to_segment_ids tensorflow.python.framework.tensor_util.constant_value tensorflow.python.ops.array_ops.zeros tensorflow.python.ops.array_ops.shape tensorflow.python.ops.math_ops.unsorted_segment_max tensorflow.python.framework.ops.control_dependencies tensorflow.python.ops.math_ops.cast tensorflow.python.framework.ops.convert_to_tensor tensorflow.python.ops.ragged.ragged_functional_ops.map_flat_values tensorflow.python.ops.array_ops.gather tensorflow.python.ops.ragged.ragged_tensor.convert_to_tensor_or_ragged_tensor tensorflow.python.ops.gen_ragged_math_ops.ragged_range tensorflow.python.ops.math_ops.reduce_max tensorflow.python.ops.math_ops.sqrt tensorflow.python.ops.ragged.ragged_tensor.is_ragged tensorflow.python.framework.ops.name_scope tensorflow.python.ops.check_ops.assert_equal tensorflow.python.ops.array_ops.ones_like tensorflow.python.ops.math_ops.cumsum
apis: tensorflow.python.data.util.nest.flatten tensorflow.python.autograph.impl.conversion.entity_to_graph tensorflow.python.autograph.operators.py_builtins.overload_of tensorflow.python.autograph.pyct.inspect_utils.SuperWrapperForDynamicAttrs tensorflow.python.util.tf_inspect.getmodule tensorflow.python.autograph.pyct.inspect_utils.isbuiltin tensorflow.python.util.tf_inspect.getcallargs tensorflow.python.util.tf_inspect.ismethod tensorflow.python.util.tf_inspect.isfunction tensorflow.python.util.tf_decorator.make_decorator tensorflow.python.autograph.impl.conversion.is_whitelisted_for_graph tensorflow.python.autograph.pyct.inspect_utils.getmethodclass tensorflow.python.util.tf_inspect.isclass tensorflow.python.autograph.pyct.compiler.ast_to_source tensorflow.python.autograph.core.converter.ConversionOptions tensorflow.python.autograph.utils.py_func.wrap_py_func tensorflow.python.autograph.pyct.compiler.ast_to_object tensorflow.python.framework.tensor_util.is_tensor