chore: import upstream snapshot with attribution

tensorflow/python/ops/ragged/__init__.py

@@ -0,0 +1,28 @@
+# 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.
+# ==============================================================================
+"""Ragged Tensors.
+
+This package defines ops for manipulating ragged tensors (`tf.RaggedTensor`),
+which are tensors with non-uniform shapes.  In particular, each `RaggedTensor`
+has one or more *ragged dimensions*, which are dimensions whose slices may have
+different lengths.  For example, the inner (column) dimension of
+`rt=[[3, 1, 4, 1], [], [5, 9, 2], [6], []]` is ragged, since the column slices
+(`rt[0, :]`, ..., `rt[4, :]`) have different lengths.  For a more detailed
+description of ragged tensors, see the `tf.RaggedTensor` class documentation
+and the [Ragged Tensor Guide](/guide/ragged_tensor).
+
+API docstring: tensorflow.ragged
+"""
+from tensorflow.python.ops.ragged import ragged_tensor

tensorflow/python/ops/ragged/convert_to_tensor_or_ragged_tensor_op_test.py

@@ -0,0 +1,234 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_tensor.convert_to_tensor_or_ragged."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedConvertToTensorOrRaggedTensorTest(test_util.TensorFlowTestCase,
+                                              parameterized.TestCase):
+
+  #=============================================================================
+  # Tests where the 'value' param is a RaggedTensor
+  #=============================================================================
+  @parameterized.parameters([
+      dict(pylist=[[1, 2], [3]]),
+      dict(pylist=[[1, 2], [3]], preferred_dtype=dtypes.float32),
+      dict(pylist=[[1, 2], [3]], preferred_dtype=dtypes.string),
+      # Note: Conversion of a single np.array is tested below. These tests
+      # check nestings consisting of multiple or irregularily-shaped np.arrays.
+      dict(
+          pylist=[np.array([1, 2]), np.array([3])],
+          preferred_dtype=dtypes.string),
+      dict(pylist=np.array([[1, 2], [3]], dtype=object),
+           preferred_dtype=dtypes.float32),
+      dict(pylist=np.array([[1, 2], [3]], dtype=object),
+           preferred_dtype=dtypes.string),
+      dict(
+          pylist=[np.array([[1], np.array([2])]), [np.array([3])]],
+          preferred_dtype=dtypes.float32),
+      dict(pylist=[np.array(1)], preferred_dtype=dtypes.string),
+  ])
+  def testConvertRaggedTensor(self, pylist, dtype=None, preferred_dtype=None):
+    rt = ragged_factory_ops.constant(pylist)
+    converted = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        rt, dtype, preferred_dtype)
+    self.assertIs(converted, rt)
+
+  @parameterized.parameters([
+      dict(
+          pylist=[[1, 2], [3, 4]],
+          dtype=dtypes.float32,
+          message=('Tensor conversion requested dtype float32 for '
+                   'RaggedTensor with dtype int32')),
+      dict(
+          pylist=np.array([[1, 2], [3, 4]]),
+          dtype=dtypes.float32,
+          message=('Tensor conversion requested dtype float32 for '
+                   'RaggedTensor with dtype int32')),
+      dict(
+          pylist=[[1, 2], [3, 4]],
+          dtype=dtypes.string,
+          message=('Tensor conversion requested dtype string for '
+                   'RaggedTensor with dtype .*')),
+  ])
+  def testConvertRaggedTensorError(self,
+                                   pylist,
+                                   message,
+                                   dtype=None,
+                                   preferred_dtype=None):
+    rt = ragged_factory_ops.constant(pylist)
+
+    with self.assertRaisesRegex(ValueError, message):
+      ragged_tensor.convert_to_tensor_or_ragged_tensor(rt, dtype,
+                                                       preferred_dtype)
+
+  #=============================================================================
+  # Tests where the 'value' param is a RaggedTensorValue
+  #=============================================================================
+  @parameterized.parameters(
+      [
+          dict(
+              value=ragged_factory_ops.constant_value([[1, 2], [3]],
+                                                      dtype=np.int32),
+              expected_dtype=dtypes.int32),
+          dict(
+              value=ragged_factory_ops.constant_value([[b'a', b'b'], [b'c']]),
+              expected_dtype=dtypes.string),
+          dict(
+              value=ragged_factory_ops.constant_value([[1, 2], [3]],
+                                                      dtype=np.int32),
+              dtype=dtypes.float32,
+              expected_dtype=dtypes.float32),
+          dict(
+              value=ragged_factory_ops.constant_value([[1, 2], [3]],
+                                                      dtype=np.int32),
+              preferred_dtype=dtypes.float32,
+              expected_dtype=dtypes.float32),
+          dict(
+              value=ragged_factory_ops.constant_value([[1, 2], [3]],
+                                                      dtype=np.int32),
+              preferred_dtype=dtypes.string,
+              expected_dtype=dtypes.int32),
+      ])
+  def testConvertRaggedTensorValue(self,
+                                   value,
+                                   dtype=None,
+                                   preferred_dtype=None,
+                                   expected_dtype=None):
+    if expected_dtype is None:
+      expected_dtype = value.dtype if dtype is None else dtype
+    converted = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        value, dtype, preferred_dtype)
+    self.assertEqual(value.ragged_rank, converted.ragged_rank)
+    self.assertEqual(dtypes.as_dtype(expected_dtype), converted.dtype)
+    self.assertAllEqual(value, converted)
+
+  @parameterized.parameters([
+      dict(
+          value=ragged_factory_ops.constant_value([['a', 'b'], ['c']],
+                                                  dtype=str),
+          dtype=dtypes.int32,
+          message=(r"invalid literal for int\(\) with base 10: "
+                   r"('a'|np.str_\('a'\))")),
+  ])
+  def testConvertRaggedTensorValueError(self,
+                                        value,
+                                        message,
+                                        dtype=None,
+                                        preferred_dtype=None):
+    with self.assertRaisesRegex(ValueError, message):
+      ragged_tensor.convert_to_tensor_or_ragged_tensor(value, dtype,
+                                                       preferred_dtype)
+
+  #=============================================================================
+  # Tests where the 'value' param is a Tensor
+  #=============================================================================
+  @parameterized.parameters([
+      dict(pylist=[[1, 2], [3, 4]]),
+      dict(pylist=[[1, 2], [3, 4]], preferred_dtype=dtypes.float32),
+      dict(pylist=[[1, 2], [3, 4]], preferred_dtype=dtypes.string),
+  ])
+  def testConvertTensor(self, pylist, dtype=None, preferred_dtype=None):
+    tensor = constant_op.constant(pylist)
+    converted = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        tensor, dtype, preferred_dtype)
+    self.assertIs(tensor, converted)
+
+  @parameterized.parameters([
+      dict(
+          pylist=[[1, 2], [3, 4]],
+          dtype=dtypes.float32,
+          message=('Tensor conversion requested dtype float32 for '
+                   'Tensor with dtype int32')),
+      dict(
+          pylist=[[1, 2], [3, 4]],
+          dtype=dtypes.string,
+          message=('Tensor conversion requested dtype string for '
+                   'Tensor with dtype int32')),
+  ])
+  def testConvertTensorError(self,
+                             pylist,
+                             message,
+                             dtype=None,
+                             preferred_dtype=None):
+    tensor = constant_op.constant(pylist)
+    with self.assertRaisesRegex(ValueError, message):
+      ragged_tensor.convert_to_tensor_or_ragged_tensor(tensor, dtype,
+                                                       preferred_dtype)
+
+  #=============================================================================
+  # Tests where the 'value' param is a np.array
+  #=============================================================================
+  @parameterized.parameters([
+      dict(
+          value=np.array([[1, 2], [3, 4]], dtype=np.int32),
+          expected_dtype=dtypes.int32),
+      dict(
+          value=np.array([[b'a', b'b'], [b'c', b'd']]),
+          expected_dtype=dtypes.string),
+      dict(
+          value=np.array([[1, 2], [3, 4]], dtype=np.int32),
+          dtype=dtypes.float32,
+          expected_dtype=dtypes.float32),
+      dict(
+          value=np.array([[1, 2], [3, 4]], dtype=np.int32),
+          preferred_dtype=dtypes.float32,
+          expected_dtype=dtypes.float32),
+      dict(
+          value=np.array([[1, 2], [3, 4]], dtype=np.int32),
+          preferred_dtype=dtypes.string,
+          expected_dtype=dtypes.int32),
+  ])
+  def testConvertNumpyArray(self,
+                            value,
+                            dtype=None,
+                            preferred_dtype=None,
+                            expected_dtype=None):
+    if expected_dtype is None:
+      expected_dtype = value.dtype if dtype is None else dtype
+    converted = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        value, dtype, preferred_dtype)
+    self.assertEqual(dtypes.as_dtype(expected_dtype), converted.dtype)
+    self.assertAllEqual(value, converted)
+
+  @parameterized.parameters([
+      dict(
+          value=np.array([['a', 'b'], ['c', 'd']], dtype=str),
+          dtype=dtypes.int32,
+          message=(r"invalid literal for int\(\) with base 10: "
+                   r"('a'|np.str_\('a'\))")),
+  ])
+  def testConvertNumpyArrayError(self,
+                                 value,
+                                 message,
+                                 dtype=None,
+                                 preferred_dtype=None):
+    with self.assertRaisesRegex(ValueError, message):
+      ragged_tensor.convert_to_tensor_or_ragged_tensor(value, dtype,
+                                                       preferred_dtype)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_autograph.py

@@ -0,0 +1,73 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Autograph-specific overrides for ragged_tensor."""
+from tensorflow.python.autograph.operators import control_flow
+from tensorflow.python.ops import cond as tf_cond
+from tensorflow.python.ops.ragged import ragged_tensor
+
+
+def _tf_ragged_for_stmt(
+    iter_, extra_test, body, get_state, set_state, symbol_names, opts
+):
+  """Overload of for_stmt that iterates over TF ragged tensors."""
+  init_vars = get_state()
+  control_flow.verify_loop_init_vars(init_vars, symbol_names)
+
+  # TODO(mdan): Move this into len()? Requires eager support.
+  if iter_.shape and iter_.shape[0] is not None:
+    n = iter_.shape[0]
+  else:
+    n = iter_.row_lengths()[0]
+
+  iterate_index = 0
+
+  def aug_get_state():
+    return (iterate_index,) + get_state()
+
+  def aug_set_state(aug_loop_vars):
+    nonlocal iterate_index
+    # TODO(b/171479293): Drop the lint override.
+    iterate_index, *loop_vars = aug_loop_vars  # pylint:disable=unused-variable
+    # The iteration index is not "output" by the for loop. If the iteration
+    # index is used outside the loop, it will appear
+    # in the loop vars separately.
+    set_state(loop_vars)
+
+  def aug_body():
+    nonlocal iterate_index
+    body(iter_[iterate_index])
+    iterate_index += 1
+
+  def aug_test():
+    main_test = iterate_index < n
+    if extra_test is not None:
+      return tf_cond.cond(main_test, extra_test, lambda: False)
+    return main_test
+
+  control_flow._add_max_iterations_hint(opts, n)  # pylint: disable=protected-access
+
+  control_flow._tf_while_stmt(  # pylint: disable=protected-access
+      aug_test,
+      aug_body,
+      aug_get_state,
+      aug_set_state,
+      ('<internal iterate>',) + symbol_names,
+      opts,
+  )
+
+
+control_flow.for_loop_registry.register(
+    ragged_tensor.RaggedTensor, _tf_ragged_for_stmt
+)

tensorflow/python/ops/ragged/ragged_batch_gather_op_test.py

@@ -0,0 +1,533 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_batch_gather_ops.batch_gather."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_batch_gather_ops
+from tensorflow.python.ops.ragged import ragged_batch_gather_with_default_op
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedBatchGatherOpTest(test_util.TensorFlowTestCase,
+                              parameterized.TestCase):
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Docstring Example
+      #=========================================================================
+      dict(
+          descr='Docstring example',
+          params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d'], [],
+                                                    ['e']]),
+          indices=ragged_factory_ops.constant_value([[1, 2, 0], [], [], [0,
+                                                                         0]]),
+          expected=ragged_factory_ops.constant_value([[b'b', b'c', b'a'], [],
+                                                      [], [b'e', b'e']])),
+      #=========================================================================
+      # 0 Batch Dimensions
+      #=========================================================================
+      dict(
+          descr='params: [P1], indices: [I], result: [I]',
+          params=['a', 'b', 'c', 'd'],
+          indices=[3, 2],
+          expected=[b'd', b'c']),
+      dict(
+          descr='params: [P1, (P2)], indices: [I], result: [I, (P2)]',
+          params=ragged_factory_ops.constant_value([['a', 'b'], [], ['c'],
+                                                    ['d', 'e']]),
+          indices=[3, 2],
+          expected=ragged_factory_ops.constant_value([[b'd', b'e'], [b'c']])),
+      #=========================================================================
+      # 1 Batch Dimension
+      #=========================================================================
+      dict(
+          descr='params: [B1, P1], indices: [B1, I], result: [B1, I]',
+          params=[['a', 'b', 'c'], ['d', 'e', 'f'], ['g', 'h', 'i']],
+          indices=[[2, 0], [0, 1], [1, 0]],
+          expected=[[b'c', b'a'], [b'd', b'e'], [b'h', b'g']]),
+      dict(
+          descr='params: [B1, (P1)], indices: [B1, I], result: [B1, I]',
+          params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d', 'e'],
+                                                    ['g']]),
+          indices=[[2, 0], [0, 1], [0, 0]],
+          expected=[[b'c', b'a'], [b'd', b'e'], [b'g', b'g']]),
+      dict(
+          descr='params: [B1, P1], indices: [B1, (I)], result: [B1, (I)]',
+          params=[['a', 'b', 'c'], ['d', 'e', 'f'], ['g', 'h', 'i']],
+          indices=ragged_factory_ops.constant_value([[2, 0, 2], [0], [1]]),
+          expected=ragged_factory_ops.constant_value([[b'c', b'a', b'c'],
+                                                      [b'd'], [b'h']])),
+      dict(
+          descr=('params: [B1, (P1), (P2), P3], indices: [B1, I], '
+                 'result: [B1, I, (P2), P3]'),
+          params=ragged_factory_ops.constant_value(
+              [[[['a']], [['b'], ['c']]], [[['d'], ['e']], [['f']]], [[['g']]]],
+              ragged_rank=2),
+          indices=[[1, 0], [0, 1], [0, 0]],
+          expected=ragged_factory_ops.constant_value(
+              [[[[b'b'], [b'c']], [[b'a']]], [[[b'd'], [b'e']], [[b'f']]],
+               [[[b'g']], [[b'g']]]],
+              ragged_rank=2)),
+      #=========================================================================
+      # 2 Batch Dimensions
+      #=========================================================================
+      dict(
+          descr=('params: [B1, B2, P1], indices: [B1, B2, I], '
+                 'result: [B1, B2, I]'),
+          params=[[['a', 'b', 'c']], [['d', 'e', 'f']], [['g', 'h', 'i']]],
+          indices=[[[2, 0]], [[0, 1]], [[1, 0]]],
+          expected=[[[b'c', b'a']], [[b'd', b'e']], [[b'h', b'g']]]),
+      dict(
+          descr=('params: [B1, (B2), P1], indices: [B1, (B2), I], '
+                 'result: [B1, (B2), I]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d', 'e', 'f']], [['g', 'h', 'i']]],
+              ragged_rank=1),
+          indices=ragged_factory_ops.constant_value(
+              [[[2, 0], [0, 1]], [[1, 0]]], ragged_rank=1),
+          expected=ragged_factory_ops.constant_value(
+              [[[b'c', b'a'], [b'd', b'e']], [[b'h', b'g']]], ragged_rank=1)),
+      dict(
+          descr=('params: [B1, (B2), (P1)], indices: [B1, (B2), I], '
+                 'result: [B1, (B2), I]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d']], [['e', 'f']]], ragged_rank=2),
+          indices=ragged_factory_ops.constant_value(
+              [[[2, 0], [0, 0]], [[1, 0]]], ragged_rank=1),
+          expected=ragged_factory_ops.constant_value(
+              [[[b'c', b'a'], [b'd', b'd']], [[b'f', b'e']]], ragged_rank=1)),
+      dict(
+          descr=('params: [B1, (B2), P1], indices: [B1, (B2), (I)], '
+                 'result: [B1, (B2), (I)]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d', 'e', 'f']], [['g', 'h', 'i']]],
+              ragged_rank=1),
+          indices=ragged_factory_ops.constant_value(
+              [[[2, 1, 0], [0]], [[1, 1]]], ragged_rank=2),
+          expected=ragged_factory_ops.constant_value(
+              [[[b'c', b'b', b'a'], [b'd']], [[b'h', b'h']]], ragged_rank=2)),
+      #=========================================================================
+      # 3 Batch Dimensions
+      #=========================================================================
+      dict(
+          descr=(
+              'params: [B1, (B2), (B3), (P1)], indices: [B1, (B2), (B3), I], '
+              'result: [B1, (B2), (B3), I]'),
+          params=ragged_factory_ops.constant_value(
+              [[[['a', 'b', 'c'], ['d']], [['e', 'f']]]], ragged_rank=3),
+          indices=ragged_factory_ops.constant_value(
+              [[[[2, 0], [0, 0]], [[1, 0]]]], ragged_rank=2),
+          expected=ragged_factory_ops.constant_value(
+              [[[[b'c', b'a'], [b'd', b'd']], [[b'f', b'e']]]], ragged_rank=2)),
+  ])
+  def testRaggedBatchGather(self, descr, params, indices, expected):
+    result = ragged_batch_gather_ops.batch_gather(params, indices)
+    self.assertAllEqual(result, expected)
+
+  @parameterized.parameters([
+      # Docstring example:
+      dict(
+          descr='Docstring example',
+          params=[['a', 'b', 'c'], ['d'], [], ['e']],
+          indices=[[1, 2, -1], [], [], [0, 10]],
+          expected=[['b', 'c', 'FOO'], [], [], ['e', 'FOO']],
+          default_value='FOO',
+      ),
+      # Dimensions:
+      # indices: [4]
+      # params: [2, (d1), (d2)]
+      dict(
+          descr='params: [2, (d1), (d2), indices: [4]',
+          indices=[1, 100, 0, -1],
+          params=[[['The', 'deal', 'came', 'about', '18', 'months', 'after',
+                    'Yahoo', '!', 'rejected', 'a', '47.5', '-', 'billion', '-',
+                    'dollar', 'takeover', 'offer', 'from', 'Microsoft', '.'],
+                   ['Trumpty', 'Dumpty', 'sat', 'on', 'a', 'wall']],
+                  [["It's", 'always', 'darkest', 'before', 'the', 'dawn']]],
+          expected=[[["It's", 'always', 'darkest', 'before', 'the', 'dawn']],
+                    [['$NONE^']],
+                    [['The', 'deal', 'came', 'about', '18', 'months', 'after',
+                      'Yahoo', '!', 'rejected', 'a', '47.5', '-', 'billion',
+                      '-', 'dollar', 'takeover', 'offer', 'from', 'Microsoft',
+                      '.'],
+                     ['Trumpty', 'Dumpty', 'sat', 'on', 'a', 'wall']],
+                    [['$NONE^']]],
+      ),
+      # Dimensions:
+      # params: [1, (d1)]
+      # indices: [3]
+      dict(
+          descr='params: rank 2, indices: rank 1',
+          params=[
+              ['Bruce', 'Wayne'],
+          ],
+          indices=[-1, 0, 1000],
+          expected=[['$NONE^'], ['Bruce', 'Wayne'], ['$NONE^']]
+      ),
+      # Dimensions:
+      # params: [1, (d1)]
+      # indices: [1, (d2)]
+      dict(
+          descr='Test underbound indices of shape [1, (d2)]',
+          params=[
+              ['The', 'deal', 'came', 'about', '18', 'months', 'after', 'Yahoo',
+               '!', 'rejected', 'a', '47.5', '-', 'billion', '-', 'dollar',
+               'takeover', 'offer', 'from', 'Microsoft', '.'],
+          ],
+          indices=[[8, -1]],
+          expected=[['!', '$NONE^']],
+      ),
+      dict(
+          descr='Test underbound indices of shape [2, (d2)]',
+          params=[
+              ['The', 'deal', 'came', 'about', '18', 'months', 'after', 'Yahoo',
+               '!', 'rejected', 'a', '47.5', '-', 'billion', '-', 'dollar',
+               'takeover', 'offer', 'from', 'Microsoft', '.'],
+              ['Who', 'let', 'the', 'dogs', 'out', '?'],
+          ],
+          indices=[[8, -1], [1, 100]],
+          expected=[['!', '$NONE^'], ['let', '$NONE^']],
+      ),
+      # Dimensions:
+      # params: [2, (d1)]
+      # indices: [2, (d2)]
+      dict(
+          descr='Test underbound indices of rank 2',
+          params=[
+              ['The', 'deal', 'came', 'about', '18', 'months', 'after', 'Yahoo',
+               '!', 'rejected', 'a', '47.5', '-', 'billion', '-', 'dollar',
+               'takeover', 'offer', 'from', 'Microsoft', '.'],
+              ['He', 'left', 'us', '.', 'Little', 'boys', 'crowded', 'together',
+               'on', 'long', 'wooden', 'benches', ',', 'and', 'in', 'the',
+               'center', 'of', 'the', 'room', 'sat', 'the', 'teacher', '.',
+               'His', 'black', 'beard', 'dripped', 'down', 'over', 'the',
+               'front', 'of', 'his', 'coat', '.', 'One', 'white', 'hand',
+               'poised', 'a', 'stick', 'above', 'his', 'desk', '.', 'He',
+               'turned', 'his', 'surly', ',', 'half', '-', 'closed', 'eyes',
+               'toward', 'us', ',', 'stared', 'for', 'a', 'second', ',', 'then',
+               'shouted', 'in', 'Yiddish', ',', '``', 'One', ',', 'two', ',',
+               'three', "''", '!', '!', 'Rapping', 'the', 'stick', 'against',
+               'the', 'desk', '.', 'The', 'little', 'boys', 'shrilled', 'out',
+               'a', 'Yiddish', 'translation', 'or', 'interpretation', 'of',
+               'the', 'Five', 'Books', 'of', 'Moses', ',', 'which', 'they',
+               'had', 'previously', 'chanted', 'in', 'Hebrew', '.']],
+          indices=[[8, -1], [3, 23, 35, 45, 75, 83, -121]],
+          expected=[['!', '$NONE^'], ['.', '.', '.', '.', '!', '.', '$NONE^']],
+      ),
+      dict(
+          descr='Test overbound indices of rank 2',
+          params=[
+              ['The', 'deal', 'came', 'about', '18', 'months', 'after', 'Yahoo',
+               '!', 'rejected', 'a', '47.5', '-', 'billion', '-', 'dollar',
+               'takeover', 'offer', 'from', 'Microsoft', '.'],
+              ['He', 'left', 'us', '.', 'Little', 'boys', 'crowded', 'together',
+               'on', 'long', 'wooden', 'benches', ',', 'and', 'in', 'the',
+               'center', 'of', 'the', 'room', 'sat', 'the', 'teacher', '.',
+               'His', 'black', 'beard', 'dripped', 'down', 'over', 'the',
+               'front', 'of', 'his', 'coat', '.', 'One', 'white', 'hand',
+               'poised', 'a', 'stick', 'above', 'his', 'desk', '.', 'He',
+               'turned', 'his', 'surly', ',', 'half', '-', 'closed', 'eyes',
+               'toward', 'us', ',', 'stared', 'for', 'a', 'second', ',', 'then',
+               'shouted', 'in', 'Yiddish', ',', '``', 'One', ',', 'two', ',',
+               'three', "''", '!', '!', 'Rapping', 'the', 'stick', 'against',
+               'the', 'desk', '.', 'The', 'little', 'boys', 'shrilled', 'out',
+               'a', 'Yiddish', 'translation', 'or', 'interpretation', 'of',
+               'the', 'Five', 'Books', 'of', 'Moses', ',', 'which', 'they',
+               'had', 'previously', 'chanted', 'in', 'Hebrew', '.']],
+          indices=[[8, 8823], [3, 23, 35, 45, 75, 83, 1234]],
+          expected=[['!', '$NONE^'], ['.', '.', '.', '.', '!', '.', '$NONE^']],
+      ),
+      # Dimensions:
+      # params: [2, (d1), 2]
+      # indices: [2, (d2)]
+      dict(
+          descr='params: rank 3, indices: rank 2',
+          params=[
+              [['The', 'deal'], ['takeover', 'offer'], ['from', 'Microsoft']],
+              [['Who', 'let'], ['the', 'dogs'], ['out', '?']],
+          ],
+          ragged_rank=1,
+          indices=[[1, -1, 2, 30], [1, 100]],
+          indices_ragged_rank=1,
+          expected=[[['takeover', 'offer'],
+                     ['$NONE^', '$NONE^'],
+                     ['from', 'Microsoft'],
+                     ['$NONE^', '$NONE^']],
+                    [['the', 'dogs'],
+                     ['$NONE^', '$NONE^']]],
+          expected_ragged_rank=1,
+          default_value=['$NONE^', '$NONE^'],
+      ),
+      # Dimensions:
+      # params: [2, (d1), (d2)]
+      # indices: [2, (d3)]
+      dict(
+          descr='params: [2, (d1), (d2)], indices: [2, (d3)]',
+          params=[
+              [['The', 'deal', 'came', 'about', '18', 'months', 'after',
+                'Yahoo', '!', 'rejected', 'a', '47.5', '-', 'billion', '-',
+                'dollar', 'takeover', 'offer', 'from', 'Microsoft', '.'],
+               ['Trumpty', 'Dumpty', 'sat', 'on', 'a', 'wall'],
+              ],
+              [['It\'s', 'always', 'darkest', 'before', 'the', 'dawn']]
+          ],
+          indices=[[1, 100], [0, -1]],
+          expected=[[['Trumpty', 'Dumpty', 'sat', 'on', 'a', 'wall'],
+                     ['$NONE^']],
+                    [["It's", 'always', 'darkest', 'before', 'the', 'dawn'],
+                     ['$NONE^']]]
+      ),
+      # Dimensions:
+      # params: [2, (d1), (d2)]
+      # indices: [2, (d1), (d3)]
+      dict(
+          descr='Test overbound indices of rank 3',
+          params=[
+              [['The', 'deal', 'came', 'about', '18', 'months', 'after',
+                'Yahoo', '!', 'rejected', 'a', '47.5', '-', 'billion', '-',
+                'dollar', 'takeover', 'offer', 'from', 'Microsoft', '.'],
+               ['Foo', 'bar', 'mar']],
+              [['He', 'left', 'us', '.', 'Little', 'boys', 'crowded',
+                'together', 'on', 'long', 'wooden', 'benches', ',', 'and', 'in',
+                'the', 'center', 'of', 'the', 'room', 'sat', 'the', 'teacher',
+                '.', 'His', 'black', 'beard', 'dripped', 'down', 'over', 'the',
+                'front', 'of', 'his', 'coat', '.', 'One', 'white', 'hand',
+                'poised', 'a', 'stick', 'above', 'his', 'desk', '.', 'He',
+                'turned', 'his', 'surly', ',', 'half', '-', 'closed', 'eyes',
+                'toward', 'us', ',', 'stared', 'for', 'a', 'second', ',',
+                'then', 'shouted', 'in', 'Yiddish', ',', '``', 'One', ',',
+                'two', ',',
+                'three', "''", '!', '!', 'Rapping', 'the', 'stick', 'against',
+                'the', 'desk', '.', 'The', 'little', 'boys', 'shrilled', 'out',
+                'a', 'Yiddish', 'translation', 'or', 'interpretation', 'of',
+                'the', 'Five', 'Books', 'of', 'Moses', ',', 'which', 'they',
+                'had', 'previously', 'chanted', 'in', 'Hebrew', '.'],
+               ['I', 'too', 'was', 'hustled', 'scammed', 'bamboozled', 'hood',
+                'winked', 'lead', 'astray']]
+          ],
+          indices=[[[8, 8823], [0, 100]], [[3, 23, 35, 45, 75, 83, 1234], [5]]],
+          expected=[[['!', '$NONE^'], ['Foo', '$NONE^']],
+                    [['.', '.', '.', '.', '!', '.', '$NONE^'],
+                     ['bamboozled']]],
+      ),
+      # params.shape = [2, (d1), 8]
+      # indices.shape = [2, (d1), 3]
+      dict(
+          descr='params = [2, (2, 1), 8], indices = [2, (2, 1), 3]',
+          params=[[['h'] * 8, ['w'] * 8], [['b'] * 8]],
+          ragged_rank=1,
+          indices=[[[0, 100, 1], [0, 1, 0]], [[1, 0, 0]]],
+          indices_ragged_rank=1,
+          expected=[[['h', '$NONE^', 'h'], ['w', 'w', 'w']], [['b', 'b', 'b']]],
+          expected_ragged_rank=1,
+      ),
+  ])
+  def testRaggedBatchGatherWithDefault(
+      self, descr, params, indices, expected, indices_ragged_rank=None,
+      expected_ragged_rank=None, ragged_rank=None, default_value='$NONE^'):
+    params = ragged_factory_ops.constant(params, ragged_rank=ragged_rank)
+    indices = ragged_factory_ops.constant(
+        indices, ragged_rank=indices_ragged_rank or ragged_rank)
+    expected = ragged_factory_ops.constant(
+        expected, ragged_rank=expected_ragged_rank or ragged_rank)
+    result = ragged_batch_gather_with_default_op.batch_gather_with_default(
+        params, indices, default_value)
+    self.assertAllEqual(result, expected)
+
+  @parameterized.parameters([
+      # Dimensions:
+      #  params: dims [2, 5], indices: [2, 2]
+      dict(
+          descr='params: dims [2, 5], indices: [2, 2]',
+          params=[
+              ['The', 'deal', 'came', 'about', '18'],
+              ['He', 'left', 'us', '.', 'Little']],
+          indices=[[0, -1], [3, 121]],
+          expected=[['The', '$NONE^'], ['.', '$NONE^']],
+          default_value='$NONE^',
+      ),
+      # Dimensions:
+      #  params: dims [2, 2, 5], indices: [2, 2]
+      dict(
+          descr='params: dims [2, 2, 5], indices: [2, 2]',
+          params=[
+              [['The', 'deal', 'came', 'about', '18'],
+               ['The', 'deal', 'came', 'about', '19'],
+              ],
+              [['He', 'left', 'us', '.', 'Little'],
+               ['The', 'deal', 'came', 'about', '20'],
+              ]
+          ],
+          indices=[[0, -1], [0, 121]],
+          expected=[[['The', 'deal', 'came', 'about', '18'],
+                     ['$NONE^', '$NONE^', '$NONE^', '$NONE^', '$NONE^']],
+                    [['He', 'left', 'us', '.', 'Little'],
+                     ['$NONE^', '$NONE^', '$NONE^', '$NONE^', '$NONE^']]],
+          default_value='$NONE^',
+      ),
+      # Test default_value with shape [5]
+      dict(
+          descr='params: dims [2, 2, 5], indices: [2, 2]',
+          params=[
+              [['The', 'deal', 'came', 'about', '18'],
+               ['The', 'deal', 'came', 'about', '19'],
+              ],
+              [['He', 'left', 'us', '.', 'Little'],
+               ['The', 'deal', 'came', 'about', '20'],
+              ]
+          ],
+          indices=[[0, -1], [0, 121]],
+          expected=[[['The', 'deal', 'came', 'about', '18'],
+                     [':FOO:', ':FOO:', ':FOO:', ':FOO:', ':FOO:']],
+                    [['He', 'left', 'us', '.', 'Little'],
+                     [':FOO:', ':FOO:', ':FOO:', ':FOO:', ':FOO:']]],
+          default_value=[':FOO:', ':FOO:', ':FOO:', ':FOO:', ':FOO:'],
+      ),
+  ])
+  def testRaggedBatchGatherWithDefaultOnTensors(
+      self, descr, params, indices, expected, default_value):
+    params = constant_op.constant(params)
+    indices = constant_op.constant(indices)
+    expected = constant_op.constant(expected)
+    result = ragged_batch_gather_with_default_op.batch_gather_with_default(
+        params, indices, default_value)
+    self.assertAllEqual(expected, result)
+
+  @parameterized.parameters([
+      dict(
+          params=[['The', 'deal', 'came', 'about', '18', 'months', 'after',
+                   'Yahoo', '!', 'rejected', 'a', '47.5', '-', 'billion', '-',
+                   'dollar', 'takeover', 'offer', 'from', 'Microsoft', '.']],
+          indices=[[[8, -1]]],
+          # Exception here because different errors are thrown in eager vs
+          # graph mode.
+          error=Exception,
+          default_value='$NONE^',
+      ),
+  ])
+  def testRankMismatch(
+      self, params, indices, default_value, error):
+    params = ragged_factory_ops.constant(params)
+    indices = ragged_factory_ops.constant(indices)
+    with self.assertRaises(error):
+      _ = ragged_batch_gather_with_default_op.batch_gather_with_default(
+          params, indices, default_value)
+
+  @parameterized.parameters([
+      # Dimensions:
+      # params: [2, (d1), 2]
+      # indices: [2, (d2)]
+      # default_value: []
+      dict(
+          descr='params: rank 3, indices: rank 2, default: rank = [], but'
+          ' should be [2]',
+          params=[
+              [['The', 'deal'], ['takeover', 'offer'], ['from', 'Microsoft']],
+              [['Who', 'let'], ['the', 'dogs'], ['out', '?']],
+          ],
+          ragged_rank=1,
+          indices=[[1, -1, 2, 30], [1, 100]],
+          indices_ragged_rank=1,
+          default_value='$NONE^',
+          error=Exception,
+      )
+  ])
+  def testInvalidDefaultValueRank(
+      self, descr, params, indices, default_value, error, ragged_rank=None,
+      indices_ragged_rank=None):
+    params = ragged_factory_ops.constant(params, ragged_rank=ragged_rank)
+    indices = ragged_factory_ops.constant(
+        indices, ragged_rank=indices_ragged_rank)
+    with self.assertRaises(error):
+      _ = ragged_batch_gather_with_default_op.batch_gather_with_default(
+          params, indices, default_value)
+
+  def testRaggedBatchGatherUnknownRankError(self):
+    if context.executing_eagerly():
+      return
+    params = [['a', 'b'], ['c', 'd']]
+    indices = array_ops.placeholder(dtypes.int32, shape=None)
+    ragged_indices = ragged_tensor.RaggedTensor.from_row_splits(
+        indices, [0, 2, 4])
+
+    with self.assertRaisesRegex(
+        ValueError, r'batch_dims=-1 may only be negative '
+        r'if rank\(indices\) is statically known.'):
+      ragged_batch_gather_ops.batch_gather(params, indices)
+
+    with self.assertRaisesRegex(
+        ValueError, r'batch_dims=-1 may only be negative '
+        r'if rank\(indices\) is statically known.'):
+      ragged_batch_gather_ops.batch_gather(params, ragged_indices)
+
+  @parameterized.parameters(
+      [
+          dict(
+              params=ragged_factory_ops.constant_value([['a'], ['b'], ['c']]),
+              indices=ragged_factory_ops.constant_value([[0], [0]]),
+              message=(r'batch shape from indices .* does not match params')),
+          dict(
+              params=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+              indices=ragged_factory_ops.constant_value([[[0, 0], [0, 0, 0]],
+                                                         [[0]]]),
+              message='batch shape from indices does not match params shape'),
+          dict(  # rank mismatch
+              params=ragged_factory_ops.constant_value([[[0, 0], [0, 0, 0]],
+                                                        [[0]]]),
+              indices=ragged_factory_ops.constant_value([[[0, 0]], [[0, 0, 0]],
+                                                         [[0]]]),
+              error=(ValueError, errors.InvalidArgumentError)),
+          dict(
+              params=ragged_factory_ops.constant_value([[[0, 0], [0, 0, 0]],
+                                                        [[0]], [[0]]]),
+              indices=ragged_factory_ops.constant_value([[[0, 0]], [[0, 0, 0]],
+                                                         [[0]]]),
+              error=(ValueError, errors.InvalidArgumentError),
+              message=(r'batch shape from indices .* does not match '
+                       r'params shape|dimension size mismatch')),
+          dict(
+              params=ragged_factory_ops.constant_value(['a', 'b', 'c']),
+              indices=ragged_factory_ops.constant_value([[0], [0]]),
+              message=r'batch_dims must be less than rank\(params\)'),
+          dict(
+              params=ragged_factory_ops.constant_value([['a']]),
+              indices=0,
+              message='batch_dims=-1 out of bounds: expected 0<=batch_dims<0'),
+          dict(
+              params=ragged_factory_ops.constant_value([['a']]),
+              indices=[[[0]]],
+              message=r'batch_dims must be less than rank\(params\)'),
+      ])
+  def testRaggedBatchGatherStaticError(self,
+                                       params,
+                                       indices,
+                                       message=None,
+                                       error=ValueError):
+    with self.assertRaisesRegex(error, message):
+      ragged_batch_gather_ops.batch_gather(params, indices)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_batch_gather_ops.py

@@ -0,0 +1,60 @@
+# 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.
+# ==============================================================================
+"""Batch gather operations for RaggedTensors."""
+
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_gather_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import dispatch
+
+
+#===============================================================================
+# ragged.batch_gather
+#===============================================================================
+@dispatch.dispatch_for_api(array_ops.batch_gather)
+def batch_gather(params: ragged_tensor.RaggedOrDense,
+                 indices: ragged_tensor.RaggedOrDense,
+                 name=None):
+  """Gathers slices from `params` according to `indices` with batch dims.
+
+  This operation is similar to `gather`, but it assumes that the leading `N`
+  dimensions of `indices` and `params` are batch dimensions, and performs a
+  gather within each batch.  In particular, when using this operation with `N`
+  batch dimensions `B1...BN`:
+
+  * `indices` has shape `[B1...BN, I]`
+  * `params` has shape `[B1...BN, P1...PM]`.
+  * `result` has shape `[B1...BN, I, P2...PM]`.
+  * `result[b1...bN, i, p2...pM] =
+    params[b1...bN, indices[b1...bN, i], p2...pM]`
+
+  Args:
+    params: A potentially ragged tensor with shape `[B1...BN, P1...PM]` (`N>=0`,
+      `M>0`).
+    indices: A potentially ragged tensor with shape `[B1...BN, I]` (`N>=0`).
+    name: A name for the operation (optional).
+
+  Returns:
+    A potentially ragged tensor with shape `[B1...BN, I, P2...PM]`.
+    `result.ragged_rank = max(indices.ragged_rank, params.ragged_rank)`.
+
+  #### Example:
+
+  >>> params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
+  >>> indices = tf.ragged.constant([[1, 2, 0], [], [], [0, 0]])
+  >>> tf.compat.v1.batch_gather(params, indices)
+  <tf.RaggedTensor [[b'b', b'c', b'a'], [], [], [b'e', b'e']]>
+  """
+  return ragged_gather_ops.gather(params, indices, batch_dims=-1, name=name)

tensorflow/python/ops/ragged/ragged_batch_gather_with_default_op.py

@@ -0,0 +1,179 @@
+# 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.
+# ==============================================================================
+"""Array operations for RaggedTensors."""
+
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_dispatch  # pylint: disable=unused-import
+from tensorflow.python.ops.ragged import ragged_operators  # pylint: disable=unused-import
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_shape
+from tensorflow.python.ops.ragged import ragged_where_op
+
+
+# ===============================================================================
+# ragged.batch_gather_with_default
+# ===============================================================================
+def batch_gather_with_default(params,
+                              indices,
+                              default_value='',
+                              name=None):
+  """Same as `batch_gather` but inserts `default_value` for invalid indices.
+
+  This operation is similar to `batch_gather` except that it will substitute
+  the value for invalid indices with `default_value` as the contents.
+  See `batch_gather` for more details.
+
+
+  Args:
+    params: A potentially ragged tensor with shape `[B1...BN, P1...PM]` (`N>=0`,
+      `M>0`).
+    indices: A potentially ragged tensor with shape `[B1...BN, I]` (`N>=0`).
+    default_value: A value to be inserted in places where `indices` are out of
+      bounds. Must be the same dtype as params and either a scalar or rank 1.
+    name: A name for the operation (optional).
+
+  Returns:
+    A potentially ragged tensor with shape `[B1...BN, I, P2...PM]`.
+    `result.ragged_rank = max(indices.ragged_rank, params.ragged_rank)`.
+
+  #### Example:
+
+  >>> params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
+  >>> indices = tf.ragged.constant([[1, 2, -1], [], [], [0, 10]])
+  >>> batch_gather_with_default(params, indices, 'FOO')
+  <tf.RaggedTensor [[b'b', b'c', b'FOO'], [], [], [b'e', b'FOO']]>
+
+  """
+  with ops.name_scope(name, 'RaggedBatchGatherWithDefault'):
+    params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        params, name='params',
+    )
+    indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        indices, name='indices',
+    )
+    default_value = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        default_value, name='default_value',
+    )
+    row_splits_dtype, (params, indices, default_value) = (
+        ragged_tensor.match_row_splits_dtypes(params, indices, default_value,
+                                              return_dtype=True))
+    # TODO(hterry): lift this restriction and support default_values of
+    #               of rank > 1
+    if default_value.shape.ndims not in (0, 1):
+      raise ValueError('"default_value" must be a scalar or vector')
+    upper_bounds = None
+    if indices.shape.ndims is None:
+      raise ValueError('Indices must have a known rank.')
+    if params.shape.ndims is None:
+      raise ValueError('Params must have a known rank.')
+
+    num_batch_dimensions = indices.shape.ndims - 1
+    pad = None
+    # The logic for this works as follows:
+    # - create a padded params, where:
+    #    padded_params[b1...bn, 0] = default_value
+    #    padded_params[b1...bn, i] = params[b1...bn, i-1] (i>0)
+    # - create an `upper_bounds` Tensor that contains the number of elements
+    #   in each innermost rank. Broadcast `upper_bounds` to be the same shape
+    #   as `indices`.
+    # - check to see which index in `indices` are out of bounds and substitute
+    #   it with the index containing `default_value` (the first).
+    # - call batch_gather with the indices adjusted.
+    with ops.control_dependencies([
+        check_ops.assert_greater_equal(array_ops.rank(params),
+                                       array_ops.rank(indices))]):
+      if ragged_tensor.is_ragged(params):
+        row_lengths = ragged_array_ops.expand_dims(
+            params.row_lengths(axis=num_batch_dimensions),
+            axis=-1)
+        upper_bounds = math_ops.cast(row_lengths, indices.dtype)
+
+        pad_shape = _get_pad_shape(params, indices, row_splits_dtype)
+
+        pad = ragged_tensor_shape.broadcast_to(
+            default_value, pad_shape)
+      else:
+        params_shape = array_ops.shape(params)
+        pad_shape = array_ops.concat([
+            params_shape[:num_batch_dimensions],
+            [1],
+            params_shape[num_batch_dimensions + 1:params.shape.ndims]
+        ], 0)
+        upper_bounds = params_shape[num_batch_dimensions]
+        pad = array_ops.broadcast_to(default_value, pad_shape)
+
+      # Add `default_value` as the first value in the innermost (ragged) rank.
+      pad = math_ops.cast(pad, params.dtype)
+      padded_params = array_ops.concat(
+          [pad, params], axis=num_batch_dimensions)
+
+      # Adjust the indices by substituting out-of-bound indices to the
+      # default-value index (which is the first element)
+      shifted_indices = indices + 1
+      is_out_of_bounds = (indices < 0) | (indices >= upper_bounds)
+      adjusted_indices = ragged_where_op.where(
+          is_out_of_bounds,
+          x=array_ops.zeros_like(indices), y=shifted_indices,
+      )
+      return array_ops.batch_gather(
+          params=padded_params, indices=adjusted_indices, name=name)
+
+
+def _get_pad_shape(params, indices, row_splits_dtype):
+  """Gets the RaggedTensorDynamicShape for the pad tensor."""
+  num_batch_dimensions = indices.shape.ndims - 1
+  params_shape = ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(
+      params, dim_size_dtype=row_splits_dtype)
+
+  # We want to create a pad tensor that can be concatenated with the params.
+  if params.shape.ndims == indices.shape.ndims:
+    # When params and indices are the same rank, the shape of the pad tensor is
+    # almost identical to params, except the last dimension which has size = 1.
+    if params_shape.num_inner_dimensions == 0:
+      pad_dims = params_shape.partitioned_dim_sizes[:-1] + (
+          array_ops.ones_like(params_shape.partitioned_dim_sizes[-1]),)
+      return ragged_tensor_shape.RaggedTensorDynamicShape(
+          pad_dims, [])
+    else:
+      return ragged_tensor_shape.RaggedTensorDynamicShape(
+          params_shape.partitioned_dim_sizes,
+          array_ops.concat([params_shape.inner_dim_sizes[:-1], [1]], axis=0))
+  else:
+    # When the rank of indices < params, the pad has the same dimension as
+    # params up to the 'num_batch_dimensions' rank. Every dimension after that
+    # has size 1.
+    pad_dims = None
+    if num_batch_dimensions == 0:
+      pad_dims = (constant_op.constant(1, dtype=row_splits_dtype),) + (
+          constant_op.constant([1], dtype=row_splits_dtype),) * (
+              params_shape.num_partitioned_dimensions -
+              num_batch_dimensions - 1)
+    else:
+      batch_dimensions = params_shape.partitioned_dim_sizes[
+          :num_batch_dimensions]
+      gather_dimension = params_shape.partitioned_dim_sizes[
+          num_batch_dimensions]
+      pad_dims = batch_dimensions + (
+          array_ops.ones_like(gather_dimension),) * (
+              params_shape.num_partitioned_dimensions - num_batch_dimensions)
+
+    return ragged_tensor_shape.RaggedTensorDynamicShape(
+        pad_dims, params_shape.inner_dim_sizes)

tensorflow/python/ops/ragged/ragged_bincount_ops.py

@@ -0,0 +1,404 @@
+# Copyright 2023 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
+# maxlengthations under the License.
+# ==============================================================================
+"""bincount ops for RaggedTensors."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import bincount_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import gen_count_ops
+from tensorflow.python.ops import gen_math_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import dispatch
+
+
+@dispatch.dispatch_for_api(bincount_ops.bincount)
+def bincount(arr: ragged_tensor.RaggedTensor,
+             weights=None,
+             minlength=None,
+             maxlength=None,
+             dtype=dtypes.int32,
+             name=None,
+             axis=None,
+             binary_output=False):
+  """Counts the number of occurrences of each value in an integer array.
+
+  If `minlength` and `maxlength` are not given, returns a vector with length
+  `tf.reduce_max(arr) + 1` if `arr` is non-empty, and length 0 otherwise.
+
+  >>> data = tf.ragged.constant([[1, 1], [2, 3, 2, 4, 4, 5]])
+  >>> tf.math.bincount(data)
+  <tf.Tensor: ... numpy=array([0, 2, 2, 1, 2, 1], dtype=int32)>
+
+  Vector length = Maximum element in vector `values` is 5. Adding 1, which is 6
+                  will be the vector length.
+
+  Each bin value in the output indicates number of occurrences of the particular
+  index. Here, index 1 in output has a value 2. This indicates value 1 occurs
+  two times in `values`.
+
+  **Bin-counting with weights**
+
+  >>> data = tf.ragged.constant([[1, 1], [2, 3, 2, 4, 4, 5]])
+  >>> weights = tf.ragged.constant([[1, 5], [0, 1, 0, 5, 4, 5]])
+  >>> tf.math.bincount(data, weights=weights)
+  <tf.Tensor: ... numpy=array([0, 6, 0, 1, 9, 5], dtype=int32)>
+
+  When `weights` is specified, bins will be incremented by the corresponding
+  weight instead of 1. Here, index 1 in output has a value 6. This is the
+  summation of `weights` corresponding to the value in `arr` (i.e. for index
+  1, the first two values `arr` are 1 so the first two weights, 1 and 5, are
+  summed).
+
+  There is an equivilance between bin-counting with weights and
+  `unsorted_segment_sum` where `data` is the weights and `segment_ids` are the
+  values.
+
+  >>> data = tf.ragged.constant([[1, 1], [2, 3, 2, 4, 4, 5]])
+  >>> weights = tf.ragged.constant([[1, 5], [0, 1, 0, 5, 4, 5]])
+  >>> tf.math.unsorted_segment_sum(weights, data, num_segments=6).numpy()
+  array([0, 6, 0, 1, 9, 5], dtype=int32)
+
+  On GPU, `bincount` with weights is only supported when XLA is enabled
+  (typically when a function decorated with `@tf.function(jit_compile=True)`).
+  `unsorted_segment_sum` can be used as a workaround for the non-XLA case on
+  GPU.
+
+  **Bin-counting matrix rows independently**
+
+  This example uses `axis=-1` with a 2 dimensional input and returns a
+  `Tensor` with bincounting where axis 0 is **not** flattened, i.e. an
+  independent bincount for each matrix row.
+
+  >>> data = tf.ragged.constant([[1, 2], [3, 0, 0, 0, 1, 2]], dtype=np.int32)
+  >>> tf.math.bincount(data, axis=-1)
+  <tf.Tensor: shape=(2, 4), dtype=int32, numpy=
+    array([[0, 1, 1, 0],
+           [3, 1, 1, 1]], dtype=int32)>
+
+  **Bin-counting with binary_output**
+
+  This example gives binary output instead of counting the occurrence.
+
+  >>> data = tf.ragged.constant([[1, 2], [3, 0, 0, 0, 1, 2]], dtype=np.int32)
+  >>> tf.math.bincount(data, axis=-1, binary_output=True)
+  <tf.Tensor: shape=(2, 4), dtype=int32, numpy=
+    array([[0, 1, 1, 0],
+           [1, 1, 1, 1]], dtype=int32)>
+
+  Args:
+    arr: A RaggedTensor whose values should be counted. These tensors must have
+      a rank of 2 if `axis=-1`.
+    weights: If non-None, must be a RaggedTensor with the same row splits as
+      `arr`. For each value in `arr`, the bin will be incremented by the
+      corresponding weight instead of 1. If non-None, `binary_output` must be
+      False.
+    minlength: If given, ensures the output has length at least `minlength`,
+      padding with zeros at the end if necessary.
+    maxlength: If given, skips values in `arr` that are equal or greater than
+      `maxlength`, ensuring that the output has length at most `maxlength`.
+    dtype: If `weights` is None, determines the type of the output bins.
+    name: A name scope for the associated operations (optional).
+    axis: The axis to slice over. Axes at and below `axis` will be flattened
+      before bin counting. Currently, only `0`, and `-1` are supported. If None,
+      all axes will be flattened (identical to passing `0`).
+    binary_output: If True, this op will output 1 instead of the number of times
+      a token appears (equivalent to one_hot + reduce_any instead of one_hot +
+      reduce_add). Defaults to False.
+
+  Returns:
+    A vector with the same dtype as `weights` or the given `dtype` containing
+    the bincount values.
+
+  Raises:
+    `InvalidArgumentError` if negative values are provided as an input.
+  """
+  name = "bincount" if name is None else name
+  with ops.name_scope(name):
+    arr = ragged_tensor.convert_to_tensor_or_ragged_tensor(arr, name="arr")
+    if weights is not None:
+      if not isinstance(weights, sparse_tensor.SparseTensor):
+        weights = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+            weights, name="weights")
+
+    if weights is not None and binary_output:
+      raise ValueError("Arguments `binary_output` and `weights` are mutually "
+                       "exclusive. Please specify only one.")
+
+    if not arr.dtype.is_integer:
+      arr = math_ops.cast(arr, dtypes.int32)
+    if axis is None:
+      axis = 0
+
+    if axis not in [0, -1]:
+      raise ValueError(f"Unsupported value for argument axis={axis}. Only 0 and"
+                       " -1 are currently supported.")
+
+    array_is_nonempty = array_ops.size(arr) > 0
+    output_size = math_ops.cast(array_is_nonempty, arr.dtype) * (
+        math_ops.reduce_max(arr) + 1)
+    if minlength is not None:
+      minlength = ops.convert_to_tensor(
+          minlength, name="minlength", dtype=arr.dtype)
+      output_size = gen_math_ops.maximum(minlength, output_size)
+    if maxlength is not None:
+      maxlength = ops.convert_to_tensor(
+          maxlength, name="maxlength", dtype=arr.dtype)
+      output_size = gen_math_ops.minimum(maxlength, output_size)
+
+    if axis == 0:
+      # Flatten RaggedTensors with multiple ragged dimensions which use a
+      # nested RaggedTensor for the values tensor.
+      while isinstance(arr, ragged_tensor.RaggedTensor):
+        if weights is not None:
+          weights = validate_ragged_weights(arr, weights, dtype)
+        arr = arr.values
+
+    if isinstance(arr, ragged_tensor.RaggedTensor):
+      weights = validate_ragged_weights(arr, weights, dtype)
+      return gen_math_ops.ragged_bincount(
+          splits=arr.row_splits,
+          values=arr.values,
+          size=output_size,
+          weights=weights,
+          binary_output=binary_output)
+    else:
+      weights = bincount_ops.validate_dense_weights(arr, weights, dtype)
+      return gen_math_ops.dense_bincount(
+          input=arr,
+          size=output_size,
+          weights=weights,
+          binary_output=binary_output)
+
+
+@dispatch.dispatch_for_api(sparse_ops.sparse_bincount)
+def sparse_bincount(values: ragged_tensor.RaggedTensor,
+                    weights=None,
+                    axis=0,
+                    minlength=None,
+                    maxlength=None,
+                    binary_output=False,
+                    name=None):
+  """Count the number of times an integer value appears in a tensor.
+
+  This op takes an N-dimensional `Tensor`, `RaggedTensor`, or `SparseTensor`,
+  and returns an N-dimensional int64 SparseTensor where element
+  `[i0...i[axis], j]` contains the number of times the value `j` appears in
+  slice `[i0...i[axis], :]` of the input tensor.  Currently, only N=0 and
+  N=-1 are supported.
+
+  Args:
+    values: A RaggedTensor whose values should be
+      counted. These tensors must have a rank of 2 if `axis=-1`.
+    weights: If non-None, must be a RaggedTensor with the same row splits as
+      `values`. For each value in `value`, the bin will be incremented by the
+      corresponding weight instead of 1.
+    axis: The axis to slice over. Axes at and below `axis` will be flattened
+      before bin counting. Currently, only `0`, and `-1` are supported. If None,
+      all axes will be flattened (identical to passing `0`).
+    minlength: If given, ensures the output has length at least `minlength`,
+      padding with zeros at the end if necessary.
+    maxlength: If given, skips values in `values` that are equal or greater than
+      `maxlength`, ensuring that the output has length at most `maxlength`.
+    binary_output: If True, this op will output 1 instead of the number of times
+      a token appears (equivalent to one_hot + reduce_any instead of one_hot +
+      reduce_add). Defaults to False.
+    name: A name for this op.
+
+  Returns:
+    A SparseTensor with `output.shape = values.shape[:axis] + [N]`, where `N` is
+      * `maxlength` (if set);
+      * `minlength` (if set, and `minlength > reduce_max(values)`);
+      * `0` (if `values` is empty);
+      * `reduce_max(values) + 1` otherwise.
+
+  Raises:
+    `InvalidArgumentError` if negative values are provided as an input.
+
+  Examples:
+
+  **Bin-counting every item in individual batches**
+
+  This example takes an input (which could be a Tensor, RaggedTensor, or
+  SparseTensor) and returns a SparseTensor where the value of (i,j) is the
+  number of times value j appears in batch i.
+
+  >>> data = tf.ragged.constant(
+  ...     [[10, 20], [30, 20, 11, 101, 11, 10001]], dtype=np.int64)
+  >>> tf.sparse.bincount(data, axis=-1)
+  SparseTensor(indices=tf.Tensor(
+    [[    0    10]
+     [    0    20]
+     [    1    11]
+     [    1    20]
+     [    1    30]
+     [    1   101]
+     [    1 10001]], shape=(7, 2), dtype=int64),
+     values=tf.Tensor([1 1 2 1 1 1 1], shape=(7,), dtype=int64),
+     dense_shape=tf.Tensor([    2 10002], shape=(2,), dtype=int64))
+
+  **Bin-counting with defined output shape**
+
+  This example takes an input (which could be a Tensor, RaggedTensor, or
+  SparseTensor) and returns a SparseTensor where the value of (i,j) is the
+  number of times value j appears in batch i. However, all values of j
+  above 'maxlength' are ignored. The dense_shape of the output sparse tensor
+  is set to 'minlength'. Note that, while the input is identical to the
+  example above, the value '10001' in batch item 2 is dropped, and the
+  dense shape is [2, 500] instead of [2,10002] or [2, 102].
+
+  >>> minlength = maxlength = 500
+  >>> data = tf.ragged.constant(
+  ...     [[10, 20], [30, 20, 11, 101, 11, 10001]], dtype=np.int64)
+  >>> tf.sparse.bincount(
+  ...    data, axis=-1, minlength=minlength, maxlength=maxlength)
+  SparseTensor(indices=tf.Tensor(
+  [[  0  10]
+   [  0  20]
+   [  1  11]
+   [  1  20]
+   [  1  30]
+   [  1 101]], shape=(6, 2), dtype=int64),
+   values=tf.Tensor([1 1 2 1 1 1], shape=(6,), dtype=int64),
+   dense_shape=tf.Tensor([  2 500], shape=(2,), dtype=int64))
+
+  **Binary bin-counting**
+
+  This example takes an input (which could be a Tensor, RaggedTensor, or
+  SparseTensor) and returns a SparseTensor where (i,j) is 1 if the value j
+  appears in batch i at least once and is 0 otherwise. Note that, even though
+  some values (like 20 in batch 1 and 11 in batch 2) appear more than once,
+  the 'values' tensor is all 1s.
+
+  >>> data = tf.ragged.constant(
+  ...     [[10, 20], [30, 20, 11, 101, 11, 10001]], dtype=np.int64)
+  >>> tf.sparse.bincount(data, binary_output=True, axis=-1)
+  SparseTensor(indices=tf.Tensor(
+  [[  0  10]
+   [  0  20]
+   [  1  11]
+   [  1  20]
+   [  1  30]
+   [  1 101]
+   [  1 10001]], shape=(7, 2), dtype=int64),
+   values=tf.Tensor([1 1 1 1 1 1 1], shape=(7,), dtype=int64),
+   dense_shape=tf.Tensor([  2 10002], shape=(2,), dtype=int64))
+
+  **Weighted bin-counting**
+
+  This example takes two inputs - a values tensor and a weights tensor. These
+  tensors must be identically shaped, and have the same row splits or indices
+  in the case of RaggedTensors or SparseTensors. When performing a weighted
+  count, the op will output a SparseTensor where the value of (i, j) is the
+  sum of the values in the weight tensor's batch i in the locations where
+  the values tensor has the value j. In this case, the output dtype is the
+  same as the dtype of the weights tensor.
+
+  >>> data = tf.ragged.constant(
+  ...    [[10, 20], [30, 20, 11, 101, 11, 10001]], dtype=np.int64)
+  >>> weights = tf.ragged.constant(
+  ...    [[2, 0.25], [15, 0.5, 2, 17, 3, 0.9]])
+  >>> tf.sparse.bincount(data, weights=weights, axis=-1)
+  SparseTensor(indices=tf.Tensor(
+  [[    0    10]
+   [    0    20]
+   [    1    11]
+   [    1    20]
+   [    1    30]
+   [    1   101]
+   [    1 10001]], shape=(7, 2), dtype=int64),
+   values=tf.Tensor([ 2. 0.25  5. 0.5 15. 17. 0.9 ], shape=(7,), dtype=float32),
+   dense_shape=tf.Tensor([    2 10002], shape=(2,), dtype=int64))
+
+  """
+  with ops.name_scope(name, "count", [values, weights]):
+    values = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        values, name="values")
+    if weights is not None:
+      if not isinstance(weights, sparse_tensor.SparseTensor):
+        weights = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+            weights, name="weights")
+
+    if weights is not None and binary_output:
+      raise ValueError("Arguments `binary_output` and `weights` are mutually "
+                       "exclusive. Please specify only one.")
+
+    if axis is None:
+      axis = 0
+
+    if axis not in [0, -1]:
+      raise ValueError(f"Unsupported value for argument axis={axis}. Only 0 and"
+                       " -1 are currently supported.")
+
+    minlength_value = minlength if minlength is not None else -1
+    maxlength_value = maxlength if maxlength is not None else -1
+
+    if axis == 0:
+      if weights is not None:
+        weights = validate_ragged_weights(values, weights)
+      values = values.values
+
+    if isinstance(values, ragged_tensor.RaggedTensor):
+      weights = validate_ragged_weights(values, weights)
+      c_ind, c_val, c_shape = gen_count_ops.ragged_count_sparse_output(
+          values.row_splits,
+          values.values,
+          weights,
+          minlength=minlength_value,
+          maxlength=maxlength_value,
+          binary_output=binary_output)
+    else:
+      weights = bincount_ops.validate_dense_weights(values, weights)
+      c_ind, c_val, c_shape = gen_count_ops.dense_count_sparse_output(
+          values,
+          weights=weights,
+          minlength=minlength_value,
+          maxlength=maxlength_value,
+          binary_output=binary_output)
+
+    return sparse_tensor.SparseTensor(c_ind, c_val, c_shape)
+
+
+def validate_ragged_weights(values, weights, dtype=None):
+  """Validates the passed weight tensor or creates an empty one."""
+  if weights is None:
+    if dtype:
+      return array_ops.constant([], dtype=dtype)
+    return array_ops.constant([], dtype=values.values.dtype)
+
+  if not isinstance(weights, ragged_tensor.RaggedTensor):
+    raise ValueError(
+        "`weights` must be a RaggedTensor if `values` is a RaggedTensor. "
+        f"Received argument weights={weights} of type: "
+        f"{type(weights).__name__}.")
+
+  checks = []
+  if weights.row_splits is not values.row_splits:
+    checks.append(
+        check_ops.assert_equal(
+            weights.row_splits,
+            values.row_splits,
+            message="'weights' and 'values' must have the same row splits."))
+  if checks:
+    with ops.control_dependencies(checks):
+      weights = array_ops.identity(weights.values)
+  else:
+    weights = weights.values
+
+  return weights

tensorflow/python/ops/ragged/ragged_bincount_ops_test.py

@@ -0,0 +1,558 @@
+# Copyright 2023 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
+# maxlengthations under the License.
+# ==============================================================================
+"""Tests for bincount ops with RaggedTensor inputs."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import config as tf_config
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import bincount_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import test
+
+
+def _ragged_factory(x):
+  return lambda: ragged_factory_ops.constant(x)
+
+
+def _adjust_expected_rank1(x, minlength, maxlength):
+  """Trim or pad an expected result based on minlength and maxlength."""
+  n = len(x)
+  if (minlength is not None) and (n < minlength):
+    x = x + [0] * (minlength - n)
+  if (maxlength is not None) and (n > maxlength):
+    x = x[:maxlength]
+  return x
+
+
+def _adjust_expected_rank2(x, minlength, maxlength):
+  return [_adjust_expected_rank1(i, minlength, maxlength) for i in x]
+
+
+class TestDenseBincount(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_ragged_input_count(self, dtype):
+    x = ragged_factory_ops.constant([[], [], [3, 0, 1], [], [5, 0, 4, 4]],
+                                    dtype)
+    # pyformat: disable
+    expected_output = [
+        [0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0],
+        [1, 1, 0, 1, 0, 0],
+        [0, 0, 0, 0, 0, 0],
+        [1, 0, 0, 0, 2, 1]]
+    # pyformat: enable
+    self.assertAllEqual(expected_output,
+                        self.evaluate(bincount_ops.bincount(arr=x, axis=-1)))
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_ragged_input_binary(self, dtype):
+    x = ragged_factory_ops.constant([[], [], [3, 0, 1], [], [5, 0, 4, 4]])
+    # pyformat: disable
+    expected_output = [
+        [0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0],
+        [1, 1, 0, 1, 0, 0],
+        [0, 0, 0, 0, 0, 0],
+        [1, 0, 0, 0, 1, 1]]
+    # pyformat: enable
+    self.assertAllEqual(
+        expected_output,
+        self.evaluate(
+            bincount_ops.bincount(arr=x, axis=-1, binary_output=True)))
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_ragged_input_count_with_weights(self, dtype):
+    x = ragged_factory_ops.constant([[], [], [3, 0, 1], [], [5, 0, 4, 4]])
+    weights = ragged_factory_ops.constant([[], [], [.1, .2, .3], [],
+                                           [.2, .5, .6, .3]])
+    # pyformat: disable
+    expected_output = [
+        [0, 0, 0, 0, 0, 0],
+        [0, 0, 0, 0, 0, 0],
+        [.2, .3, 0, .1, 0, 0],
+        [0, 0, 0, 0, 0, 0],
+        [.5, 0, 0, 0, .9, .2]]
+    # pyformat: enable
+    self.assertAllClose(
+        expected_output,
+        self.evaluate(bincount_ops.bincount(arr=x, weights=weights, axis=-1)))
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_ragged_input_count_np(self, dtype):
+    np.random.seed(42)
+    num_rows = 128
+    num_cols = 27
+    size = 1000
+    inp = np.random.randint(0, size, (num_rows, num_cols), dtype=dtype)
+    np_out = np.reshape(
+        np.concatenate(
+            [np.bincount(inp[j, :], minlength=size) for j in range(num_rows)],
+            axis=0), (num_rows, size))
+    x = ragged_tensor.RaggedTensor.from_tensor(inp)
+    self.assertAllEqual(
+        np_out,
+        self.evaluate(bincount_ops.bincount(arr=x, minlength=size, axis=-1)))
+
+  @parameterized.parameters([{
+      "dtype": np.int32,
+  }, {
+      "dtype": np.int64,
+  }])
+  def test_ragged_input_count_np_with_weights(self, dtype):
+    np.random.seed(42)
+    num_rows = 128
+    num_cols = 27
+    size = 1000
+    inp = np.random.randint(0, size, (num_rows, num_cols), dtype=dtype)
+    np_weight = np.random.random((num_rows, num_cols))
+    np_out = np.reshape(
+        np.concatenate([
+            np.bincount(inp[j, :], weights=np_weight[j, :], minlength=size)
+            for j in range(num_rows)
+        ],
+                       axis=0), (num_rows, size))
+    x = ragged_tensor.RaggedTensor.from_tensor(inp)
+    weights = ragged_tensor.RaggedTensor.from_tensor(np_weight)
+    self.assertAllEqual(
+        np_out,
+        self.evaluate(
+            bincount_ops.bincount(
+                arr=x, weights=weights, minlength=size, axis=-1)))
+
+  @parameterized.product(
+      (
+          dict(
+              tid="_r2",
+              x_factory=_ragged_factory([[], [1], [2, 2], [3, 3, 3]]),
+              expected=[0, 1, 2, 3],  # no implied zeros
+          ),
+          dict(
+              tid="_r3",
+              x_factory=_ragged_factory([[[], [1]], [[2, 2], [3, 3, 3]]]),
+              expected=[0, 1, 2, 3],  # no implied zeros
+          ),
+      ),
+      (
+          dict(minlength=None, maxlength=None),
+          dict(minlength=3, maxlength=None),
+          dict(minlength=5, maxlength=None),
+          dict(minlength=None, maxlength=3),
+          dict(minlength=None, maxlength=5),
+          dict(minlength=2, maxlength=3),
+          dict(minlength=3, maxlength=5),
+          dict(minlength=5, maxlength=10),
+          dict(minlength=None, maxlength=0),
+      ),
+  )
+  def test_default(self, x_factory, minlength, maxlength, expected, tid=None):
+    x = x_factory()
+    expected = _adjust_expected_rank1(expected, minlength, maxlength)
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(x, minlength=minlength, maxlength=maxlength)
+        ),
+    )
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(
+                x, minlength=minlength, maxlength=maxlength, axis=0
+            )
+        ),
+    )
+
+  @parameterized.product(
+      (
+          dict(
+              tid="_r2",
+              x_factory=_ragged_factory([[], [1], [2, 2], [3, 3, 3]]),
+              # no implied zeros
+              expected=[[0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 2, 0], [0, 0, 0, 3]],
+          ),
+      ),
+      (
+          dict(minlength=None, maxlength=None),
+          dict(minlength=3, maxlength=None),
+          dict(minlength=5, maxlength=None),
+          dict(minlength=None, maxlength=3),
+          dict(minlength=None, maxlength=5),
+          dict(minlength=2, maxlength=3),
+          dict(minlength=3, maxlength=5),
+          dict(minlength=5, maxlength=10),
+          dict(minlength=None, maxlength=0),
+      ),
+  )
+  def test_axis_neg_1(self, tid, x_factory, minlength, maxlength, expected):
+    x = x_factory()
+    expected = _adjust_expected_rank2(expected, minlength, maxlength)
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(
+                x, minlength=minlength, maxlength=maxlength, axis=-1
+            )
+        ),
+    )
+
+  @parameterized.product(
+      (
+          dict(
+              tid="_r2",
+              x_factory=_ragged_factory([[], [1], [2, 2], [3, 3, 3]]),
+              weights_factory=_ragged_factory([[], [1], [2, 3], [4, 5, 6]]),
+              axis=None,
+              expected=[0, 1, 5, 15],  # no implied zeros
+          ),
+          dict(
+              tid="_r3",
+              x_factory=_ragged_factory([[[], [1]], [[2, 2], [3, 3, 3]]]),
+              weights_factory=_ragged_factory([[[], [1]], [[2, 3], [4, 5, 6]]]),
+              expected=[0, 1, 5, 15],  # no implied zeros
+              axis=None,
+          ),
+          dict(
+              tid="_r2_axis_neg_1",
+              x_factory=_ragged_factory([[], [1], [2, 2], [3, 3, 3]]),
+              weights_factory=_ragged_factory([[], [1], [2, 3], [4, 5, 6]]),
+              # no implied zeros
+              expected=[
+                  [0, 0, 0, 0],
+                  [0, 1, 0, 0],
+                  [0, 0, 5, 0],
+                  [0, 0, 0, 15],
+              ],
+              axis=-1,
+          ),
+      ),
+      (
+          dict(minlength=None, maxlength=None),
+          dict(minlength=3, maxlength=None),
+          dict(minlength=5, maxlength=None),
+          dict(minlength=None, maxlength=3),
+          dict(minlength=None, maxlength=5),
+          dict(minlength=2, maxlength=3),
+          dict(minlength=3, maxlength=5),
+          dict(minlength=5, maxlength=10),
+          dict(minlength=None, maxlength=0),
+      ),
+  )
+  def test_weights(
+      self,
+      tid,
+      x_factory,
+      weights_factory,
+      minlength,
+      maxlength,
+      expected,
+      axis,
+  ):
+    device_set = set([d.device_type for d in tf_config.list_physical_devices()])
+    if "GPU" in device_set and not test_util.is_xla_enabled():
+      self.skipTest(
+          "b/263004039 The DenseBincount GPU kernel does not support weights."
+          " unsorted_segment_sum should be used instead on GPU."
+      )
+    x = x_factory()
+    weights = weights_factory()
+    if axis == -1:
+      expected = _adjust_expected_rank2(expected, minlength, maxlength)
+    else:
+      expected = _adjust_expected_rank1(expected, minlength, maxlength)
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(
+                x,
+                weights=weights,
+                minlength=minlength,
+                maxlength=maxlength,
+                axis=axis,
+            )
+        ),
+    )
+
+  @parameterized.product(
+      (
+          dict(
+              tid="_r2",
+              x_factory=_ragged_factory([[], [1], [2, 2], [3, 3, 3]]),
+              expected=[0, 1, 1, 1],  # no implied zeros
+              axis=None,
+          ),
+          dict(
+              tid="_r3",
+              x_factory=_ragged_factory([[[], [1]], [[2, 2], [3, 3, 3]]]),
+              expected=[0, 1, 1, 1],  # no implied zeros
+              axis=None,
+          ),
+          dict(
+              tid="_r2_axis_neg_1",
+              x_factory=_ragged_factory([[], [1], [2, 2], [3, 3, 3]]),
+              # no implied zeros
+              expected=[[0, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1]],
+              axis=-1,
+          ),
+      ),
+      (
+          dict(minlength=None, maxlength=None),
+          dict(minlength=3, maxlength=None),
+          dict(minlength=5, maxlength=None),
+          dict(minlength=None, maxlength=3),
+          dict(minlength=None, maxlength=5),
+          dict(minlength=2, maxlength=3),
+          dict(minlength=3, maxlength=5),
+          dict(minlength=5, maxlength=10),
+          dict(minlength=None, maxlength=0),
+      ),
+  )
+  def test_binary_output(
+      self,
+      tid,
+      x_factory,
+      minlength,
+      maxlength,
+      expected,
+      axis=None,
+  ):
+    x = x_factory()
+    if axis == -1:
+      expected = _adjust_expected_rank2(expected, minlength, maxlength)
+    else:
+      expected = _adjust_expected_rank1(expected, minlength, maxlength)
+    self.assertAllEqual(
+        expected,
+        self.evaluate(
+            bincount_ops.bincount(
+                x,
+                minlength=minlength,
+                maxlength=maxlength,
+                binary_output=True,
+                axis=axis,
+            )
+        ),
+    )
+
+
+class TestSparseCount(test_util.TensorFlowTestCase, parameterized.TestCase):
+  @parameterized.named_parameters(
+      {
+          "testcase_name": "_no_maxlength",
+          "x": [[], [], [3, 0, 1], [], [5, 0, 4, 4]],
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [4, 0], [4, 4], [4, 5]],
+          "expected_values": [1, 1, 1, 1, 2, 1],
+          "expected_shape": [5, 6],
+      }, {
+          "testcase_name": "_maxlength",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "maxlength": 7,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [4, 0], [4, 4], [4, 5]],
+          "expected_values": [1, 1, 1, 1, 2, 1],
+          "expected_shape": [5, 7],
+      }, {
+          "testcase_name": "_maxlength_zero",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "maxlength": 0,
+          "expected_indices": np.empty([0, 2], dtype=np.int64),
+          "expected_values": [],
+          "expected_shape": [5, 0],
+      }, {
+          "testcase_name": "_minlength",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "minlength": 9,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [3, 7], [4, 0], [4, 4],
+                               [4, 5]],
+          "expected_values": [1, 1, 1, 1, 1, 2, 1],
+          "expected_shape": [5, 9],
+      }, {
+          "testcase_name": "_minlength_larger_values",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "minlength": 3,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [3, 7], [4, 0], [4, 4],
+                               [4, 5]],
+          "expected_values": [1, 1, 1, 1, 1, 2, 1],
+          "expected_shape": [5, 8],
+      }, {
+          "testcase_name": "_no_maxlength_binary",
+          "x": [[], [], [3, 0, 1], [], [5, 0, 4, 4]],
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [4, 0], [4, 4], [4, 5]],
+          "expected_values": [1, 1, 1, 1, 1, 1],
+          "expected_shape": [5, 6],
+          "binary_output": True,
+      }, {
+          "testcase_name": "_maxlength_binary",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "maxlength": 7,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [4, 0], [4, 4], [4, 5]],
+          "expected_values": [1, 1, 1, 1, 1, 1],
+          "expected_shape": [5, 7],
+          "binary_output": True,
+      }, {
+          "testcase_name": "_minlength_binary",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "minlength": 9,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [3, 7], [4, 0], [4, 4],
+                               [4, 5]],
+          "expected_values": [1, 1, 1, 1, 1, 1, 1],
+          "expected_shape": [5, 9],
+          "binary_output": True,
+      }, {
+          "testcase_name": "_minlength_larger_values_binary",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "minlength": 3,
+          "binary_output": True,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [3, 7], [4, 0], [4, 4],
+                               [4, 5]],
+          "expected_values": [1, 1, 1, 1, 1, 1, 1],
+          "expected_shape": [5, 8],
+      }, {
+          "testcase_name": "_no_maxlength_weights",
+          "x": [[], [], [3, 0, 1], [], [5, 0, 4, 4]],
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [4, 0], [4, 4], [4, 5]],
+          "expected_values": [0.5, 2, 6, 0.25, 8, 10],
+          "expected_shape": [5, 6],
+          "weights": [[], [], [6, 0.5, 2], [], [10, 0.25, 5, 3]],
+      }, {
+          "testcase_name": "_maxlength_weights",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "maxlength": 7,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [4, 0], [4, 4], [4, 5]],
+          "expected_values": [0.5, 2, 6, 0.25, 8, 10],
+          "expected_shape": [5, 7],
+          "weights": [[], [], [6, 0.5, 2], [14], [10, 0.25, 5, 3]],
+      }, {
+          "testcase_name": "_minlength_weights",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "minlength": 9,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [3, 7], [4, 0], [4, 4],
+                               [4, 5]],
+          "expected_values": [0.5, 2, 6, 14, 0.25, 8, 10],
+          "expected_shape": [5, 9],
+          "weights": [[], [], [6, 0.5, 2], [14], [10, 0.25, 5, 3]],
+      }, {
+          "testcase_name": "_minlength_larger_values_weights",
+          "x": [[], [], [3, 0, 1], [7], [5, 0, 4, 4]],
+          "minlength": 3,
+          "expected_indices": [[2, 0], [2, 1], [2, 3], [3, 7], [4, 0], [4, 4],
+                               [4, 5]],
+          "expected_values": [0.5, 2, 6, 14, 0.25, 8, 10],
+          "expected_shape": [5, 8],
+          "weights": [[], [], [6, 0.5, 2], [14], [10, 0.25, 5, 3]],
+      }, {
+          "testcase_name": "_1d",
+          "x": [3, 0, 1, 1],
+          "expected_indices": [[0], [1], [3]],
+          "expected_values": [1, 2, 1],
+          "expected_shape": [4],
+      }, {
+          "testcase_name": "_all_axes",
+          "x": [[], [], [3, 0, 1], [], [5, 0, 4, 4]],
+          "expected_indices": [[0], [1], [3], [4], [5]],
+          "expected_values": [2, 1, 1, 2, 1],
+          "expected_shape": [6],
+          "axis": None,
+      }, {
+          "testcase_name": "_large_inputs",
+          "x": [[1941591354222760687, 1748591354222760687],
+                [1941591354222760687, 1241591354229760689, 1241591354229760687]
+               ],
+          "expected_indices": [[1241591354229760687], [1241591354229760689],
+                               [1748591354222760687], [1941591354222760687]],
+          "expected_values": [1, 1, 1, 2],
+          "expected_shape": [1941591354222760687 + 1],
+          "axis": None
+      })
+  def test_ragged_input(self,
+                        x,
+                        expected_indices,
+                        expected_values,
+                        expected_shape,
+                        maxlength=None,
+                        minlength=None,
+                        binary_output=False,
+                        weights=None,
+                        axis=-1):
+    x_ragged = ragged_factory_ops.constant(x)
+    w = ragged_factory_ops.constant(weights) if weights is not None else None
+    y = sparse_ops.sparse_bincount(
+        x_ragged,
+        weights=w,
+        minlength=minlength,
+        maxlength=maxlength,
+        binary_output=binary_output,
+        axis=axis)
+    self.assertAllEqual(expected_indices, y.indices)
+    self.assertAllEqual(expected_values, y.values)
+    self.assertAllEqual(expected_shape, y.dense_shape)
+
+
+class TestSparseCountFailureModes(test_util.TensorFlowTestCase):
+  def test_dense_input_ragged_weights_fails(self):
+    x = np.array([[3, 2, 1], [5, 4, 4]], dtype=np.int32)
+    weights = ragged_factory_ops.constant([[6, 0.5, 2], [14], [10, 0.25, 5, 3]])
+    with self.assertRaisesRegex(ValueError, "must be a tf.Tensor"):
+      self.evaluate(sparse_ops.sparse_bincount(x, weights=weights, axis=-1))
+
+  def test_sparse_input_ragged_weights_fails(self):
+    x = sparse_ops.from_dense(
+        np.array([[3, 0, 1, 0], [0, 0, 0, 0], [5, 0, 4, 4]], dtype=np.int32))
+    weights = ragged_factory_ops.constant([[6, 0.5, 2], [14], [10, 0.25, 5, 3]])
+    with self.assertRaisesRegex(ValueError, "must be a SparseTensor"):
+      self.evaluate(sparse_ops.sparse_bincount(x, weights=weights, axis=-1))
+
+  def test_ragged_input_dense_weights_fails(self):
+    x = ragged_factory_ops.constant([[6, 1, 2], [14], [10, 1, 5, 3]])
+    weights = np.array([[3, 2, 1], [5, 4, 4]], dtype=np.int32)
+    with self.assertRaisesRegex(ValueError, "must be a RaggedTensor"):
+      self.evaluate(sparse_ops.sparse_bincount(x, weights=weights, axis=-1))
+
+  def test_ragged_input_sparse_weights_fails(self):
+    x = ragged_factory_ops.constant([[6, 1, 2], [14], [10, 1, 5, 3]])
+    weights = sparse_ops.from_dense(
+        np.array([[3, 0, 1, 0], [0, 0, 0, 0], [5, 0, 4, 4]], dtype=np.int32))
+    with self.assertRaisesRegex(ValueError, "must be a RaggedTensor"):
+      self.evaluate(sparse_ops.sparse_bincount(x, weights=weights, axis=-1))
+
+  def test_ragged_input_different_shape_fails(self):
+    x = ragged_factory_ops.constant([[6, 1, 2], [14], [10, 1, 5, 3]])
+    weights = ragged_factory_ops.constant([[6, 0.5, 2], [], [10, 0.25, 5, 3]])
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                "must have the same row splits"):
+      self.evaluate(sparse_ops.sparse_bincount(x, weights=weights, axis=-1))
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/ragged/ragged_bitcast_op_test.py

@@ -0,0 +1,123 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.bitcast."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSplitOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Cast to same-size dtype.
+      #=========================================================================
+      dict(
+          descr='int32 to int32 cast',
+          inputs=ragged_factory_ops.constant_value(
+              [[1, 2], [3]],
+              dtype=dtypes.int32,
+          ),
+          outputs=ragged_factory_ops.constant_value(
+              [[1, 2], [3]],
+              dtype=dtypes.int32,
+          )),
+      dict(
+          descr='int32 to uint32 cast',
+          inputs=ragged_factory_ops.constant_value(
+              [[1, 2], [-1]],
+              dtype=dtypes.int32,
+          ),
+          outputs=ragged_factory_ops.constant_value(
+              [[1, 2], [4294967295]],
+              dtype=dtypes.uint32,
+          )),
+      dict(
+          descr='uint32 to int32 cast',
+          inputs=ragged_factory_ops.constant_value(
+              [[1, 2], [4294967295]],
+              dtype=dtypes.uint32,
+          ),
+          outputs=ragged_factory_ops.constant_value(
+              [[1, 2], [-1]],
+              dtype=dtypes.int32,
+          )),
+      #=========================================================================
+      # Cast to larger dtype.
+      #=========================================================================
+      dict(
+          descr='int32 to int64 cast',
+          inputs=ragged_factory_ops.constant_value(
+              [[[1, 0], [2, 0]], [[3, 0]]],
+              dtype=dtypes.int32,
+              ragged_rank=1,
+          ),
+          outputs=ragged_factory_ops.constant_value(
+              [[1, 2], [3]],
+              dtype=dtypes.int64,
+          )),
+      #=========================================================================
+      # Cast to smaller dtype.
+      #=========================================================================
+      dict(
+          descr='int64 to int32 cast',
+          inputs=ragged_factory_ops.constant_value(
+              [[1, 2], [3]],
+              dtype=dtypes.int64,
+          ),
+          outputs=ragged_factory_ops.constant_value(
+              [[[1, 0], [2, 0]], [[3, 0]]],
+              dtype=dtypes.int32,
+              ragged_rank=1,
+          )),
+  ])  # pyformat: disable
+  def testBitcast(self, descr, inputs, outputs, name=None):
+    result = ragged_array_ops.bitcast(inputs, outputs.dtype, name)
+    self.assertEqual(result.dtype, outputs.dtype)
+    self.assertEqual(result.ragged_rank, outputs.ragged_rank)
+    self.assertAllEqual(result, outputs)
+
+  @parameterized.parameters([
+      dict(
+          descr='Upcast requires uniform inner dimension',
+          inputs=ragged_factory_ops.constant_value(
+              [[[1, 0], [2, 0]], [[3, 0]]],
+              dtype=dtypes.int32,
+              ragged_rank=2,
+          ),
+          cast_to_dtype=dtypes.int64,
+          exception=ValueError,
+          message='`input.flat_values` is required to have rank >= 2'),
+  ])  # pyformat: disable
+  def testBitcastError(self,
+                       descr,
+                       inputs,
+                       cast_to_dtype,
+                       exception,
+                       message,
+                       name=None):
+    with self.assertRaisesRegex(exception, message):
+      result = ragged_array_ops.bitcast(inputs, cast_to_dtype, name)
+      self.evaluate(result)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_boolean_mask_op_test.py

@@ -0,0 +1,278 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.boolean_mask."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedBooleanMaskOpTest(test_util.TensorFlowTestCase,
+                              parameterized.TestCase):
+  # Define short constants for true & false, so the data & mask can be lined
+  # up in the examples below.  This makes it easier to read the examples, to
+  # see which values should be kept vs. masked.
+  T = True
+  F = False
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Docstring examples
+      #=========================================================================
+      dict(
+          descr='Docstring example 1',
+          data=[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+          mask=[[T, F, T], [F, F, F], [T, F, F]],
+          expected=ragged_factory_ops.constant_value([[1, 3], [], [7]])),
+      dict(
+          descr='Docstring example 2',
+          data=ragged_factory_ops.constant_value([[1, 2, 3], [4], [5, 6]]),
+          mask=ragged_factory_ops.constant_value([[F, F, T], [F], [T, T]]),
+          expected=ragged_factory_ops.constant_value([[3], [], [5, 6]])),
+      dict(
+          descr='Docstring example 3',
+          data=ragged_factory_ops.constant_value([[1, 2, 3], [4], [5, 6]]),
+          mask=[True, False, True],
+          expected=ragged_factory_ops.constant_value([[1, 2, 3], [5, 6]])),
+      #=========================================================================
+      # Uniform data and uniform mask.
+      #=========================================================================
+      dict(
+          descr='data.shape=[7]; mask.shape=[7]',
+          data=[1, 2, 3, 4, 5, 6, 7],
+          mask=[T, F, T, T, F, F, F],
+          expected=[1, 3, 4]),
+      dict(
+          descr='data.shape=[5, 3]; mask.shape=[5]',
+          data=[[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12], [13, 14, 15]],
+          mask=[True, False, True, True, False],
+          expected=[[1, 2, 3], [7, 8, 9], [10, 11, 12]]),
+      dict(
+          descr='data.shape=[5, 3]; mask.shape=[5, 3]',
+          data=[[1, 2, 3], [4, 5, 6], [7, 8, 9], [0, 1, 2], [3, 4, 5]],
+          mask=[[F, F, F], [T, F, T], [T, T, T], [F, F, F], [T, T, F]],
+          expected=ragged_factory_ops.constant_value(
+              [[], [4, 6], [7, 8, 9], [], [3, 4]])),
+      dict(
+          descr='data.shape=[3, 2, 2]; mask.shape=[3]',
+          data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
+          mask=[F, F, T],
+          expected=[[[2, 4], [6, 8]]]),
+      dict(
+          descr='data.shape=[3, 2, 2]; mask.shape=[3]',
+          data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
+          mask=[F, F, T],
+          expected=[[[2, 4], [6, 8]]]),
+      dict(
+          descr='data.shape=[3, 2, 2]; mask.shape=[3, 2]',
+          data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
+          mask=[[T, F], [T, T], [F, F]],
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2]], [[5, 6], [7, 8]], []],
+              ragged_rank=1)),
+      dict(
+          descr='data.shape=[3, 2, 2]; mask.shape=[3, 2, 2]',
+          data=[[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]],
+          mask=[[[T, T], [F, T]], [[F, F], [F, F]], [[T, F], [T, T]]],
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2], [4]], [[], []], [[2], [6, 8]]])),
+      dict(
+          descr='data.shape=mask.shape=[2, 2, 2, 2]',
+          data=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                [[[2, 4], [6, 8]], [[1, 3], [5, 7]]]],
+          mask=[[[[T, T], [F, F]], [[T, F], [F, F]]],
+                [[[F, F], [F, F]], [[T, T], [T, F]]]],
+          expected=ragged_factory_ops.constant_value(
+              [[[[1, 2], []], [[5], []]], [[[], []], [[1, 3], [5]]]])),
+
+      #=========================================================================
+      # Ragged data and ragged mask.
+      #=========================================================================
+      dict(
+          descr='data.shape=[5, (D2)]; mask.shape=[5, (D2)]',
+          data=ragged_factory_ops.constant_value(
+              [[1, 2], [3, 4, 5, 6], [7, 8, 9], [], [1, 2, 3]]),
+          mask=ragged_factory_ops.constant_value(
+              [[F, F], [F, T, F, T], [F, F, F], [], [T, F, T]]),
+          expected=ragged_factory_ops.constant_value(
+              [[], [4, 6], [], [], [1, 3]])),
+      dict(
+          descr='data.shape=[3, (D2), (D3)]; mask.shape=[3, (D2)]',
+          data=ragged_factory_ops.constant_value(
+              [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4], [6, 8]]]),
+          mask=ragged_factory_ops.constant_value([[T, F], [T, T], [F, F]]),
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2]], [[5, 6], [7, 8]], []])),
+      dict(
+          descr='data.shape=[3, (D2), D3]; mask.shape=[3, (D2)]',
+          data=ragged_factory_ops.constant_value(
+              [[[1, 2], [3, 4]], [[5, 6], [7, 8], [2, 4]], [[6, 8]]],
+              ragged_rank=1),
+          mask=ragged_factory_ops.constant_value([[T, F], [T, T, F], [F]]),
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2]], [[5, 6], [7, 8]], []],
+              ragged_rank=1)),
+      dict(
+          descr='data.shape=[3, (D2), (D3)]; mask.shape=[3, (D2), (D3)]',
+          data=ragged_factory_ops.constant_value(
+              [[[1, 2], [3, 4]], [[5, 6], [7, 8]], [[2, 4]]]),
+          mask=ragged_factory_ops.constant_value(
+              [[[T, T], [F, T]], [[F, F], [F, F]], [[T, F]]]),
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2], [4]], [[], []], [[2]]])),
+      dict(
+          descr=('data.shape=[3, (D2), (D3), (D4)]; '
+                 'mask.shape=[3, (D2), (D3), (D4)]'),
+          data=ragged_factory_ops.constant_value(
+              [[[[1, 2], [3, 4]], [[5, 6]]], [[[2, 4], [6, 8]]]]),
+          mask=ragged_factory_ops.constant_value(
+              [[[[T, T], [F, F]], [[T, F]]], [[[F, F], [T, T]]]]),
+          expected=ragged_factory_ops.constant_value(
+              [[[[1, 2], []], [[5]]], [[[], [6, 8]]]])),
+
+      #=========================================================================
+      # Ragged mask and uniform data
+      #=========================================================================
+      dict(
+          descr='data.shape=[2, 3]; mask.shape=[2, (3)]',
+          data=[[1, 2, 3], [4, 5, 6]],
+          mask=ragged_factory_ops.constant_value([[T, F, F], [F, T, T]]),
+          expected=ragged_factory_ops.constant_value([[1], [5, 6]])),
+      dict(
+          descr='data.shape=[2, 3, 2]; mask.shape=[2, (3)]',
+          data=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 0], [2, 4]]],
+          mask=ragged_factory_ops.constant_value([[T, F, F], [F, T, T]]),
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2]], [[9, 0], [2, 4]]],
+              ragged_rank=1)),
+      dict(
+          descr='data.shape=[2, 3, 2]; mask.shape=[2, (3), 2]',
+          data=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 0], [2, 4]]],
+          mask=ragged_factory_ops.constant_value(
+              [[[T, F], [F, F], [T, T]], [[T, F], [F, T], [F, F]]],
+              ragged_rank=1),
+          expected=ragged_factory_ops.constant_value(
+              [[[1], [], [5, 6]], [[7], [0], []]])),
+
+      #=========================================================================
+      # Ragged data and uniform mask.
+      #=========================================================================
+      dict(
+          descr='data.shape=[4, (D2)]; mask.shape=[4]',
+          data=ragged_factory_ops.constant_value([[1, 2, 3], [4], [], [5, 6]]),
+          mask=[T, F, T, F],
+          expected=ragged_factory_ops.constant_value([[1, 2, 3], []])),
+      dict(
+          descr='data.shape=[4, (D2), (D3)]; mask.shape=[4]',
+          data=ragged_factory_ops.constant_value(
+              [[[1, 2, 3]], [[4], []], [[5, 6]], []]),
+          mask=[T, F, T, T],
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2, 3]], [[5, 6]], []])),
+      dict(
+          descr='data.shape=[4, (D2), 2]; mask.shape=[4]',
+          data=ragged_factory_ops.constant_value(
+              [[[1, 2], [3, 4]], [], [[5, 6]], [[7, 8], [9, 0], [1, 2]]],
+              ragged_rank=1),
+          mask=[T, F, F, T],
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2], [3, 4]], [[7, 8], [9, 0], [1, 2]]],
+              ragged_rank=1)),
+      dict(
+          descr='data.shape=[4, (D2), 2]; mask.shape=[4]',
+          data=ragged_factory_ops.constant_value(
+              [[[1, 2], [3, 4]], [], [[5, 6]], [[7, 8], [9, 0], [1, 2]]],
+              ragged_rank=1),
+          mask=[T, F, F, T],
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2], [3, 4]], [[7, 8], [9, 0], [1, 2]]],
+              ragged_rank=1)),
+      dict(
+          descr='data.shape=[1, (2)]; mask.shape=[1, 2]',
+          data=ragged_factory_ops.constant_value([[1, 2]]),
+          mask=[[T, F]],
+          expected=ragged_factory_ops.constant_value([[1]])),
+      dict(
+          descr='data.shape=[2, (2), (D3)]; mask.shape=[2, 2]',
+          data=ragged_factory_ops.constant_value(
+              [[[1], [2, 3]], [[], [4, 5, 6]]]),
+          mask=[[T, F], [T, T]],
+          expected=ragged_factory_ops.constant_value([[[1]], [[], [4, 5, 6]]])),
+      dict(
+          descr='data.shape=[2, (2), 3]; mask.shape=[2, 2]',
+          data=ragged_factory_ops.constant_value(
+              [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [2, 4, 6]]],
+              ragged_rank=1),
+          mask=[[T, F], [T, T]],
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2, 3]], [[7, 8, 9], [2, 4, 6]]],
+              ragged_rank=1)),
+      dict(
+          descr='data.shape=[2, (2), 3]; mask.shape=[2, 2, 3]',
+          data=ragged_factory_ops.constant_value(
+              [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [2, 4, 6]]],
+              ragged_rank=1),
+          mask=[[[T, F, F], [T, F, T]], [[T, F, T], [F, F, F]]],
+          expected=ragged_factory_ops.constant_value(
+              [[[1], [4, 6]], [[7, 9], []]])),
+  ])  # pyformat: disable
+  def testBooleanMask(self, descr, data, mask, expected):
+    actual = ragged_array_ops.boolean_mask(data, mask)
+    self.assertAllEqual(actual, expected)
+
+  def testErrors(self):
+    if not context.executing_eagerly():
+      self.assertRaisesRegex(ValueError,
+                             r'mask\.shape\.ndims must be known statically',
+                             ragged_array_ops.boolean_mask, [[1, 2]],
+                             array_ops.placeholder(dtypes.bool))
+
+    self.assertRaises(TypeError, ragged_array_ops.boolean_mask, [[1, 2]],
+                      [[0, 1]])
+    self.assertRaisesRegex(
+        ValueError, 'Tensor conversion requested dtype bool for '
+        'RaggedTensor with dtype int32', ragged_array_ops.boolean_mask,
+        ragged_factory_ops.constant([[1, 2]]),
+        ragged_factory_ops.constant([[0, 0]]))
+
+    self.assertRaisesRegex(ValueError,
+                           r'Shapes \(1, 2\) and \(1, 3\) are incompatible',
+                           ragged_array_ops.boolean_mask, [[1, 2]],
+                           [[True, False, True]])
+
+    self.assertRaisesRegex(errors.InvalidArgumentError,
+                           r'Inputs must have identical ragged splits',
+                           ragged_array_ops.boolean_mask,
+                           ragged_factory_ops.constant([[1, 2]]),
+                           ragged_factory_ops.constant([[True, False, True]]))
+
+    self.assertRaisesRegex(ValueError, 'mask cannot be scalar',
+                           ragged_array_ops.boolean_mask, [[1, 2]], True)
+
+    self.assertRaisesRegex(ValueError, 'mask cannot be scalar',
+                           ragged_array_ops.boolean_mask,
+                           ragged_factory_ops.constant([[1, 2]]), True)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_check_ops.py

@@ -0,0 +1,27 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Asserts and Boolean Checks for RaggedTensors."""
+
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import dispatch
+
+
+@dispatch.dispatch_for_api(check_ops.assert_type)
+def assert_type(tensor: ragged_tensor.Ragged, tf_type, message=None, name=None):
+  return check_ops.assert_type(tensor.flat_values, tf_type,
+                               message=message, name=name)
+
+

tensorflow/python/ops/ragged/ragged_concat_op_test.py

@@ -0,0 +1,325 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.concat."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedConcatOpTest(test_util.TensorFlowTestCase,
+                         parameterized.TestCase):
+
+  def _rt_inputs_to_tensors(self, rt_inputs, ragged_ranks=None):
+    if ragged_ranks is None:
+      ragged_ranks = [None] * len(rt_inputs)
+    return [  # pylint: disable=g-long-ternary
+        ragged_factory_ops.constant(rt_input, ragged_rank=rrank)
+        if rrank != 0 else constant_op.constant(rt_input)
+        for (rt_input, rrank) in zip(rt_inputs, ragged_ranks)
+    ]
+
+  @parameterized.parameters(
+      dict(
+          descr='Two rank-2 inputs with empty value axis=1',
+          rt_inputs=([[]], [[]]),
+          axis=1,
+          expected=[[]]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=0',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21']],   # shape=(3, None)
+              [['b00'], ['b10']]),                    # shape=(2, None)
+          axis=0,
+          expected=[[b'a00', b'a01'], [], [b'a20', b'a21'], [b'b00'],
+                    [b'b10']]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=1',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00'], ['b10', 'b11', 'b12'], ['b20']]),    # shape=(3, None)
+          axis=1,
+          expected=[
+              [b'a00', b'a01', b'b00'],
+              [b'b10', b'b11', b'b12'],
+              [b'a20', b'a21', b'a22', b'b20']]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=-2',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21']],   # shape=(3, None)
+              [['b00'], ['b10']]),                    # shape=(2, None)
+          axis=-2,
+          expected=[[b'a00', b'a01'], [], [b'a20', b'a21'], [b'b00'],
+                    [b'b10']]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=-1',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00'], ['b10', 'b11', 'b12'], ['b20']]),    # shape=(3, None)
+          axis=-1,
+          expected=[
+              [b'a00', b'a01', b'b00'],
+              [b'b10', b'b11', b'b12'],
+              [b'a20', b'a21', b'a22', b'b20']],
+          expected_shape=[3, None]),
+      dict(
+          descr='Three rank-2 inputs (ragged_rank=1), axis=0',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00'], ['b10']],                            # shape=(2, None)
+              [['c00'], ['c10', 'c11'], ['c21']]),           # shape=(3, None)
+          axis=0,
+          expected=[[b'a00', b'a01'], [], [b'a20', b'a21', b'a22'], [b'b00'],
+                    [b'b10'], [b'c00'], [b'c10', b'c11'], [b'c21']]),
+      dict(
+          descr='Three rank-2 inputs (ragged_rank=1), axis=1',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00'], ['b10', 'b11', 'b12'], ['b20']],     # shape=(3, None)
+              [[], ['c10', 'c11'], ['c20', 'c21']]),         # shape=(3, None)
+          axis=1,
+          expected=[
+              [b'a00', b'a01', b'b00'],
+              [b'b10', b'b11', b'b12', b'c10', b'c11'],
+              [b'a20', b'a21', b'a22', b'b20', b'c20', b'c21']]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=0',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']],
+               [['a100', 'a101', 'a102'], ['a110', 'a111']]],
+              [[['b000']], [['b100', 'b101'], ['b110']]],
+              [[], [['c100', 'c101', 'c102', 'c103']], [[], ['c210', 'c211']]]),
+          axis=0,
+          expected=[
+              [[b'a000', b'a001'], [b'a010']],
+              [[b'a100', b'a101', b'a102'], [b'a110', b'a111']],
+              [[b'b000']],
+              [[b'b100', b'b101'], [b'b110']],
+              [],
+              [[b'c100', b'c101', b'c102', b'c103']],
+              [[], [b'c210', b'c211']]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=1',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']],
+               [['a100', 'a101', 'a102'], ['a110', 'a111']]],
+              [[['b000']], [['b100', 'b101'], ['b110']]],
+              [[], [[], ['c110', 'c111']]]),
+          axis=1,
+          expected=[
+              [[b'a000', b'a001'], [b'a010'], [b'b000']],
+              [[b'a100', b'a101', b'a102'], [b'a110', b'a111'],
+               [b'b100', b'b101'], [b'b110'], [], [b'c110', b'c111']]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=2',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']],
+               [['a100', 'a101', 'a102'], ['a110', 'a111']]],
+              [[[], ['b010', 'b011']], [['b100', 'b101'], ['b110']]],
+              [[['c000'], ['c010']], [[], ['c110', 'c111']]]),
+          axis=2,
+          expected=[
+              [[b'a000', b'a001', b'c000'],
+               [b'a010', b'b010', b'b011', b'c010']],
+              [[b'a100', b'a101', b'a102', b'b100', b'b101'],
+               [b'a110', b'a111', b'b110', b'c110', b'c111']]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=-1',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']],
+               [['a100', 'a101', 'a102'], ['a110', 'a111']]],
+              [[[], ['b010', 'b011']], [['b100', 'b101'], ['b110']]],
+              [[['c000'], ['c010']], [[], ['c110', 'c111']]]),
+          axis=-1,
+          expected=[
+              [[b'a000', b'a001', b'c000'],
+               [b'a010', b'b010', b'b011', b'c010']],
+              [[b'a100', b'a101', b'a102', b'b100', b'b101'],
+               [b'a110', b'a111', b'b110', b'c110', b'c111']]]),
+      dict(
+          descr='ragged_concat([uniform, ragged, uniform], axis=1)',
+          ragged_ranks=[0, 1, 0],
+          rt_inputs=(
+              [['0('], ['1('], ['2(']],                   # shape=(3, 1)
+              [['b00'], ['b10', 'b11', 'b12'], ['b20']],  # shape=(3, None)
+              [[')0'], [')1'], [')2']]),                  # shape=(3, 1)
+          axis=1,
+          expected=[
+              [b'0(', b'b00', b')0'],
+              [b'1(', b'b10', b'b11', b'b12', b')1'],
+              [b'2(', b'b20', b')2']]),
+      dict(
+          descr='ragged_concat([uniform, uniform], axis=0)',
+          ragged_ranks=[0, 0],
+          rt_inputs=(
+              [['a00', 'a01'], ['a10', 'a11'], ['a20', 'a21']],  # shape=(3, 2)
+              [['b00', 'b01', 'b02'], ['b10', 'b11', 'b12']]),   # shape=(2, 3)
+          axis=0,
+          expected=[
+              [b'a00', b'a01'], [b'a10', b'a11'], [b'a20', b'a21'],
+              [b'b00', b'b01', b'b02'], [b'b10', b'b11', b'b12']],
+          expected_ragged_rank=1),
+      dict(
+          descr='ragged_concat([uniform, ragged], axis=0)',
+          ragged_ranks=[0, 1],
+          rt_inputs=(
+              [['a00', 'a01'], ['a10', 'a11'], ['a20', 'a21']],  # shape=(3, 2)
+              [['b00', 'b01', 'b02'], ['b10', 'b11', 'b12']]),   # shape=(2, 3)
+          axis=0,
+          expected=[
+              [b'a00', b'a01'], [b'a10', b'a11'], [b'a20', b'a21'],
+              [b'b00', b'b01', b'b02'], [b'b10', b'b11', b'b12']]),
+      dict(
+          descr='ragged_concat([uniform, ragged], axis=0) with rank-3 inputs',
+          ragged_ranks=[0, 2],
+          rt_inputs=(
+              [[[0, 1], [2, 3]], [[4, 5], [6, 7]]],  # shape = (2, 2, 2)
+              [[[8], [8, 8]]]),                      # shape = (2, None, None)
+          axis=0,
+          expected=[[[0, 1], [2, 3]], [[4, 5], [6, 7]], [[8], [8, 8]]]),
+      dict(
+          descr='Two rank-3 inputs with ragged_rank=1, axis=-1',
+          ragged_ranks=[1, 1],
+          rt_inputs=(
+              [[[0, 1], [2, 3], [4, 5]], [], [[6, 7], [8, 9]]],
+              [[[9, 8], [7, 6], [5, 4]], [], [[3, 2], [1, 0]]]),
+          axis=-1,
+          expected=[
+              [[0, 1, 9, 8], [2, 3, 7, 6], [4, 5, 5, 4]], [],
+              [[6, 7, 3, 2], [8, 9, 1, 0]]],
+          expected_ragged_rank=1),
+      dict(
+          descr='ragged_concat([vector, vector], axis=0)',
+          ragged_ranks=[0, 0],
+          rt_inputs=([1, 2, 3], [4, 5, 6]),
+          axis=0,
+          expected=[1, 2, 3, 4, 5, 6]),
+      dict(
+          descr='One input (so ragged_concat is a noop)',
+          rt_inputs=([['a00', 'a01'], [], ['a20', 'a21']],),
+          axis=0,
+          expected=[[b'a00', b'a01'], [], [b'a20', b'a21']]),
+  )  # pyformat: disable
+  def testRaggedConcat(self,
+                       descr,
+                       rt_inputs,
+                       axis,
+                       expected,
+                       ragged_ranks=None,
+                       expected_ragged_rank=None,
+                       expected_shape=None):
+    rt_inputs = self._rt_inputs_to_tensors(rt_inputs, ragged_ranks)
+    concatenated = ragged_concat_ops.concat(rt_inputs, axis)
+    if expected_ragged_rank is not None:
+      self.assertEqual(concatenated.ragged_rank, expected_ragged_rank)
+    if expected_shape is not None:
+      self.assertEqual(concatenated.shape.as_list(), expected_shape)
+    self.assertAllEqual(concatenated, expected)
+
+  @parameterized.parameters(
+      dict(
+          rt_inputs=(),
+          axis=0,
+          error=ValueError,
+          message=r'rt_inputs may not be empty\.'),
+      dict(
+          rt_inputs=([[1, 2]], [[3, 4]]),
+          axis=r'foo',
+          error=TypeError,
+          message='axis must be an int'),
+      dict(
+          rt_inputs=([[1, 2]], [[3, 4]]),
+          axis=-3,
+          error=ValueError,
+          message='axis=-3 out of bounds: expected -2<=axis<2'),
+      dict(
+          rt_inputs=([[1, 2]], [[3, 4]]),
+          axis=2,
+          error=ValueError,
+          message='axis=2 out of bounds: expected -2<=axis<2'),
+      dict(
+          ragged_ranks=(0, 0),
+          rt_inputs=([[1, 2]], [[3, 4], [5, 6]]),
+          axis=1,
+          error=(ValueError, errors.InvalidArgumentError)),
+  )
+  def testStaticError(self,
+                      rt_inputs,
+                      axis,
+                      error,
+                      message=None,
+                      ragged_ranks=None):
+    rt_inputs = self._rt_inputs_to_tensors(rt_inputs, ragged_ranks)
+    self.assertRaisesRegex(error, message, ragged_concat_ops.concat, rt_inputs,
+                           axis)
+
+  @parameterized.parameters([
+      dict(
+          ragged_ranks=(1, 1),
+          rt_inputs=([[1, 2]], [[3, 4], [5, 6]]),
+          axis=1,
+          error=errors.InvalidArgumentError,
+          message=(
+              r'Input tensors at index 0 \(=x\) and 1 \(=y\) have incompatible'
+              r' shapes\.'
+          ),
+      ),
+  ])
+  def testRuntimeError(self, rt_inputs, axis, error, message,
+                       ragged_ranks=None):
+    if context.executing_eagerly():
+      return
+    rt_inputs = [
+        array_ops.placeholder_with_default(rt, shape=None) for rt in rt_inputs
+    ]
+    concatenated = ragged_concat_ops.concat(rt_inputs, axis)
+    with self.assertRaisesRegex(error, message):
+      self.evaluate(concatenated)
+
+  def testNegativeAxisWithUnknownRankError(self):
+    if context.executing_eagerly():
+      return
+    rt_inputs = [
+        array_ops.placeholder(dtypes.int64),
+        array_ops.placeholder(dtypes.int64)
+    ]
+    self.assertRaisesRegex(
+        ValueError,
+        r'axis=-1 may only be negative if ndims is statically known.',
+        ragged_concat_ops.concat, rt_inputs, -1)
+
+  def testSingleTensorInput(self):
+    """Tests ragged_concat with a single tensor input.
+
+    Usually, we pass a list of values in for rt_inputs.  However, you can
+    also pass in a single value (as with tf.concat), in which case it simply
+    returns that tensor.  This test exercises that path.
+    """
+    rt_inputs = ragged_factory_ops.constant([[1, 2], [3, 4]])
+    concatenated = ragged_concat_ops.concat(rt_inputs, 0)
+    self.assertAllEqual(concatenated, [[1, 2], [3, 4]])
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_concat_ops.py

@@ -0,0 +1,330 @@
+# 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.
+# ==============================================================================
+"""Concat and stack operations for RaggedTensors."""
+
+import typing
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_gather_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_util
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@dispatch.dispatch_for_api(array_ops.concat)
+def concat(values: typing.List[ragged_tensor.RaggedOrDense], axis, name=None):
+  """Concatenates potentially ragged tensors along one dimension.
+
+  Given a list of tensors with the same rank `K` (`K >= axis`), returns a
+  rank-`K` `RaggedTensor` `result` such that `result[i0...iaxis]` is the
+  concatenation of `[rt[i0...iaxis] for rt in values]`.
+
+  Args:
+    values: A list of potentially ragged tensors.  May not be empty. All
+      `values` must have the same rank and the same dtype; but unlike
+      `tf.concat`, they can have arbitrary shapes.
+    axis: A python integer, indicating the dimension along which to concatenate.
+      (Note: Unlike `tf.concat`, the `axis` parameter must be statically known.)
+        Negative values are supported only if the rank of at least one
+        `values` value is statically known.
+    name: A name prefix for the returned tensor (optional).
+
+  Returns:
+    A `RaggedTensor` with rank `K`.
+    `result.ragged_rank=max(axis, max(rt.ragged_rank for rt in values]))`.
+
+  Raises:
+    ValueError: If `values` is empty, if `axis` is out of bounds or if
+      the input tensors have different ranks.
+
+  #### Example:
+
+  >>> t1 = tf.ragged.constant([[1, 2], [3, 4, 5]])
+  >>> t2 = tf.ragged.constant([[6], [7, 8, 9]])
+  >>> tf.concat([t1, t2], axis=0)
+  <tf.RaggedTensor [[1, 2], [3, 4, 5], [6], [7, 8, 9]]>
+  >>> tf.concat([t1, t2], axis=1)
+  <tf.RaggedTensor [[1, 2, 6], [3, 4, 5, 7, 8, 9]]>
+  """
+  if not isinstance(values, (list, tuple)):
+    values = [values]
+  with ops.name_scope(name, 'RaggedConcat', values):
+    return _ragged_stack_concat_helper(values, axis, stack_values=False)
+
+
+@tf_export('ragged.stack')
+@dispatch.add_dispatch_support
+@dispatch.dispatch_for_api(array_ops_stack.stack)
+def stack(values: typing.List[ragged_tensor.RaggedOrDense],
+          axis=0,
+          name=None):
+  """Stacks a list of rank-`R` tensors into one rank-`(R+1)` `RaggedTensor`.
+
+  Given a list of tensors or ragged tensors with the same rank `R`
+  (`R >= axis`), returns a rank-`R+1` `RaggedTensor` `result` such that
+  `result[i0...iaxis]` is `[value[i0...iaxis] for value in values]`.
+
+  #### Examples:
+
+  >>> # Stacking two ragged tensors.
+  >>> t1 = tf.ragged.constant([[1, 2], [3, 4, 5]])
+  >>> t2 = tf.ragged.constant([[6], [7, 8, 9]])
+  >>> tf.ragged.stack([t1, t2], axis=0)
+  <tf.RaggedTensor [[[1, 2], [3, 4, 5]], [[6], [7, 8, 9]]]>
+  >>> tf.ragged.stack([t1, t2], axis=1)
+  <tf.RaggedTensor [[[1, 2], [6]], [[3, 4, 5], [7, 8, 9]]]>
+
+  >>> # Stacking two dense tensors with different sizes.
+  >>> t3 = tf.constant([[1, 2, 3], [4, 5, 6]])
+  >>> t4 = tf.constant([[5], [6], [7]])
+  >>> tf.ragged.stack([t3, t4], axis=0)
+  <tf.RaggedTensor [[[1, 2, 3], [4, 5, 6]], [[5], [6], [7]]]>
+
+  Args:
+    values: A list of `tf.Tensor` or `tf.RaggedTensor`.  May not be empty. All
+      `values` must have the same rank and the same dtype; but unlike
+      `tf.stack`, they can have arbitrary dimension sizes.
+    axis: A python integer, indicating the dimension along which to stack.
+      (Note: Unlike `tf.stack`, the `axis` parameter must be statically known.)
+      Negative values are supported only if the rank of at least one
+      `values` value is statically known.
+    name: A name prefix for the returned tensor (optional).
+
+  Returns:
+    A `RaggedTensor` with rank `R+1` (if `R>0`).
+    If `R==0`, then the result will be returned as a 1D `Tensor`, since
+    `RaggedTensor` can only be used when `rank>1`.
+    `result.ragged_rank=1+max(axis, max(rt.ragged_rank for rt in values]))`.
+
+  Raises:
+    ValueError: If `values` is empty, if `axis` is out of bounds or if
+      the input tensors have different ranks.
+  """
+  if not isinstance(values, (list, tuple)):
+    values = [values]
+  with ops.name_scope(name, 'RaggedConcat', values):
+    return _ragged_stack_concat_helper(values, axis, stack_values=True)
+
+
+def _ragged_stack_concat_helper(rt_inputs, axis, stack_values):
+  """Helper function to concatenate or stack ragged tensors.
+
+  Args:
+    rt_inputs: A list of RaggedTensors or Tensors to combine.
+    axis: The axis along which to concatenate or stack.
+    stack_values: A boolean -- if true, then stack values; otherwise,
+      concatenate them.
+
+  Returns:
+    A RaggedTensor.
+  Raises:
+    ValueError: If rt_inputs is empty, or if axis is out of range.
+  """
+  # Validate parameters.
+  if not rt_inputs:
+    raise ValueError('rt_inputs may not be empty.')
+
+  # Convert input tensors.
+  rt_inputs = [
+      ragged_tensor.convert_to_tensor_or_ragged_tensor(
+          rt_input, name='rt_input') for rt_input in rt_inputs
+  ]
+  row_splits_dtype, rt_inputs = ragged_tensor.match_row_splits_dtypes(
+      *rt_inputs, return_dtype=True)
+  rt_inputs = list(rt_inputs)
+
+  # Special case: if there's only one input, then return it as-is.
+  if len(rt_inputs) == 1 and not stack_values:
+    return rt_inputs[0]
+
+  # Check the rank (number of dimensions) of the input tensors.
+  ndims = None
+  for rt in rt_inputs:
+    if ndims is None:
+      ndims = rt.shape.ndims
+    else:
+      rt.shape.assert_has_rank(ndims)
+
+  out_ndims = ndims if (ndims is None or not stack_values) else ndims + 1
+  axis = array_ops.get_positive_axis(axis, out_ndims)
+
+  if stack_values and ndims == 1 and axis == 0:
+    return ragged_tensor.RaggedTensor.from_row_lengths(
+        values=array_ops.concat(rt_inputs, axis=0),
+        row_lengths=array_ops.concat([array_ops.shape(r) for r in rt_inputs],
+                                     axis=0))
+
+  # If all the inputs are Tensors, and we're combining the final dimension,
+  # then we can delegate to the tf.stack/tf.concat operation, and return a
+  # Tensor.
+  if all(not ragged_tensor.is_ragged(rt) for rt in rt_inputs):
+    if ndims is not None and (axis == out_ndims - 1 or axis == ndims - 1):
+      if stack_values:
+        return array_ops_stack.stack(rt_inputs, axis)
+      else:
+        return array_ops.concat(rt_inputs, axis)
+
+  # Convert any Tensor inputs to RaggedTensors.  This makes it
+  # possible to concatenate Tensors and RaggedTensors together.
+  for i in range(len(rt_inputs)):
+    if not ragged_tensor.is_ragged(rt_inputs[i]):
+      rt_inputs[i] = ragged_tensor.RaggedTensor.from_tensor(
+          rt_inputs[i], ragged_rank=1, row_splits_dtype=row_splits_dtype)
+
+  # Convert the input tensors to all have the same ragged_rank.
+  ragged_rank = max(max(rt.ragged_rank for rt in rt_inputs), 1)
+  rt_inputs = [_increase_ragged_rank_to(rt, ragged_rank, row_splits_dtype)
+               for rt in rt_inputs]
+
+  if axis == 0:
+    return _ragged_stack_concat_axis_0(rt_inputs, stack_values)
+  elif axis == 1:
+    return _ragged_stack_concat_axis_1(rt_inputs, stack_values)
+  else:  # axis > 1: recurse.
+    values = [rt.values for rt in rt_inputs]
+    splits = [[rt_input.row_splits] for rt_input in rt_inputs]
+    with ops.control_dependencies(ragged_util.assert_splits_match(splits)):
+      return ragged_tensor.RaggedTensor.from_row_splits(
+          _ragged_stack_concat_helper(values, axis - 1, stack_values),
+          splits[0][0], validate=False)
+
+
+def _ragged_stack_concat_axis_0(rt_inputs, stack_values):
+  """Helper function to concatenate or stack ragged tensors along axis 0.
+
+  Args:
+    rt_inputs: A list of RaggedTensors, all with the same rank and ragged_rank.
+    stack_values: Boolean.  If true, then stack values; otherwise, concatenate
+      them.
+
+  Returns:
+    A RaggedTensor.
+  """
+  # Concatenate the inner values together.
+  flat_values = [rt.flat_values for rt in rt_inputs]
+  concatenated_flat_values = array_ops.concat(flat_values, axis=0)
+
+  # Concatenate the splits together for each ragged dimension (adjusting
+  # split offsets as necessary).
+  nested_splits = [rt.nested_row_splits for rt in rt_inputs]
+  ragged_rank = rt_inputs[0].ragged_rank
+  concatenated_nested_splits = [
+      _concat_ragged_splits([ns[dim]
+                             for ns in nested_splits])
+      for dim in range(ragged_rank)
+  ]
+
+  # If we are performing a stack operation, then add another splits.
+  if stack_values:
+    stack_lengths = array_ops_stack.stack([rt.nrows() for rt in rt_inputs])
+    stack_splits = ragged_util.lengths_to_splits(stack_lengths)
+    concatenated_nested_splits.insert(0, stack_splits)
+
+  return ragged_tensor.RaggedTensor.from_nested_row_splits(
+      concatenated_flat_values, concatenated_nested_splits, validate=False)
+
+
+def _ragged_stack_concat_axis_1(rt_inputs, stack_values):
+  """Helper function to concatenate or stack ragged tensors along axis 1.
+
+  Args:
+    rt_inputs: A list of RaggedTensors, all with the same rank and ragged_rank.
+    stack_values: Boolean.  If true, then stack values; otherwise, concatenate
+      them.
+
+  Returns:
+    A RaggedTensor.
+  """
+  num_inputs = len(rt_inputs)
+
+  nrows_checks = []
+  rt_nrows = rt_inputs[0].nrows()
+  for index, rt in enumerate(rt_inputs[1:]):
+    nrows_checks.append(
+        check_ops.assert_equal(
+            rt_nrows,
+            rt.nrows(),
+            message=(
+                f'Input tensors at index 0 (=x) and {index+1} (=y) have'
+                ' incompatible shapes.'
+            ),
+        )
+    )
+
+  with ops.control_dependencies(nrows_checks):
+    # Concatenate the inputs together to put them in a single ragged tensor.
+    concatenated_rt = _ragged_stack_concat_axis_0(rt_inputs, stack_values=False)
+
+    # Use ragged.gather to permute the rows of concatenated_rt.  In particular,
+    #   permuted_rt = [rt_inputs[0][0], ..., rt_inputs[N][0],
+    #                  rt_inputs[0][1], ..., rt_inputs[N][1],
+    #                      ...,
+    #                  rt_inputs[0][M], ..., rt_input[N][M]]
+    # where `N=num_inputs-1` and `M=rt_nrows-1`.
+    row_indices = math_ops.range(rt_nrows * num_inputs)
+    row_index_matrix = array_ops.reshape(row_indices, [num_inputs, -1])
+    transposed_row_index_matrix = array_ops.transpose(row_index_matrix)
+    row_permutation = array_ops.reshape(transposed_row_index_matrix, [-1])
+    permuted_rt = ragged_gather_ops.gather(concatenated_rt, row_permutation)
+
+    if stack_values:
+      # Add a new splits tensor to group together the values.
+      stack_splits = math_ops.range(0, rt_nrows * num_inputs + 1, num_inputs)
+      _copy_row_shape(rt_inputs, stack_splits)
+      return ragged_tensor.RaggedTensor.from_row_splits(
+          permuted_rt, stack_splits, validate=False)
+    else:
+      # Merge together adjacent rows by dropping the row-split indices that
+      # separate them.
+      concat_splits = permuted_rt.row_splits[::num_inputs]
+      _copy_row_shape(rt_inputs, concat_splits)
+      return ragged_tensor.RaggedTensor.from_row_splits(
+          permuted_rt.values, concat_splits, validate=False)
+
+
+def _copy_row_shape(rt_inputs, splits):
+  """Sets splits.shape to [rt[shape[0]+1] for each rt in rt_inputs."""
+  for rt in rt_inputs:
+    if rt.shape[0] is not None:
+      splits.set_shape(tensor_shape.TensorShape(rt.shape[0] + 1))
+
+
+def _increase_ragged_rank_to(rt_input, ragged_rank, row_splits_dtype):
+  """Adds ragged dimensions to `rt_input` so it has the desired ragged rank."""
+  if ragged_rank > 0:
+    if not ragged_tensor.is_ragged(rt_input):
+      rt_input = ragged_tensor.RaggedTensor.from_tensor(
+          rt_input, row_splits_dtype=row_splits_dtype)
+    if rt_input.ragged_rank < ragged_rank:
+      rt_input = rt_input.with_values(
+          _increase_ragged_rank_to(rt_input.values, ragged_rank - 1,
+                                   row_splits_dtype))
+  return rt_input
+
+
+def _concat_ragged_splits(splits_list):
+  """Concatenates a list of RaggedTensor splits to form a single splits."""
+  pieces = [splits_list[0]]
+  splits_offset = splits_list[0][-1]
+  for splits in splits_list[1:]:
+    pieces.append(splits[1:] + splits_offset)
+    splits_offset += splits[-1]
+  return array_ops.concat(pieces, axis=0)

tensorflow/python/ops/ragged/ragged_config.py

@@ -0,0 +1,29 @@
+# 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.
+# ==============================================================================
+"""Configuration parameters for RaggedTensors."""
+
+
+def auto_cast_partition_dtype():
+  """Whether incompatible row-partitioning dtypes should be auto-converted.
+
+  If true, then operations that combine RaggedTensors but have different
+  row-partitioning tensor dtypes will be automatically cast to a
+  compatible dtype (`tf.int64`).  If false, then such operations will result
+  in an error.
+
+  Returns:
+    `bool`
+  """
+  return False

tensorflow/python/ops/ragged/ragged_const_op_test.py

@@ -0,0 +1,411 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_factory_ops.constant."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import ragged
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedConstOpTest(test_util.TensorFlowTestCase,
+                        parameterized.TestCase):
+
+  @parameterized.parameters(
+      #=========================================================================
+      # 0-dimensional tensors.
+      dict(pylist=b'x', expected_shape=()),
+
+      #=========================================================================
+      # 1-dimensional tensors.
+      dict(pylist=[1, 2, 3], expected_shape=(3,)),
+
+      #=========================================================================
+      # 2-dimensional tensors.
+      dict(pylist=[[1, 2, 3], [4], [5, 6]], expected_shape=(3, None)),
+      dict(pylist=[[1, 2, 3], [4, 5, 6], [7, 8, 9]], expected_shape=(3, None)),
+
+      #=========================================================================
+      # 3-dimensional tensors.
+      dict(
+          pylist=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          expected_shape=(3, None, None)),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          inner_shape=(2,),
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          inner_shape=(2,),
+          expected_shape=(3, None, 2)),
+      # 3-dimensional tensors with numpy arrays
+      dict(
+          pylist=[[[1, 2], np.array([3, np.array(4)])],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          expected_shape=(3, None, None)),
+      dict(
+          pylist=[[[1, 2], np.array([3, np.array(4)])],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[[1, 2], np.array([3, np.array(4)])],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          inner_shape=(2,),
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[[1, 2], np.array([3, np.array(4)])],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          inner_shape=(2,),
+          expected_shape=(3, None, 2)),
+      #=========================================================================
+      # 4-dimensional tensors.
+      dict(
+          pylist=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                  [[[2, 4], [6, 8]], [[1, 5], [7, 9]]]],
+          expected_shape=(2, None, None, None)),
+      dict(
+          pylist=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                  [[[2, 4], [6, 8]], [[1, 5], [7, 9]]]],
+          ragged_rank=1,
+          expected_shape=(2, None, 2, 2)),
+      dict(
+          pylist=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                  [[[2, 4], [6, 8]], [[1, 5], [7, 9]]]],
+          inner_shape=(2,),
+          expected_shape=(2, None, None, 2)),
+      dict(
+          pylist=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                  [[[2, 4], [6, 8]], [[1, 5], [7, 9]]]],
+          inner_shape=(2, 2),
+          expected_shape=(2, None, 2, 2)),
+      # 4-dimensional tensors with numpy arrays
+      dict(
+          pylist=np.array([[[np.array([1, 2]), [3, 4]], [[5, 6], [7, 8]]],
+                           np.array([[[2, 4], [6, 8]], [[1, 5], [7, 9]]])]),
+          expected_shape=(2, None, None, None)),
+
+      #=========================================================================
+      # Empty tensors (no scalar values) w/ default ragged_rank and inner_shape
+      dict(pylist=[], expected_shape=(0,)),
+      dict(pylist=[[], [], np.array([])], expected_shape=(3, None)),
+      dict(
+          pylist=[[[], []], [], [[], [[]]]],
+          expected_shape=(3, None, None, None)),
+      dict(
+          pylist=np.array([np.array([[], []]),
+                           np.array([]), [[], [[]]]], dtype=object),
+          expected_shape=(3, None, None, None)),
+
+      #=========================================================================
+      # Empty tensors (no scalar values) w/ explicit ragged_rank or inner_shape
+      dict(pylist=[], ragged_rank=1, expected_shape=(0, None)),
+      dict(pylist=[], ragged_rank=2, expected_shape=(0, None, None)),
+      dict(pylist=[], inner_shape=(0, 100, 20), expected_shape=(0, 100, 20)),
+      dict(
+          pylist=[],
+          ragged_rank=1,
+          inner_shape=(100, 20),
+          expected_shape=(0, None, 100, 20)),
+      dict(
+          pylist=[],
+          ragged_rank=2,
+          inner_shape=(100, 20),
+          expected_shape=(0, None, None, 100, 20)),
+      dict(pylist=[[], [], []], ragged_rank=2, expected_shape=(3, None, None)),
+      dict(pylist=[], inner_shape=(0,), expected_shape=(0,)),
+      dict(pylist=[[]], inner_shape=(1, 0), expected_shape=(1, 0)),
+      dict(
+          pylist=np.array([]),
+          ragged_rank=1,
+          inner_shape=(100, 20),
+          expected_shape=(0, None, 100, 20)),
+
+      #=========================================================================
+      # default/inferred dtypes
+      dict(pylist=[], expected_dtype=dtypes.float32),
+      dict(pylist=[[[], [[[]], []]]], expected_dtype=dtypes.float32),
+      dict(pylist=[[1, 2], [3], [4, 5, 6]], expected_dtype=dtypes.int32),
+      dict(pylist=[[1., 2.], [], [4., 5., 6.]], expected_dtype=dtypes.float32),
+      dict(pylist=[[1, 2], [3.], [4, 5, 6]], expected_dtype=dtypes.float32),
+      dict(pylist=[[b'a', b'b'], [b'c']], expected_dtype=dtypes.string),
+      dict(pylist=[[True]], expected_dtype=dtypes.bool),
+      dict(
+          pylist=[np.array([1, 2]), np.array([3.]), [4, 5, 6]],
+          expected_dtype=dtypes.float32),
+
+      #=========================================================================
+      # explicit dtypes
+      dict(pylist=[], dtype=dtypes.float32),
+      dict(pylist=[], dtype=dtypes.string),
+      dict(pylist=[[1, 2], [3], [4, 5, 6]], dtype=dtypes.int64),
+      dict(pylist=[[1, 2], [3], [4, 5, 6]], dtype=dtypes.int32),
+      dict(pylist=[[1, 2], [3], [4, 5, 6]], dtype=dtypes.float32),
+      dict(pylist=[[1., 2.], [3.], [4., 5., 6.]], dtype=dtypes.float16),
+      dict(pylist=[[1., 2.], [3.], [4., 5., 6.]], dtype=dtypes.float32),
+      dict(
+          pylist=[[b'a', b'b'], [b'c'], [b'd', b'e', b'f']],
+          dtype=dtypes.string),
+  )
+  def testRaggedConst(self,
+                      pylist,
+                      dtype=None,
+                      ragged_rank=None,
+                      inner_shape=None,
+                      expected_shape=None,
+                      expected_dtype=None):
+    """Tests that `ragged_const(pylist).eval().tolist() == pylist`.
+
+    Args:
+      pylist: The `pylist` argument for `ragged_const()`.
+      dtype: The `dtype` argument for `ragged_const()`.  If not None, then also
+        test that the resulting ragged tensor has this `dtype`.
+      ragged_rank: The `ragged_rank` argument for `ragged_const()`.  If not
+        None, then also test that the resulting ragged tensor has this
+        `ragged_rank`.
+      inner_shape: The `inner_shape` argument for `ragged_const()`.  If not
+        None, then also test that the resulting ragged tensor has this
+        `inner_shape`.
+      expected_shape: The expected shape for the resulting ragged tensor.
+      expected_dtype: The expected dtype for the resulting ragged tensor (used
+        to test default/inferred types when dtype=None).
+    """
+    rt = ragged_factory_ops.constant(
+        pylist, dtype=dtype, ragged_rank=ragged_rank, inner_shape=inner_shape)
+    # Normalize the pylist, i.e., convert all np.arrays to list.
+    # E.g., [np.array((1,2))] --> [[1,2]]
+    pylist = _normalize_pylist(pylist)
+
+    # If dtype was explicitly specified, check it.
+    if dtype is not None:
+      self.assertEqual(rt.dtype, dtype)
+    if expected_dtype is not None:
+      self.assertEqual(rt.dtype, expected_dtype)
+
+    # If ragged_rank was explicitly specified, check it.
+    if ragged_rank is not None:
+      if isinstance(rt, ragged_tensor.RaggedTensor):
+        self.assertEqual(rt.ragged_rank, ragged_rank)
+      else:
+        self.assertEqual(0, ragged_rank)
+
+    # If inner_shape was explicitly specified, check it.
+    if inner_shape is not None:
+      if isinstance(rt, ragged_tensor.RaggedTensor):
+        self.assertEqual(rt.flat_values.shape.as_list()[1:], list(inner_shape))
+      else:
+        self.assertEqual(rt.shape.as_list(), list(inner_shape))
+
+    if expected_shape is not None:
+      self.assertEqual(tuple(rt.shape.as_list()), expected_shape)
+      if (expected_shape and expected_shape[0] == 0 and
+          None not in expected_shape):
+        pylist = np.zeros(expected_shape, rt.dtype.as_numpy_dtype)
+
+    self.assertAllEqual(rt, pylist)
+
+  @parameterized.parameters(
+      dict(
+          pylist=12,
+          ragged_rank=1,
+          exception=ValueError,
+          message='Invalid pylist=12: incompatible with ragged_rank=1'),
+      dict(
+          pylist=12,
+          inner_shape=(1,),
+          exception=ValueError,
+          message='Invalid pylist=12: incompatible with '
+          'dim\\(inner_shape\\)=1'),
+      dict(
+          pylist=[[[1], [2]]],
+          ragged_rank=-1,
+          exception=ValueError,
+          message='Invalid ragged_rank=-1: must be nonnegative'),
+      dict(
+          pylist=[[1, [2]]],
+          exception=ValueError,
+          message='all scalar values must have the same nesting depth'),
+      dict(
+          pylist=[[[1]], [[[2]]]],
+          exception=ValueError,
+          message='all scalar values must have the same nesting depth'),
+      dict(
+          pylist=[[1], [[]]],
+          exception=ValueError,
+          message='Invalid pylist=.*: empty list nesting is greater '
+          'than scalar value nesting'),
+      dict(
+          pylist=[1, 2, 3],
+          ragged_rank=1,
+          exception=ValueError,
+          message='pylist has scalar values depth 1, but ragged_rank=1 '
+          'requires scalar value depth greater than 1'),
+      dict(
+          pylist=[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+          ragged_rank=2,
+          exception=ValueError,
+          message='pylist has scalar values depth 2, but ragged_rank=2 '
+          'requires scalar value depth greater than 2'),
+      dict(pylist=[1, 2, 3], inner_shape=(1, 1), exception=TypeError),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+          inner_shape=(2, 2),
+          ragged_rank=1,
+          exception=ValueError,
+          message='Invalid pylist=.*: incompatible with ragged_rank=1 and '
+          'dim\\(inner_shape\\)=2'),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [[5, 6], [7, 8, 9]]],
+          ragged_rank=1,
+          exception=ValueError,
+          message='inner values have inconsistent shape'),
+      dict(
+          pylist=[[[], [[]]]],
+          ragged_rank=1,
+          exception=ValueError,
+          message='inner values have inconsistent shape'),
+  )
+  def testRaggedConstError(self,
+                           pylist,
+                           dtype=None,
+                           ragged_rank=None,
+                           inner_shape=None,
+                           exception=None,
+                           message=None):
+    """Tests that `ragged_const()` raises an expected exception."""
+    self.assertRaisesRegex(
+        exception,
+        message,
+        ragged_factory_ops.constant,
+        pylist,
+        dtype=dtype,
+        ragged_rank=ragged_rank,
+        inner_shape=inner_shape)
+
+  @parameterized.parameters([
+      dict(pylist=9, scalar_depth=0, max_depth=0),
+      dict(pylist=[9], scalar_depth=1, max_depth=1),
+      dict(pylist=[1, 2, 3], scalar_depth=1, max_depth=1),
+      dict(pylist=[[1], [2]], scalar_depth=2, max_depth=2),
+      dict(pylist=[[[1], [2]], [[3]]], scalar_depth=3, max_depth=3),
+      dict(pylist=[], scalar_depth=None, max_depth=1),
+      dict(pylist=[[]], scalar_depth=None, max_depth=2),
+      dict(pylist=[[], [], []], scalar_depth=None, max_depth=2),
+      dict(pylist=[[[], []], [[], [[[]]]], []], scalar_depth=None, max_depth=5),
+      dict(
+          pylist=[1, [2]],
+          exception=ValueError,
+          message='all scalar values must have the same nesting depth'),
+      dict(
+          pylist=[[1], 2],
+          exception=ValueError,
+          message='all scalar values must have the same nesting depth'),
+      dict(
+          pylist=[[[[1]], []], [[2]]],
+          exception=ValueError,
+          message='all scalar values must have the same nesting depth'),
+  ])
+  def testScalarAndMaxDepthHelper(self,
+                                  pylist,
+                                  scalar_depth=None,
+                                  max_depth=None,
+                                  exception=None,
+                                  message=None):
+    """Tests for the _find_scalar_and_max_depth helper function."""
+    if exception is not None:
+      self.assertRaisesRegex(exception, message,
+                             ragged_factory_ops._find_scalar_and_max_depth,
+                             pylist)
+    else:
+      self.assertEqual(
+          ragged_factory_ops._find_scalar_and_max_depth(pylist),
+          (scalar_depth, max_depth))
+
+  @parameterized.parameters([
+      dict(pylist=[[1], [2, 3]], ragged_rank=1, inner_shape=()),
+      dict(
+          pylist=[[[1], [2]], [[3], [4], [5]]], ragged_rank=1,
+          inner_shape=(1,)),
+      dict(pylist=[[[1], [2]], [[3], [4], [5]]], ragged_rank=2, inner_shape=()),
+      dict(
+          pylist=[[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [2, 4, 6]]]],
+          ragged_rank=1,
+          inner_shape=(2, 3)),
+      dict(
+          pylist=[[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [2, 4, 6]]]],
+          ragged_rank=2,
+          inner_shape=(3,)),
+      dict(
+          pylist=[[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [2, 4, 6]]]],
+          ragged_rank=3,
+          inner_shape=()),
+      dict(
+          pylist=[[[1], [2, 3]]],
+          ragged_rank=1,
+          exception=ValueError,
+          message='inner values have inconsistent shape'),
+      dict(
+          pylist=[[[1], [[2]]]],
+          ragged_rank=1,
+          exception=ValueError,
+          message='inner values have inconsistent shape'),
+      dict(
+          pylist=[[[[1]], [2]]],
+          ragged_rank=1,
+          exception=ValueError,
+          message='inner values have inconsistent shape'),
+  ])
+  def testDefaultInnerShapeForPylistHelper(self,
+                                           pylist,
+                                           ragged_rank,
+                                           inner_shape=None,
+                                           exception=None,
+                                           message=None):
+    """Tests for the _default_inner_shape_for_pylist helper function."""
+    if exception is not None:
+      self.assertRaisesRegex(
+          exception, message,
+          ragged.ragged_factory_ops._default_inner_shape_for_pylist, pylist,
+          ragged_rank)
+    else:
+      self.assertEqual(
+          ragged.ragged_factory_ops._default_inner_shape_for_pylist(
+              pylist, ragged_rank), inner_shape)
+
+
+def _normalize_pylist(item):
+  """Convert all (possibly nested) np.arrays contained in item to list."""
+  # convert np.arrays in current level to list
+  if not isinstance(item, (list, np.ndarray)):
+    return item
+  level = (x.tolist() if isinstance(x, np.ndarray) else x for x in item)
+  return [_normalize_pylist(el) if isinstance(item, (list, np.ndarray))
+          else el for el in level]
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_constant_value_op_test.py

@@ -0,0 +1,331 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_factory_ops.constant_value."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedConstantValueOpTest(test_util.TensorFlowTestCase,
+                                parameterized.TestCase):
+  @parameterized.parameters(
+      #=========================================================================
+      # 0-dimensional tensors.
+      dict(pylist='x', expected_shape=()),
+
+      #=========================================================================
+      # 1-dimensional tensors.
+      dict(pylist=[1, 2, 3], expected_shape=(3,)),
+
+      #=========================================================================
+      # 2-dimensional tensors.
+      dict(pylist=[[1, 2, 3], [4], [5, 6]], expected_shape=(3, None)),
+      dict(pylist=[[1, 2, 3], [4, 5, 6], [7, 8, 9]], expected_shape=(3, None)),
+
+      #=========================================================================
+      # 3-dimensional tensors.
+      dict(
+          pylist=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          expected_shape=(3, None, None)),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          inner_shape=(2,),
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          inner_shape=(2,),
+          expected_shape=(3, None, 2)),
+      # 3-dimensional tensors with numpy arrays
+      dict(
+          pylist=[[[1, 2], np.array([3, np.array(4)])],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          expected_shape=(3, None, None)),
+      dict(
+          pylist=[[[1, 2], np.array([3, np.array(4)])],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[np.array([3, np.array(4)]), [1, 2]],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[[1, 2], np.array([3, np.array(4)])],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          inner_shape=(2,),
+          expected_shape=(3, None, 2)),
+      dict(
+          pylist=[[[1, 2], np.array([3, np.array(4)])],
+                  np.array([]), [[5, 6], [7, 8], [9, 0]]],
+          ragged_rank=1,
+          inner_shape=(2,),
+          expected_shape=(3, None, 2)),
+      #=========================================================================
+      # 4-dimensional tensors.
+      dict(
+          pylist=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                  [[[2, 4], [6, 8]], [[1, 5], [7, 9]]]],
+          expected_shape=(2, None, None, None)),
+      dict(
+          pylist=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                  [[[2, 4], [6, 8]], [[1, 5], [7, 9]]]],
+          ragged_rank=1,
+          expected_shape=(2, None, 2, 2)),
+      dict(
+          pylist=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                  [[[2, 4], [6, 8]], [[1, 5], [7, 9]]]],
+          inner_shape=(2,),
+          expected_shape=(2, None, None, 2)),
+      dict(
+          pylist=[[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                  [[[2, 4], [6, 8]], [[1, 5], [7, 9]]]],
+          inner_shape=(2, 2),
+          expected_shape=(2, None, 2, 2)),
+      # 4-dimensional tensors with numpy arrays
+      dict(
+          pylist=np.array([[[np.array([1, 2]), [3, 4]], [[5, 6], [7, 8]]],
+                           np.array([[[2, 4], [6, 8]], [[1, 5], [7, 9]]])]),
+          expected_shape=(2, None, None, None)),
+
+      #=========================================================================
+      # Empty tensors (no scalar values) w/ default ragged_rank and inner_shape
+      dict(pylist=[], expected_shape=(0,)),
+      dict(pylist=[[], [], np.array([])], expected_shape=(3, None)),
+      dict(
+          pylist=[[[], []], [], [[], [[]]]],
+          expected_shape=(3, None, None, None)),
+      dict(
+          pylist=np.array([np.array([[], []]),
+                           np.array([]), [[], [[]]]], dtype=object),
+          expected_shape=(3, None, None, None)),
+
+      #=========================================================================
+      # Empty tensors (no scalar values) w/ explicit ragged_rank or inner_shape
+      dict(pylist=[], ragged_rank=1, expected_shape=(0, None)),
+      dict(pylist=[], ragged_rank=2, expected_shape=(0, None, None)),
+      dict(pylist=[], inner_shape=(0, 100, 20), expected_shape=(0, 100, 20)),
+      dict(
+          pylist=[],
+          ragged_rank=1,
+          inner_shape=(100, 20),
+          expected_shape=(0, None, 100, 20)),
+      dict(
+          pylist=[],
+          ragged_rank=2,
+          inner_shape=(100, 20),
+          expected_shape=(0, None, None, 100, 20)),
+      dict(pylist=[[], [], []], ragged_rank=2, expected_shape=(3, None, None)),
+      dict(pylist=[], inner_shape=(0,), expected_shape=(0,)),
+      dict(pylist=[[]], inner_shape=(1, 0), expected_shape=(1, 0)),
+      dict(
+          pylist=np.array([]),
+          ragged_rank=1,
+          inner_shape=(100, 20),
+          expected_shape=(0, None, 100, 20)),
+
+      #=========================================================================
+      # default/inferred dtypes.
+      #
+      # Note: numpy has different default/inferred types than tensorflow.
+      # Since we are using values, not tensors, we get the default numpy types
+      # here.
+      dict(pylist=[], expected_dtype=np.float64),
+      dict(pylist=[[[], [[[]], []]]], expected_dtype=np.float64),
+      dict(pylist=[[1, 2], [3], [4, 5, 6]], expected_dtype=np.int64),
+      dict(pylist=[[1., 2.], [], [4., 5., 6.]], expected_dtype=np.float64),
+      dict(pylist=[[1, 2], [3.], [4, 5, 6]], expected_dtype=np.float64),
+      dict(pylist=[[b'a', b'b'], [b'c']], expected_dtype=np.dtype('S1')),
+      dict(pylist=[[True]], expected_dtype=np.bool_),
+      dict(
+          pylist=[np.array([1, 2]), np.array([3.]), [4, 5, 6]],
+          expected_dtype=np.float64),
+
+      #=========================================================================
+      # explicit dtypes
+      dict(pylist=[], dtype=np.float32),
+      dict(pylist=[], dtype=np.dtype('S1')),
+      dict(pylist=[[1, 2], [3], [4, 5, 6]], dtype=np.int64),
+      dict(pylist=[[1, 2], [3], [4, 5, 6]], dtype=np.int32),
+      dict(pylist=[[1, 2], [3], [4, 5, 6]], dtype=np.float32),
+      dict(pylist=[[1., 2.], [3.], [4., 5., 6.]], dtype=np.float16),
+      dict(pylist=[[1., 2.], [3.], [4., 5., 6.]], dtype=np.float32),
+      dict(
+          pylist=[[b'a', b'b'], [b'c'], [b'd', b'e', b'f']],
+          dtype=np.dtype('S1')),
+      dict(pylist=[], dtype=dtypes.float32, expected_dtype=np.float32),
+      dict(pylist=[], dtype=dtypes.int32, expected_dtype=np.int32),
+  )
+  def testRaggedValues(self,
+                       pylist,
+                       dtype=None,
+                       ragged_rank=None,
+                       inner_shape=None,
+                       expected_shape=None,
+                       expected_dtype=None):
+    """Tests that `ragged_value(pylist).to_list() == pylist`."""
+    rt = ragged_factory_ops.constant_value(
+        pylist, dtype=dtype, ragged_rank=ragged_rank, inner_shape=inner_shape)
+    # Normalize the pylist, i.e., convert all np.arrays to list.
+    # E.g., [np.array((1,2))] --> [[1,2]]
+    pylist = _normalize_pylist(pylist)
+    # If dtype was explicitly specified, check it.
+    if expected_dtype is not None:
+      self.assertEqual(rt.dtype, expected_dtype)
+    elif dtype is not None:
+      self.assertEqual(rt.dtype, dtype)
+
+    # If ragged_rank was explicitly specified, check it.
+    if ragged_rank is not None:
+      if isinstance(rt, ragged_tensor_value.RaggedTensorValue):
+        self.assertEqual(rt.ragged_rank, ragged_rank)
+      else:
+        self.assertEqual(0, ragged_rank)
+
+    # If inner_shape was explicitly specified, check it.
+    if inner_shape is not None:
+      if isinstance(rt, ragged_tensor_value.RaggedTensorValue):
+        self.assertEqual(rt.flat_values.shape[1:], inner_shape)
+      else:
+        self.assertEqual(rt.shape, inner_shape)
+
+    if expected_shape is not None:
+      self.assertEqual(tuple(rt.shape), expected_shape)
+
+    if rt.shape:
+      if isinstance(rt, ragged_tensor_value.RaggedTensorValue):
+        self.assertEqual(rt.to_list(), pylist)
+      else:
+        self.assertEqual(rt.tolist(), pylist)
+      if expected_shape is not None:
+        self.assertEqual(rt.shape, expected_shape)
+    else:
+      self.assertEqual(rt, pylist)
+      if expected_shape is not None:
+        self.assertEqual((), expected_shape)
+
+  @parameterized.parameters(
+      dict(
+          pylist=12,
+          ragged_rank=1,
+          exception=ValueError,
+          message='Invalid pylist=12: incompatible with ragged_rank=1'),
+      dict(
+          pylist=np.array(12),
+          ragged_rank=1,
+          exception=ValueError,
+          message='Invalid pylist=array\\(12\\): incompatible with '
+          'ragged_rank=1'),
+      dict(
+          pylist=12,
+          inner_shape=(1,),
+          exception=ValueError,
+          message='Invalid pylist=12: incompatible with '
+          'dim\\(inner_shape\\)=1'),
+      dict(
+          pylist=[[[1], [2]]],
+          ragged_rank=-1,
+          exception=ValueError,
+          message='Invalid ragged_rank=-1: must be nonnegative'),
+      dict(
+          pylist=[[1, [2]]],
+          exception=ValueError,
+          message='all scalar values must have the same nesting depth'),
+      dict(
+          pylist=[[[1]], [[[2]]]],
+          exception=ValueError,
+          message='all scalar values must have the same nesting depth'),
+      dict(
+          pylist=[[1], [[]]],
+          exception=ValueError,
+          message='Invalid pylist=.*: empty list nesting is greater '
+          'than scalar value nesting'),
+      dict(
+          pylist=[1, 2, 3],
+          ragged_rank=1,
+          exception=ValueError,
+          message='pylist has scalar values depth 1, but ragged_rank=1 '
+          'requires scalar value depth greater than 1'),
+      dict(
+          pylist=[[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+          ragged_rank=2,
+          exception=ValueError,
+          message='pylist has scalar values depth 2, but ragged_rank=2 '
+          'requires scalar value depth greater than 2'),
+      dict(
+          pylist=[1, 2, 3],
+          inner_shape=(1, 1),
+          exception=ValueError,
+          message='cannot reshape array'),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+          inner_shape=(2, 2),
+          ragged_rank=1,
+          exception=ValueError,
+          message='Invalid pylist=.*: incompatible with ragged_rank=1 and '
+          'dim\\(inner_shape\\)=2'),
+      dict(
+          pylist=[[[1, 2], [3, 4]], [[5, 6], [7, 8, 9]]],
+          ragged_rank=1,
+          exception=ValueError,
+          message='inner values have inconsistent shape'),
+      dict(
+          pylist=[[[], [[]]]],
+          ragged_rank=1,
+          exception=ValueError,
+          message='inner values have inconsistent shape'),
+  )
+  def testRaggedValuesError(self,
+                            pylist,
+                            dtype=None,
+                            ragged_rank=None,
+                            inner_shape=None,
+                            exception=None,
+                            message=None):
+    """Tests that `constant_value()` raises an expected exception."""
+    self.assertRaisesRegex(
+        exception,
+        message,
+        ragged_factory_ops.constant_value,
+        pylist,
+        dtype=dtype,
+        ragged_rank=ragged_rank,
+        inner_shape=inner_shape)
+
+
+def _normalize_pylist(item):
+  """Convert all (possibly nested) np.arrays contained in item to list."""
+  # convert np.arrays in current level to list
+  if not isinstance(item, (list, np.ndarray)):
+    return item
+  level = (x.tolist() if isinstance(x, np.ndarray) else x for x in item)
+  return [_normalize_pylist(el) if isinstance(item, (list, np.ndarray))
+          else el for el in level]
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_conversion_ops.py

@@ -0,0 +1,180 @@
+# 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.
+# ==============================================================================
+"""Ops to convert between RaggedTensors and other tensor types."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_ragged_conversion_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+
+
+def from_tensor(tensor,
+                lengths=None,
+                padding=None,
+                ragged_rank=1,
+                row_splits_dtype=dtypes.int64,
+                name=None):
+  if ragged_tensor.is_ragged(tensor):
+    return tensor
+  else:
+    return ragged_tensor.RaggedTensor.from_tensor(
+        tensor,
+        lengths=lengths,
+        padding=padding,
+        ragged_rank=ragged_rank,
+        row_splits_dtype=row_splits_dtype,
+        name=name)
+
+
+def to_tensor(rt_input, default_value=None, name=None):
+  if ragged_tensor.is_ragged(rt_input):
+    return rt_input.to_tensor(default_value, name)
+  else:
+    return rt_input
+
+
+def ragged_to_dense(rt_input, default_value=None, shape=None):
+  """Create a dense tensor from a ragged tensor."""
+  return rt_input.to_tensor(default_value=default_value, shape=shape)
+
+
+@ops.RegisterGradient("RaggedTensorToTensor")
+def _ragged_tensor_to_tensor_grad(op, grad):
+  """Gradient for RaggedToTensor op."""
+  # Extract inputs from the op.
+  flat_values = op.inputs[1]
+  default_value = op.inputs[2]
+  row_partition_tensors = op.inputs[3:]
+  row_partition_types = op.get_attr("row_partition_types")
+  flat_value_shape = array_ops.shape(flat_values)
+  ragged_rank = sum(
+      1 for typ in row_partition_types if typ != b"FIRST_DIM_SIZE")
+
+  # Create two tensors that correspond 1:1 with grad (and op.output):
+  # * indices[i1...iN] is the index in `flat_values` of the value used to
+  #   populate output[i1...iN] (if the value came from `flat_values`) or
+  #   -1 (if the value came from `default_value`).
+  # * mask[i1...iN] is true if output[i1...iN] came from `flat_values`, or
+  #   false if it came from `default_value`.
+  indices = gen_ragged_conversion_ops.ragged_tensor_to_tensor(
+      shape=array_ops.shape(grad)[:1 + ragged_rank],
+      values=math_ops.range(flat_value_shape[0]),
+      default_value=-1,
+      row_partition_types=row_partition_types,
+      row_partition_tensors=row_partition_tensors)
+  mask = math_ops.not_equal(indices, -1)
+
+  # Select out the gradients & indices that came from `flat_values`, and use
+  # those to construct the gradient for `flat_values` (as an IndexedSlices).
+  values_grad = indexed_slices.IndexedSlices(
+      values=array_ops.boolean_mask(grad, mask),
+      indices=array_ops.boolean_mask(indices, mask),
+      dense_shape=flat_value_shape)
+
+  # Select out the gradients that came from `default_value`, and sum them to
+  # get the gradient for the default.  Note that the default_value may have
+  # been broadcast as part of the RaggedTensorToTensor operation, so we also
+  # need to reduce any dimensions that might have been broadcast.
+  default_grads = array_ops.boolean_mask(grad, ~mask)
+  dims_to_reduce = math_ops.range(
+      array_ops.rank(default_grads) -
+      _rank_ignoring_leading_dims_with_size_1(default_value))
+  default_grad = math_ops.reduce_sum(default_grads, axis=dims_to_reduce)
+
+  # Restore any leading dims with size one.
+  default_grad = array_ops.reshape(default_grad, array_ops.shape(default_value))
+
+  return ([None, values_grad, default_grad] +
+          [None for _ in row_partition_tensors])
+
+
+def _rank_ignoring_leading_dims_with_size_1(value):
+  """Returns `rank(value)`, ignoring any leading dimensions with size 1."""
+  # Compute the result using static shape, if possible.
+  if value.shape.rank is not None:
+    ndims = value.shape.rank
+    for dim in value.shape.dims:
+      if dim.value == 1:
+        ndims -= 1
+      elif dim.value is None:
+        ndims = None  # Can't compute the result using static shape.
+        break
+      else:
+        break
+    if ndims is not None:
+      return ndims
+
+  # Otherwise, we need to compute the result dynamically.  The math we use to
+  # do this is a bit round-about, so here's an example to illustrate:
+  #              shape = [1, 1, 3, 5, 1, 4]  # shape(value)
+  #         dim_is_one = [1, 1, 0, 0, 1, 0]  # equal(shape, 1)
+  #       leading_ones = [1, 1, 0, 0, 0, 0]  # cumprod(dim_is_one)
+  #   num_leading_ones = 2                   # reduce_sum(leading_ones)
+  #             result = 4                   # rank(value) - num_leading_ones
+  shape = array_ops.shape(value)
+  dim_is_one = math_ops.cast(math_ops.equal(shape, 1), dtypes.int32)
+  leading_ones = math_ops.cumprod(dim_is_one)
+  num_leading_ones = math_ops.reduce_sum(leading_ones)
+  return array_ops.rank(value) - num_leading_ones
+
+
+def to_sparse(rt_input, name=None):
+  return rt_input.to_sparse(name)
+
+
+def from_sparse(st_input, name=None):
+  return ragged_tensor.RaggedTensor.from_sparse(st_input, name)
+
+
+@ops.RegisterGradient("RaggedTensorFromVariant")
+def _ragged_tensor_from_variant_grad(op, *grads):
+  """Gradient for RaggedTensorFromVariant op."""
+
+  variant_rank = op.inputs[0].shape.rank
+  if variant_rank == 0:
+    batched_input = False
+  elif variant_rank == 1:
+    batched_input = True
+  elif variant_rank is None:
+    batched_input = (op.get_attr("output_ragged_rank") > 0)
+  else:
+    # TODO(edloper): Add a batch_dims argument to RaggedTensorToVariant, so
+    # we can support this.
+    raise ValueError("Unable to compute gradient: RaggedTensorToVariant "
+                     "can currently only generate 0D or 1D output.")
+  return [
+      gen_ragged_conversion_ops.ragged_tensor_to_variant(
+          rt_nested_splits=op.outputs[:-1],
+          rt_dense_values=grads[-1],
+          batched_input=batched_input)
+  ]
+
+
+@ops.RegisterGradient("RaggedTensorToVariant")
+def _ragged_tensor_to_variant_grad(op, encoded_ragged_grad):
+  """Gradient for RaggedTensorToVariant op."""
+  dense_values = op.inputs[-1]
+  ragged_rank = len(op.inputs) - 1
+  row_splits = 0 if ragged_rank == 0 else op.inputs[0]
+  values_grad = gen_ragged_conversion_ops.ragged_tensor_to_variant_gradient(
+      encoded_ragged_grad=encoded_ragged_grad,
+      row_splits=row_splits,
+      dense_values_shape=array_ops.shape(dense_values),
+      Tvalues=op.inputs[-1].dtype)
+  result = [None] * ragged_rank + [values_grad]
+  return result

tensorflow/python/ops/ragged/ragged_cross_op_test.py

@@ -0,0 +1,511 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Tests for tf.ragged.cross and tf.ragged.cross_hashed."""
+
+from absl.testing import parameterized
+
+import numpy as np
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_ragged_array_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+ragged_const = ragged_factory_ops.constant_value
+dense_const = np.array
+
+
+def sparse_const(matrix):
+  indices = []
+  values = []
+  for i, row in enumerate(matrix):
+    for j, val in enumerate(row):
+      indices.append([i, j])
+      values.append(val)
+  shape = [len(matrix), max(len(row) for row in matrix)] if matrix else [0, 0]
+  if not values:
+    indices = np.zeros([0, 2], dtype=np.int64)
+    values = np.zeros([0], dtype=np.int64)
+  return sparse_tensor.SparseTensorValue(indices, values, shape)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedCrossOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters([
+      dict(
+          testcase_name='NoInputs',
+          inputs=[],
+          expected=ragged_const([], ragged_rank=1, dtype=dtypes.int32)),
+      dict(
+          testcase_name='OneInput_RaggedStr',
+          inputs=[ragged_const([['a', 'b'], [], ['c']])],
+          expected=ragged_const([[b'a', b'b'], [], [b'c']])),
+      dict(
+          testcase_name='OneInput_RaggedInt',
+          inputs=[ragged_const([[1, 2, 3], [4, 5]])],
+          expected=ragged_const([[b'1', b'2', b'3'], [b'4', b'5']])),
+      dict(
+          testcase_name='OneInput_DenseInt',
+          inputs=[dense_const([[1, 2, 3], [4, 5, 6]])],
+          expected=ragged_const([[b'1', b'2', b'3'], [b'4', b'5', b'6']])),
+      dict(
+          testcase_name='OneInput_SparseStr',
+          inputs=[sparse_const([['a', 'b'], [], ['c']])],
+          expected=ragged_const([[b'a', b'b'], [], [b'c']])),
+      dict(
+          testcase_name='TwoInputs_RaggedStr_RaggedStr',
+          inputs=[
+              ragged_const([['a', 'b'], [], ['c']]),
+              ragged_const([['d', 'e'], ['f'], ['g']])
+          ],
+          expected=ragged_const([[b'a_X_d', b'a_X_e', b'b_X_d', b'b_X_e'], [],
+                                 [b'c_X_g']])),
+      dict(
+          testcase_name='TwoInputs_RaggedInt_RaggedInt',
+          inputs=[
+              ragged_const([[1, 2], [], [3]]),
+              ragged_const([[4, 5, 6], [], [7]])
+          ],
+          expected=ragged_const(
+              [[b'1_X_4', b'1_X_5', b'1_X_6', b'2_X_4', b'2_X_5', b'2_X_6'], [],
+               [b'3_X_7']])),
+      dict(
+          testcase_name='TwoInputs_RaggedStr_RaggedInt',
+          inputs=[
+              ragged_const([['a', 'b'], [], ['c']]),
+              ragged_const([['1', '2'], ['3'], ['4']])
+          ],
+          expected=ragged_const([[b'a_X_1', b'a_X_2', b'b_X_1', b'b_X_2'], [],
+                                 [b'c_X_4']])),
+      dict(
+          testcase_name='TwoInputs_SparseStr_SparseStr',
+          inputs=[
+              sparse_const([['a', 'b'], [], ['c']]),
+              sparse_const([['d', 'e'], ['f'], ['g']])
+          ],
+          expected=ragged_const([[b'a_X_d', b'a_X_e', b'b_X_d', b'b_X_e'], [],
+                                 [b'c_X_g']])),
+      dict(
+          testcase_name='TwoInputs_DenseInt_DenseInt',
+          inputs=[dense_const([[1, 2], [3, 4]]),
+                  dense_const([[5, 6], [7, 8]])],
+          expected=ragged_const([[b'1_X_5', b'1_X_6', b'2_X_5', b'2_X_6'],
+                                 [b'3_X_7', b'3_X_8', b'4_X_7', b'4_X_8']])),
+      dict(
+          testcase_name='TwoInputs_DenseInt_DenseStr',
+          inputs=[
+              dense_const([[1, 2], [3, 4]]),
+              dense_const([[b'5', b'6'], [b'7', b'8']])
+          ],
+          expected=ragged_const([[b'1_X_5', b'1_X_6', b'2_X_5', b'2_X_6'],
+                                 [b'3_X_7', b'3_X_8', b'4_X_7', b'4_X_8']])),
+      dict(
+          testcase_name='TwoInputs_RaggedInt_DenseInt',
+          inputs=[
+              ragged_const([[], [], [1, 2], [3]]),
+              dense_const([[1, 2], [3, 4], [5, 6], [7, 8]])
+          ],
+          expected=ragged_const([[], [],
+                                 [b'1_X_5', b'1_X_6', b'2_X_5', b'2_X_6'],
+                                 [b'3_X_7', b'3_X_8']])),
+      dict(
+          # This test exercises `input_order`.
+          testcase_name='TwoInputs_DenseInt_RaggedStr',
+          inputs=[
+              dense_const([[1, 2], [3, 4], [5, 6]]),
+              ragged_const([['d', 'e'], ['f'], ['g']])
+          ],
+          expected=ragged_const([[b'1_X_d', b'1_X_e', b'2_X_d', b'2_X_e'],
+                                 [b'3_X_f', b'4_X_f'], [b'5_X_g', b'6_X_g']]),
+          matches_sparse_cross=False  # sparse doesn't preserve input order.
+      ),
+      dict(
+          # This test exercises `input_order`.
+          testcase_name='TwoInputs_SparseInt_RaggedStr',
+          inputs=[
+              sparse_const([[1, 2], [3, 4], [5, 6]]),
+              ragged_const([['d', 'e'], ['f'], ['g']])
+          ],
+          expected=ragged_const([[b'1_X_d', b'1_X_e', b'2_X_d', b'2_X_e'],
+                                 [b'3_X_f', b'4_X_f'], [b'5_X_g', b'6_X_g']]),
+          matches_sparse_cross=False  # sparse doesn't preserve input order.
+      ),
+      dict(
+          testcase_name='ThreeInputs_RaggedInt_RaggedInt_RaggedInt',
+          inputs=[
+              ragged_const([[11], [12, 13], [], [14, 15]]),
+              ragged_const([[21, 22], [23], [24, 25], [26, 27]]),
+              ragged_const([[31], [32, 33], [34, 35], [36, 37]])
+          ],
+          expected=ragged_const([[b'11_X_21_X_31', b'11_X_22_X_31'],
+                                 [
+                                     b'12_X_23_X_32', b'12_X_23_X_33',
+                                     b'13_X_23_X_32', b'13_X_23_X_33'
+                                 ], [],
+                                 [
+                                     b'14_X_26_X_36', b'14_X_26_X_37',
+                                     b'14_X_27_X_36', b'14_X_27_X_37',
+                                     b'15_X_26_X_36', b'15_X_26_X_37',
+                                     b'15_X_27_X_36', b'15_X_27_X_37'
+                                 ]])),
+      dict(
+          testcase_name='ThreeInputs_RaggedInt_SparseInt_DenseInt',
+          inputs=[
+              ragged_const([[11], [12, 13], [], [14, 15]]),
+              sparse_const([[21, 22], [23], [24, 25], [26, 27]]),
+              dense_const([[31], [32], [33], [34]])
+          ],
+          expected=ragged_const([[b'11_X_21_X_31', b'11_X_22_X_31'],
+                                 [
+                                     b'12_X_23_X_32',
+                                     b'13_X_23_X_32',
+                                 ], [],
+                                 [
+                                     b'14_X_26_X_34',
+                                     b'14_X_27_X_34',
+                                     b'15_X_26_X_34',
+                                     b'15_X_27_X_34',
+                                 ]])),
+      dict(
+          testcase_name='FiveInputs',
+          inputs=[
+              ragged_const([[1]]),
+              dense_const([[2]]),
+              ragged_const([[3]]),
+              sparse_const([[4]]),
+              ragged_const([[5]])
+          ],
+          expected=ragged_const([[b'1_X_2_X_3_X_4_X_5']]),
+          matches_sparse_cross=False  # sparse doesn't preserve input order.
+      ),
+      dict(
+          testcase_name='Permutation_3x3x3',
+          inputs=[[['11', '12', '13']], [['21', '22', '23']],
+                  [['31', '32', '33']]],
+          expected=[[
+              b'11_X_21_X_31', b'11_X_21_X_32', b'11_X_21_X_33',
+              b'11_X_22_X_31', b'11_X_22_X_32', b'11_X_22_X_33',
+              b'11_X_23_X_31', b'11_X_23_X_32', b'11_X_23_X_33',
+              b'12_X_21_X_31', b'12_X_21_X_32', b'12_X_21_X_33',
+              b'12_X_22_X_31', b'12_X_22_X_32', b'12_X_22_X_33',
+              b'12_X_23_X_31', b'12_X_23_X_32', b'12_X_23_X_33',
+              b'13_X_21_X_31', b'13_X_21_X_32', b'13_X_21_X_33',
+              b'13_X_22_X_31', b'13_X_22_X_32', b'13_X_22_X_33',
+              b'13_X_23_X_31', b'13_X_23_X_32', b'13_X_23_X_33'
+          ]]),
+      dict(
+          testcase_name='BatchSizeZero',
+          inputs=[
+              ragged_const([], ragged_rank=1, dtype=dtypes.int32),
+              sparse_const([]),
+              np.zeros([0, 3], dtype=np.int32),
+          ],
+          expected=ragged_const([], ragged_rank=1, dtype=dtypes.int32)),
+      dict(
+          testcase_name='ThreeInputs_OneEmpty',
+          inputs=[
+              ragged_const([[1, 2]]),
+              ragged_const([[]], dtype=dtypes.int32),
+              ragged_const([[3, 4]])
+          ],
+          expected=ragged_const([[]], dtype=dtypes.string)),
+      dict(
+          testcase_name='ThreeInputs_AllEmpty',
+          inputs=[
+              ragged_const([[]], dtype=dtypes.int64),
+              ragged_const([[]], dtype=dtypes.string),
+              ragged_const([[]], dtype=dtypes.int32)
+          ],
+          expected=ragged_const([[]], ragged_rank=1, dtype=dtypes.string)),
+      dict(
+          testcase_name='HashedZeroBucketsDefaultKey',
+          inputs=[
+              ragged_const([['batch1-FC1-F1']]),
+              ragged_const([['batch1-FC2-F1']]),
+              ragged_const([['batch1-FC3-F1']])
+          ],
+          expected_hashed=ragged_const([[1971693436396284976]])),
+      dict(
+          testcase_name='Hashed100BucketsDefaultKey',
+          inputs=[
+              ragged_const([['batch1-FC1-F1']]),
+              ragged_const([['batch1-FC2-F1']]),
+              ragged_const([['batch1-FC3-F1']])
+          ],
+          num_buckets=100,
+          expected_hashed=ragged_const([[83]])),
+      dict(
+          testcase_name='HashedZeroBucketsCustomKey',
+          inputs=[
+              ragged_const([['batch1-FC1-F1']]),
+              ragged_const([['batch1-FC2-F1']]),
+              ragged_const([['batch1-FC3-F1']])
+          ],
+          hash_key=ragged_array_ops._DEFAULT_CROSS_HASH_KEY + 1,
+          expected_hashed=ragged_const([[4847552627144134031]])),
+      dict(
+          testcase_name='Hashed100BucketsCustomKey',
+          inputs=[
+              ragged_const([['batch1-FC1-F1']]),
+              ragged_const([['batch1-FC2-F1']]),
+              ragged_const([['batch1-FC3-F1']])
+          ],
+          num_buckets=100,
+          hash_key=ragged_array_ops._DEFAULT_CROSS_HASH_KEY + 1,
+          expected_hashed=ragged_const([[31]])),
+      dict(
+          testcase_name='HashedZeroKey',
+          inputs=[
+              ragged_const([['batch1-FC1-F1']]),
+              ragged_const([['batch1-FC2-F1']]),
+              ragged_const([['batch1-FC3-F1']])
+          ],
+          hash_key=0,
+          expected_hashed=ragged_const([[9077905385164735582]]),
+          matches_sparse_cross=False  # sparse treats hash_key=0 as None.
+      ),
+      dict(
+          testcase_name='UInt64',
+          inputs=[ragged_const([[2**64 - 1]], dtype=dtypes.uint64)],
+          expected=ragged_const([[b'-1']])),
+  ])
+  def testRaggedCross(self,
+                      inputs,
+                      num_buckets=0,
+                      hash_key=None,
+                      expected=None,
+                      expected_hashed=None,
+                      matches_sparse_cross=True):
+    ragged_cross = ragged_array_ops.cross(inputs)
+    ragged_cross_hashed = ragged_array_ops.cross_hashed(inputs, num_buckets,
+                                                        hash_key)
+
+    if expected is not None:
+      self.assertAllEqual(ragged_cross, expected)
+    if expected_hashed is not None:
+      self.assertAllEqual(ragged_cross_hashed, expected_hashed)
+
+    if matches_sparse_cross:
+      # Check that ragged.cross & sparse.cross match.
+      sparse_inputs = [self._ragged_to_sparse(t) for t in inputs]
+      sparse_cross = sparse_ops.sparse_cross(sparse_inputs)
+      self.assertAllEqual(ragged_cross,
+                          ragged_tensor.RaggedTensor.from_sparse(sparse_cross))
+
+      # Check that ragged.cross_hashed & sparse.cross_hashed match.
+      sparse_inputs = [self._ragged_to_sparse(t) for t in inputs]
+      sparse_cross_hashed = sparse_ops.sparse_cross_hashed(
+          sparse_inputs, num_buckets, hash_key)
+      self.assertAllEqual(
+          ragged_cross_hashed,
+          ragged_tensor.RaggedTensor.from_sparse(sparse_cross_hashed))
+
+  def testRaggedCrossLargeBatch(self):
+    batch_size = 5000
+    inputs = [
+        ragged_const([[1, 2, 3]] * batch_size),
+        ragged_const([[b'4']] * batch_size),
+        dense_const([[5]] * batch_size),
+        sparse_const([[6, 7]] * batch_size)
+    ]
+
+    expected = [[
+        b'1_X_4_X_5_X_6', b'1_X_4_X_5_X_7', b'2_X_4_X_5_X_6', b'2_X_4_X_5_X_7',
+        b'3_X_4_X_5_X_6', b'3_X_4_X_5_X_7'
+    ]] * batch_size
+
+    ragged_cross = ragged_array_ops.cross(inputs)
+
+    # Note: we don't use assertAllEqual here because if they don't match,
+    # then the code in assertAllEqual that tries to build the error message
+    # is very slow, causing the test to timeout.
+    # pylint: disable=g-generic-assert
+    self.assertTrue(self.evaluate(ragged_cross).to_list() == expected)
+
+  @parameterized.named_parameters([
+      dict(
+          testcase_name='BadDType',
+          inputs=[ragged_const([[1.1], [2.2, 3.3]])],
+          message=r'Unexpected dtype for inputs\[0\]',
+      ),
+      dict(
+          testcase_name='StaticBatchSizeMismatch1',
+          inputs=[ragged_const([[1]]), ragged_const([[2], [3]])],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='inputs must all have the same batch dimension size',
+      ),
+      dict(
+          testcase_name='StaticBatchSizeMismatch2',
+          inputs=[ragged_const([[1]]), dense_const([[2], [3]])],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='inputs must all have the same batch dimension size',
+      ),
+      dict(
+          testcase_name='3DDenseTensor',
+          inputs=[dense_const([[[1]]])],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2',
+      ),
+      dict(
+          testcase_name='0DDenseTensor',
+          inputs=[dense_const(1)],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2',
+      ),
+  ])
+  def testStaticError(self, inputs, exception=ValueError, message=None):
+    with self.assertRaisesRegex(exception, message):
+      ragged_array_ops.cross(inputs)
+
+  @parameterized.named_parameters([
+      dict(
+          testcase_name='3DRaggedTensor',
+          inputs=[ragged_const([[[1]]], ragged_rank=1)],
+          message='tf.ragged.cross only supports inputs with rank=2',
+      ),
+      dict(
+          testcase_name='0DDenseTensor',
+          inputs=[dense_const(1)],
+          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2',
+      ),
+      dict(
+          testcase_name='1DDenseTensor',
+          inputs=[dense_const([1])],
+          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2',
+      ),
+      dict(
+          testcase_name='3DDenseTensor',
+          inputs=[dense_const([[[1]]])],
+          signature=[[tensor_spec.TensorSpec(None, dtypes.int32)]],
+          exception=(ValueError, errors.InvalidArgumentError),
+          message='tf.ragged.cross only supports inputs with rank=2',
+      ),
+  ])
+  def testRuntimeError(self,
+                       inputs,
+                       exception=errors.InvalidArgumentError,
+                       message=None,
+                       signature=None):
+    @def_function.function(input_signature=signature)
+    def fn(x):
+      return ragged_array_ops.cross(x)
+
+    with self.assertRaisesRegex(exception, message):
+      self.evaluate(fn(inputs))
+
+  def _ragged_to_sparse(self, t):
+    if ragged_tensor.is_ragged(t):
+      return ragged_tensor.convert_to_tensor_or_ragged_tensor(t).to_sparse()
+    elif sparse_tensor.is_sparse(t):
+      return sparse_tensor.SparseTensor.from_value(t)
+    else:
+      return ops.convert_to_tensor(t)
+
+  def testSparseValuesAndIndicesMustMatch(self):
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError),
+        'sparse indices and values must have the same length'):
+      self.evaluate(gen_ragged_array_ops.RaggedCross(
+          ragged_values=[],
+          ragged_row_splits=[],
+          sparse_indices=[[5]],
+          sparse_values=[],
+          sparse_shape=[5],
+          dense_inputs=[['a']],
+          input_order='RD',
+          hashed_output=False,
+          num_buckets=5,
+          hash_key=2,
+          out_values_type=dtypes.string,
+          out_row_splits_type=dtypes.int64))
+
+  def testRaggedValuesAndSplitsMustMatch(self):
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError),
+        'ragged values and splits must have the same length'):
+      self.evaluate(gen_ragged_array_ops.RaggedCross(
+          ragged_values=[['a']],
+          ragged_row_splits=[],
+          sparse_indices=[],
+          sparse_values=[],
+          sparse_shape=[],
+          dense_inputs=[['a']],
+          input_order='RD',
+          hashed_output=False,
+          num_buckets=5,
+          hash_key=2,
+          out_values_type=dtypes.string,
+          out_row_splits_type=dtypes.int64))
+
+  @parameterized.named_parameters([
+      dict(testcase_name='EmptySplits', ragged_splits=[]),
+      dict(
+          testcase_name='NegativeSplits', ragged_splits=[-216, -114, -58, -54]
+      ),
+      dict(testcase_name='TooLargeValueSplits', ragged_splits=[0, 1, 2, 10]),
+      dict(testcase_name='UnsortedSplits', ragged_splits=[0, 2, 2, 1]),
+  ])
+  def testRaggedCrossInvalidRaggedSplits(self, ragged_splits):
+    # Test case in GitHub isseu 59114.
+    with self.assertRaisesRegex(
+        (ValueError, errors.InvalidArgumentError), 'Invalid ragged splits'
+    ):
+      ragged_values_0_tensor = ops.convert_to_tensor(np.ones([3], dtype=str))
+      ragged_values_0 = array_ops.identity(ragged_values_0_tensor)
+      ragged_values = [
+          ragged_values_0,
+      ]
+      ragged_row_splits_0_tensor = ragged_const(
+          ragged_splits, dtype=dtypes.int64
+      )
+      ragged_row_splits_0 = array_ops.identity(ragged_row_splits_0_tensor)
+      ragged_row_splits = [
+          ragged_row_splits_0,
+      ]
+      self.evaluate(
+          gen_ragged_array_ops.RaggedCross(
+              ragged_values=ragged_values,
+              ragged_row_splits=ragged_row_splits,
+              sparse_indices=[],
+              sparse_values=[],
+              sparse_shape=[],
+              dense_inputs=[],
+              input_order='R',
+              hashed_output=False,
+              num_buckets=0,
+              hash_key=956888297470,
+              out_values_type=7,
+              out_row_splits_type=9,
+          )
+      )
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_dispatch.py

@@ -0,0 +1,160 @@
+# 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.
+# ==============================================================================
+"""Operator dispatch for RaggedTensors."""
+
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_shape
+from tensorflow.python.util import dispatch
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_export
+from tensorflow.python.util import tf_inspect
+
+
+@dispatch.dispatch_for_unary_elementwise_apis(ragged_tensor.Ragged)
+def ragged_unary_elementwise_op(op, x):
+  """Unary elementwise api handler for RaggedTensors."""
+  x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x)
+  return x.with_values(op(x.values))
+
+
+# TODO(martinz): This is deprecated. Delete.
+def ragged_binary_elementwise_op(op, x, y):
+  """Binary elementwise api handler for RaggedTensors."""
+  x_is_ragged = ragged_tensor.is_ragged(x)
+  y_is_ragged = ragged_tensor.is_ragged(y)
+
+  # Convert args to tensors.
+  x = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+      x, preferred_dtype=(y.dtype if y_is_ragged else None))
+  y = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+      y, preferred_dtype=x.dtype)
+
+  if x_is_ragged and y_is_ragged:
+    x, y = ragged_tensor.match_row_splits_dtypes(x, y)
+
+  # Perform broadcasting, when appropriate
+  if ((x_is_ragged and y_is_ragged) or
+      (x_is_ragged and x.flat_values.shape.ndims <= y.shape.ndims) or
+      (y_is_ragged and y.flat_values.shape.ndims <= x.shape.ndims)):
+    # If both x and y are ragged, they must have the same row_splits_dtype now.
+    if x_is_ragged:
+      dim_size_dtype = x.row_splits.dtype
+    else:
+      dim_size_dtype = y.row_splits.dtype
+
+    shape_x = ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(
+        x, dim_size_dtype=dim_size_dtype)
+    shape_y = ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(
+        y, dim_size_dtype=dim_size_dtype)
+    bcast_shape = ragged_tensor_shape.broadcast_dynamic_shape(shape_x, shape_y)
+    x = ragged_tensor_shape.broadcast_to(
+        x, bcast_shape, broadcast_inner_dimensions=False)
+    y = ragged_tensor_shape.broadcast_to(
+        y, bcast_shape, broadcast_inner_dimensions=False)
+
+  x_values = x.flat_values if ragged_tensor.is_ragged(x) else x
+  y_values = y.flat_values if ragged_tensor.is_ragged(y) else y
+  mapped_values = op(x_values, y_values)
+  if isinstance(mapped_values, bool):
+    return mapped_values  # Special case for tensor_equals.
+  if ragged_tensor.is_ragged(x):
+    return x.with_flat_values(mapped_values)
+  else:
+    return y.with_flat_values(mapped_values)
+
+
+# TODO(edloper): Update the documentation generation tools to automatically
+# build lists of which types are supported by which ops (and then delete all
+# the following code).
+
+
+# We don't need to register a separate delegation handler for these v1 ops,
+# since they delegate to the v2 ops (which already have a handler).  But we
+# still want to include them in the ragged_op_list() output.
+_V2_OPS_THAT_ARE_DELEGATED_TO_FROM_V1_OPS = [
+    math_ops.reduce_sum,
+    math_ops.reduce_prod,
+    math_ops.reduce_min,
+    math_ops.reduce_max,
+    math_ops.reduce_mean,
+    math_ops.reduce_variance,
+    math_ops.reduce_std,
+    math_ops.reduce_any,
+    math_ops.reduce_all,
+    string_ops.string_to_number,
+    string_ops.string_to_hash_bucket,
+    string_ops.reduce_join_v2,
+]
+
+
+def _ragged_op_signature(op, ragged_args, ragged_varargs=False):
+  """Returns a signature for the given op, marking ragged args in bold."""
+  op_name = tf_export.get_canonical_name_for_symbol(op)
+  argspec = tf_inspect.getfullargspec(op)
+  arg_names = argspec.args
+
+  # Mark ragged arguments in bold.
+  for pos in ragged_args:
+    arg_names[pos] = '**' + arg_names[pos] + '**'
+
+  # Add argument defaults.
+  if argspec.defaults is not None:
+    for pos in range(-1, -len(argspec.defaults) - 1, -1):
+      arg_names[pos] += '=`{!r}`'.format(argspec.defaults[pos])
+
+  # Add varargs and keyword args
+  if argspec.varargs:
+    if ragged_varargs:
+      arg_names.append('***' + argspec.varargs + '**')
+    else:
+      arg_names.append('*' + argspec.varargs)
+  if argspec.varkw:
+    arg_names.append('**' + argspec.varkw)
+
+  return '* `tf.{}`({})'.format(op_name, ', '.join(arg_names))
+
+
+def _op_is_in_tf_version(op, version):
+  if version == 1:
+    return (tf_export.get_v1_names(tf_decorator.unwrap(op)[1]) or
+            op in _V2_OPS_THAT_ARE_DELEGATED_TO_FROM_V1_OPS)
+  elif version == 2:
+    return tf_export.get_v2_names(tf_decorator.unwrap(op)[1])
+  else:
+    raise ValueError('Expected version 1 or 2.')
+
+
+def ragged_op_list(tf_version=2):
+  """Returns a string listing operations that have dispathers registered."""
+  lines = []
+  api_signatures = dispatch.type_based_dispatch_signatures_for(
+      ragged_tensor.RaggedTensor)
+  for api, signatures in api_signatures.items():
+    arg_names = tf_inspect.getargspec(api).args
+    ragged_args = set()
+    for signature in signatures:
+      for arg in signature:
+        ragged_args.add(arg if isinstance(arg, int) else arg_names.index(arg))
+    if _op_is_in_tf_version(api, tf_version):
+      lines.append(_ragged_op_signature(api, ragged_args))
+
+  lines.append(
+      _ragged_op_signature(logging_ops.print_v2, [], ragged_varargs=True))
+  return ('\n\n### Additional ops that support `RaggedTensor`\n\n'
+          'Arguments that accept `RaggedTensor`s are marked in **bold**.\n\n' +
+          '\n'.join(sorted(lines)) + 'n')

tensorflow/python/ops/ragged/ragged_dynamic_partition_op_test.py

@@ -0,0 +1,253 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for ragged_array_ops_stack.stack_dynamic_partitions."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import data_flow_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSegmentStackOpTest(test_util.TensorFlowTestCase,
+                               parameterized.TestCase):
+
+  @parameterized.parameters([
+      dict(  # empty inputs
+          data=[],
+          partitions=[],
+          num_partitions=0,
+          expected=[],
+          expected_ragged_rank=1),
+      dict(  # empty data, num_partitions>0
+          data=[],
+          partitions=[],
+          num_partitions=3,
+          expected=[[], [], []]),
+      dict(  # 1D data, 1D partitions (docstring example)
+          data=['a', 'b', 'c', 'd', 'e'],
+          partitions=[3, 0, 2, 2, 3],
+          num_partitions=5,
+          expected=[['b'], [], ['c', 'd'], ['a', 'e'], []]),
+      dict(  # 2D data, 1D partitions
+          data=[['a', 'b'], ['c', 'd'], ['e', 'f'], ['g', 'h']],
+          data_ragged_rank=0,
+          partitions=[2, 1, 2, 3],
+          num_partitions=4,
+          expected=[[], [['c', 'd']], [['a', 'b'], ['e', 'f']], [['g', 'h']]],
+          expected_ragged_rank=1),
+      dict(  # 2D ragged data, 1D partitions
+          data=[['a'], ['b', 'c', 'd'], [], ['e', 'f']],
+          data_ragged_rank=1,
+          partitions=[2, 1, 2, 3],
+          num_partitions=4,
+          expected=[[], [['b', 'c', 'd']], [['a'], []], [['e', 'f']]],
+          expected_ragged_rank=2),
+      dict(  # 2D data, 2D partitions
+          data=[['a', 'b'], ['c', 'd'], ['e', 'f'], ['g', 'h']],
+          data_ragged_rank=0,
+          partitions=[[3, 0], [2, 2], [4, 3], [2, 0]],
+          num_partitions=5,
+          expected=[['b', 'h'], [], ['c', 'd', 'g'], ['a', 'f'], ['e']]),
+      dict(  # 2D ragged data, 2D ragged partitions
+          data=[['a', 'b'], ['c', 'd'], ['e', 'f'], ['g', 'h']],
+          data_ragged_rank=0,
+          partitions=[[3, 0], [2, 2], [4, 3], [2, 0]],
+          num_partitions=5,
+          expected=[['b', 'h'], [], ['c', 'd', 'g'], ['a', 'f'], ['e']]),
+      dict(  # 3D data, 1d partitions
+          data=[[['a', 'b'], ['c', 'd']], [['e', 'f'], ['g', 'h']]],
+          data_ragged_rank=0,
+          partitions=[1, 0],
+          num_partitions=2,
+          expected=[[[['e', 'f'], ['g', 'h']]], [[['a', 'b'], ['c', 'd']]]],
+          expected_ragged_rank=1),
+      dict(  # 3D data (ragged_rank=1), 1d partitions
+          data=[[['a', 'b'], ['c', 'd']], [['e', 'f']]],
+          data_ragged_rank=1,
+          partitions=[2, 0],
+          num_partitions=3,
+          expected=[[[['e', 'f']]], [], [[['a', 'b'], ['c', 'd']]]],
+          expected_ragged_rank=2),
+      dict(  # 3D data (ragged_rank=2), 1d partitions
+          data=[[['a', 'b'], ['c', 'd']], [['e', 'f', 'g', 'h']]],
+          data_ragged_rank=2,
+          partitions=[2, 0],
+          num_partitions=3,
+          expected=[[[['e', 'f', 'g', 'h']]], [], [[['a', 'b'], ['c', 'd']]]],
+          expected_ragged_rank=3),
+      dict(  # 3D data, 2d partitions
+          data=[[['a', 'b'], ['c', 'd']], [['e', 'f'], ['g', 'h']]],
+          data_ragged_rank=0,
+          partitions=[[1, 0], [0, 3]],
+          segment_ids_ragged_rank=0,
+          num_partitions=4,
+          expected=[[['c', 'd'], ['e', 'f']], [['a', 'b']], [], [['g', 'h']]],
+          expected_ragged_rank=1),
+      dict(  # 3D data (ragged_rank=1), 2d partitions
+          data=[[['a', 'b'], ['c', 'd']], [['e', 'f']]],
+          data_ragged_rank=1,
+          partitions=[[1, 0], [0]],
+          segment_ids_ragged_rank=1,
+          num_partitions=2,
+          expected=[[['c', 'd'], ['e', 'f']], [['a', 'b']]],
+          expected_ragged_rank=1),
+      dict(  # 3D data (ragged_rank=2), 2d partitions
+          data=[[['a', 'b'], ['c', 'd']], [['e', 'f', 'g', 'h']]],
+          data_ragged_rank=2,
+          partitions=[[1, 0], [0]],
+          segment_ids_ragged_rank=1,
+          num_partitions=3,
+          expected=[[['c', 'd'], ['e', 'f', 'g', 'h']], [['a', 'b']], []],
+          expected_ragged_rank=2),
+      dict(  # 3D data (ragged_rank=2), 3d partitions (ragged_rank=2)
+          data=[[['a', 'b'], ['c', 'd']], [['e', 'f', 'g', 'h']]],
+          data_ragged_rank=2,
+          partitions=[[[3, 0], [1, 2]], [[1, 1, 0, 1]]],
+          segment_ids_ragged_rank=2,
+          num_partitions=4,
+          expected=[['b', 'g'], ['c', 'e', 'f', 'h'], ['d'], ['a']]),
+      dict(  # 0D data, 0D partitions
+          data='a',
+          partitions=3,
+          num_partitions=5,
+          expected=[[], [], [], ['a'], []]),
+      dict(  # 1D data, 0D partitions
+          data=['a', 'b', 'c'],
+          partitions=3,
+          num_partitions=5,
+          expected=[[], [], [], [['a', 'b', 'c']], []],
+          expected_ragged_rank=1),
+      dict(  # 2D data, 0D partitions
+          data=[['a', 'b'], ['c', 'd']],
+          data_ragged_rank=0,
+          partitions=3,
+          num_partitions=5,
+          expected=[[], [], [], [[['a', 'b'], ['c', 'd']]], []],
+          expected_ragged_rank=1),
+      dict(  # 2D data (ragged_rank=1), 0D partitions
+          data=[['a', 'b'], ['c']],
+          data_ragged_rank=1,
+          partitions=3,
+          num_partitions=5,
+          expected=[[], [], [], [[['a', 'b'], ['c']]], []],
+          expected_ragged_rank=3),
+  ])
+  def testRaggedSegmentStack(self,
+                             data,
+                             partitions,
+                             num_partitions,
+                             expected,
+                             data_ragged_rank=None,
+                             segment_ids_ragged_rank=None,
+                             expected_ragged_rank=None):
+    for seg_dtype in [dtypes.int32, dtypes.int64]:
+      data_tensor = ragged_factory_ops.constant(
+          data, row_splits_dtype=seg_dtype, ragged_rank=data_ragged_rank)
+      segment_ids_tensor = ragged_factory_ops.constant(
+          partitions,
+          dtype=seg_dtype,
+          row_splits_dtype=seg_dtype,
+          ragged_rank=segment_ids_ragged_rank)
+      expected_tensor = ragged_factory_ops.constant(
+          expected,
+          row_splits_dtype=seg_dtype,
+          ragged_rank=expected_ragged_rank)
+      result = ragged_array_ops.stack_dynamic_partitions(
+          data_tensor, segment_ids_tensor, num_partitions)
+      self.assertAllEqual(result, expected_tensor)
+
+      # Check that it's equivalent to tf.stack(dynamic_partition(...)),
+      # where applicable.
+      if (data_ragged_rank == 0 and segment_ids_ragged_rank == 0 and
+          seg_dtype == dtypes.int32):
+        equiv = ragged_concat_ops.stack(
+            data_flow_ops.dynamic_partition(data_tensor, segment_ids_tensor,
+                                            num_partitions))
+        self.assertAllEqual(result, self.evaluate(equiv).to_list())
+
+  @parameterized.parameters([
+      dict(
+          data=['a', 'b', 'c'],
+          partitions=[2, -1, 0],
+          num_partitions=10,
+          error='must be non-negative'),
+      dict(
+          data=['a', 'b', 'c'],
+          partitions=[2, 10, 0],
+          num_partitions=1,
+          error='partitions must be less than num_partitions'),
+      dict(
+          data=['a', 'b', 'c'],
+          partitions=[2, 10, 0],
+          num_partitions=10,
+          error='partitions must be less than num_partitions'),
+      dict(
+          data=[['a', 'b'], ['c']],
+          partitions=[[2], [3, 0]],
+          num_partitions=10,
+          error='data and partitions have incompatible ragged shapes'),
+  ])
+  def testRuntimeError(self, data, partitions, num_partitions, error):
+    data = ragged_factory_ops.constant(data)
+    partitions = ragged_factory_ops.constant(partitions, dtype=dtypes.int64)
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                error):
+      self.evaluate(
+          ragged_array_ops.stack_dynamic_partitions(data, partitions,
+                                                    num_partitions))
+
+  @parameterized.parameters([
+      dict(
+          data=['a', 'b', 'c'],
+          partitions=[1, 2],
+          num_partitions=10,
+          error=r'Shapes \(2,\) and \(3,\) are incompatible'),
+      dict(
+          data=[['a', 'b'], ['c', 'd']],
+          partitions=[[1, 2, 3], [4, 5, 6]],
+          num_partitions=10,
+          error=r'Shapes \(2, 3\) and \(2, 2\) are incompatible'),
+      dict(
+          data=['a', 'b', 'c'],
+          partitions=[1, 2, 3],
+          num_partitions=[1, 2, 3],
+          error='must have rank 0'),
+  ])
+  def testStaticError(self, data, partitions, num_partitions, error):
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                error):
+      ragged_array_ops.stack_dynamic_partitions(data, partitions,
+                                                num_partitions)
+
+  def testUnknownRankError(self):
+    if context.executing_eagerly():
+      return
+    partitions = array_ops.placeholder(dtypes.int32, None)
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                'partitions must have known rank'):
+      ragged_array_ops.stack_dynamic_partitions(['a', 'b', 'c'], partitions, 10)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_eager_test.py

@@ -0,0 +1,55 @@
+# 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.
+# ==============================================================================
+"""Tests for tf.ragged in eager execution mode."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+class RaggedTensorTest(test_util.TensorFlowTestCase,
+                       parameterized.TestCase):
+
+  @parameterized.parameters([
+      dict(pylist=[[b'a', b'b'], [b'c']]),
+      dict(pylist=[[[1, 2], [3]], [[4, 5, 6], [], [7]]]),
+      dict(pylist=[[[1, 2], [3, 4]], [[5, 6], [], [7, 8]]], ragged_rank=1),
+  ])
+  def testRaggedTensorToList(self, pylist, ragged_rank=None):
+    rt = ragged_factory_ops.constant(pylist, ragged_rank)
+    self.assertAllEqual(rt, pylist)
+
+  @parameterized.parameters([
+      dict(pylist=[[b'a', b'b'], [b'c']],
+           expected_str="[[b'a', b'b'], [b'c']]"),
+      dict(pylist=[[[1, 2], [3]], [[4, 5, 6], [], [7]]],
+           expected_str='[[[1, 2], [3]], [[4, 5, 6], [], [7]]]'),
+      dict(pylist=[[0, 1], np.arange(2, 2000)],
+           expected_str='[[0, 1], [2, 3, 4, ..., 1997, 1998, 1999]]'),
+      dict(pylist=[[[0, 1]], [np.arange(2, 2000)]],
+           expected_str='[[[0, 1]],\n [[2, 3, 4, ..., 1997, 1998, 1999]]]'),
+  ])
+  def testRaggedTensorStr(self, pylist, expected_str):
+    rt = ragged_factory_ops.constant(pylist)
+    self.assertEqual(str(rt), f'<tf.RaggedTensor {expected_str}>')
+
+
+if __name__ == '__main__':
+  ops.enable_eager_execution()
+  googletest.main()

tensorflow/python/ops/ragged/ragged_embedding_ops.py

@@ -0,0 +1,432 @@
+# Copyright 2023 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.
+# ==============================================================================
+"""Embedding operations."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import embedding_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import resource_variable_ops
+from tensorflow.python.ops import variables
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import dispatch
+
+
+@dispatch.dispatch_for_api(embedding_ops.embedding_lookup)
+def embedding_lookup(
+    params,
+    ids: ragged_tensor.Ragged,
+    partition_strategy="mod",
+    name=None,
+    validate_indices=True,  # pylint: disable=unused-argument
+    max_norm=None,
+):
+  """Look up the ragged ids in a list of embedding tensors.
+
+  Args:
+    params: A tensor representing the complete embedding tensor having the shape
+      [e1, ...eM]
+    ragged_ids: A 'RaggedTensor' with type 'int32' or 'int64' containing the ids
+      to be looked up in 'params' of shape [r0, ..rN]. Values must be in the
+      range '[0, params.shape[0]]'.
+    partition_strategy: A string specifying the partitioning strategy.
+    max_norm: If not `None`, each embedding is clipped if its l2-norm is larger
+      than this value.
+    name: A name for the operation (optional)
+
+  Returns:
+    A ragged tensor of shape [r0, r1, ...rN, e1, ...eM].
+
+  Raises:
+    ValueError: When params is empty or the type of the ids is not int32 or
+      int64.
+  """
+  if params is None:
+    raise ValueError("params must be specified.")
+  if isinstance(params, (list, tuple)) and not params:
+    raise ValueError("params should not be empty.")
+  if ids.dtype != dtypes.int32 and ids.dtype != dtypes.int64:
+    raise ValueError(
+        "The values contained by the inputs have type "
+        f"{str(ids.dtype)}"
+        " and cannot be processed. All values"
+        " should be indices, either of type `int32` or `int64`."
+    )
+
+  with ops.name_scope(name, "embedding_lookup_ragged") as name:
+    looked_up_ragged = ragged_functional_ops.map_flat_values(
+        embedding_ops.embedding_lookup,
+        params=params,
+        ids=ids,
+        partition_strategy=partition_strategy,
+        max_norm=max_norm,
+    )
+
+    return looked_up_ragged
+
+
+@dispatch.dispatch_for_api(embedding_ops.embedding_lookup_sparse)
+def embedding_lookup_sparse(
+    params,
+    sp_ids: ragged_tensor.Ragged,
+    sp_weights,
+    partition_strategy="mod",
+    name=None,
+    combiner=None,
+    max_norm=None,
+    allow_fast_lookup=False,
+):
+  """Looks up embeddings for the given ids and weights from a list of tensors.
+
+  This op assumes that there is at least one id for each row in the dense tensor
+  represented by sp_ids (i.e. there are no rows with empty features), and that
+  all the indices of sp_ids are in canonical row-major order.
+
+  `sp_ids` and `sp_weights` (if not None) are `RaggedTensor`s with rank of 2.
+  Embeddings are always aggregated along the last dimension.
+
+  It also assumes that all id values lie in the range [0, p0), where p0
+  is the sum of the size of params along dimension 0.
+
+  Args:
+    params: A single tensor representing the complete embedding tensor, or a
+      list tensors all of same shape except for the first dimension,
+      representing sharded embedding tensors. Alternatively, a
+      `PartitionedVariable`, created by partitioning along dimension 0. Each
+      element must be appropriately sized for the given `partition_strategy`.
+    sp_ids: `RaggedTensor` with rank 2. The rank is not verified for performance
+      reasons.
+    sparse_weights: `RaggedTensor` of same type and shape as `sparse_ids`,
+      containing float / double weights corresponding to `sparse_ids`, or `None`
+      if all weights are assumed to be 1.0.
+    partition_strategy: A string specifying the partitioning strategy, relevant
+      if `len(params) > 1`. Currently `"div"` and `"mod"` are supported. Default
+      is `"mod"`. See `tf.nn.embedding_lookup` for more details.
+    name: Optional name for the op.
+    combiner: A string specifying the reduction op. Currently "mean", "sqrtn"
+      and "sum" are supported. "sum" computes the weighted sum of the embedding
+      results for each row. "mean" is the weighted sum divided by the total
+      weight. "sqrtn" is the weighted sum divided by the square root of the sum
+      of the squares of the weights. Defaults to `mean`.
+    max_norm: If not `None`, each embedding is clipped if its l2-norm is larger
+      than this value, before combining.
+    allow_fast_lookup: An optional boolean specifying whether to allow
+      simplified embedding lookups when `params` is a single tensor and
+      `max_norm` is `None`. Setting this flag to `True` during training can
+      cause the use of dense gradients with increased memory footprint.
+
+  Returns:
+    A dense tensor representing the combined embeddings for the
+    sparse ids. For each row in the dense tensor represented by `sp_ids`, the op
+    looks up the embeddings for all ids in that row, multiplies them by the
+    corresponding weight, and combines these embeddings as specified.
+
+    In other words, if
+
+      `shape(combined params) = [p0, p1, ..., pm]`
+
+    and
+
+      `shape(sp_ids) = shape(sp_weights) = [d0, d1]`
+
+    then
+
+      `shape(output) = [d0, p1, ..., pm]`.
+
+    For instance, if params is a 10x20 matrix, and sp_ids / sp_weights are
+
+      ```python
+      [0, 0]: id 1, weight 2.0
+      [0, 1]: id 3, weight 0.5
+      [1, 0]: id 0, weight 1.0
+      [2, 3]: id 1, weight 3.0
+      ```
+
+    with `combiner`="mean", then the output will be a 3x20 matrix where
+
+      ```python
+      output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
+      output[1, :] = (params[0, :] * 1.0) / 1.0
+      output[2, :] = (params[1, :] * 3.0) / 3.0
+      ```
+
+  Raises:
+    TypeError: If `sp_weights` is neither `None` nor of the same type as
+      `sp_ids`.
+    ValueError: If `combiner` is not one of {"mean", "sqrtn", "sum"}.
+  """
+  rt_ids = sp_ids
+  rt_weights = sp_weights
+  if combiner is None:
+    combiner = "mean"
+  if combiner not in ("mean", "sqrtn", "sum"):
+    raise ValueError(
+        f"combiner must be one of 'mean', 'sqrtn' or 'sum', got {combiner}"
+    )
+  if isinstance(params, variables.PartitionedVariable):
+    params = list(params)  # Iterate to get the underlying Variables.
+  if not isinstance(params, list):
+    params = [params]
+  ignore_weights = rt_weights is None
+  if not ignore_weights:
+    if not isinstance(rt_weights, ragged_tensor.RaggedTensor):
+      raise TypeError(
+          f"sp_ids must be of the same type as sp_weights, "
+          f"received {{type(sp_ids).__name__!r}} for sp_ids and "
+          f"{{type(sp_weights).__name__!r}} for sp_weights."
+      )
+    rt_ids.values.get_shape().assert_is_compatible_with(
+        rt_weights.values.get_shape()
+    )
+    rt_ids.get_shape().assert_is_compatible_with(rt_weights.get_shape())
+
+  with ops.name_scope(
+      name, "embedding_lookup_sparse", params + [rt_ids]
+  ) as name:
+    segment_ids = rt_ids.value_rowids()
+    ids = rt_ids.flat_values
+
+    return embedding_ops.embedding_lookup_sparse_impl(
+        params,
+        segment_ids,
+        sp_weights,
+        ids,
+        combiner,
+        ignore_weights,
+        max_norm,
+        allow_fast_lookup,
+        partition_strategy,
+        name,
+    )
+
+
+@dispatch.dispatch_for_api(embedding_ops.safe_embedding_lookup_sparse)
+def safe_embedding_lookup_sparse(
+    embedding_weights,
+    sparse_ids: ragged_tensor.Ragged,
+    sparse_weights=None,
+    combiner="mean",
+    default_id=None,
+    name=None,
+    partition_strategy="div",
+    max_norm=None,
+    allow_fast_lookup=False,
+):
+  """Lookup embedding results, accounting for invalid IDs and empty features.
+
+  The partitioned embedding in `embedding_weights` must all be the same shape
+  except for the first dimension. The first dimension is allowed to vary as the
+  vocabulary size is not necessarily a multiple of `P`.  `embedding_weights`
+  may be a `PartitionedVariable` as returned by using
+  `tf.compat.v1.get_variable()` with a
+  partitioner.
+
+  Invalid IDs (< 0) are pruned from input IDs and weights, as well as any IDs
+  with non-positive weight. For an entry with no features, the embedding vector
+  for `default_id` is returned, or the 0-vector if `default_id` is not supplied.
+
+  The ids and weights may be multi-dimensional `SparseTensor`s or
+  `RaggedTensor`s with rank of 2. For `SpareTensor`s with left-aligned non-zero
+  entries which can be described as `RaggedTensor`s, use of `RaggedTensor`s can
+  yield higher performance. Embeddings are always aggregated along the last
+  dimension.
+
+  Args:
+    embedding_weights: A single tensor representing the complete embedding
+      tensor, or a list tensors all of same shape except for the first
+      dimension, representing sharded embedding tensors. Alternatively, a
+      `PartitionedVariable`, created by partitioning along dimension 0. Each
+      element must be appropriately sized for the given `partition_strategy`.
+    sp_ids: `RaggedTensor` with rank 2. The rank is not verified for performance
+      reasons.
+    sparse_weights: `RaggedTensor` of same type and shape as `sparse_ids`,
+      containing float weights corresponding to `sparse_ids`, or `None` if all
+      weights are assumed to be 1.0.
+    combiner: A string specifying how to combine embedding results for each
+      entry. Currently "mean", "sqrtn" and "sum" are supported, with "mean" the
+      default.
+    default_id: The id to use for an entry with no features.
+    name: A name for this operation (optional).
+    partition_strategy: A string specifying the partitioning strategy. Currently
+      `"div"` and `"mod"` are supported. Default is `"div"`.
+    max_norm: If not `None`, all embeddings are l2-normalized to max_norm before
+      combining.
+    allow_fast_lookup: An optional boolean specifying whether to allow
+      simplified embedding lookups when `params` is a single tensor and
+      `max_norm` is `None`. Setting this flag to `True` during training can
+      cause the use of dense gradients with increased memory footprint.
+
+  Returns:
+    A dense tensor representing the combined embeddings for the
+    sparse ids. For each row in the dense tensor represented by `sp_ids`, the op
+    looks up the embeddings for all ids in that row, multiplies them by the
+    corresponding weight, and combines these embeddings as specified.
+
+    In other words, if
+
+      `shape(combined embedding_weights) = [p0, p1, ..., pm]`
+
+    and
+
+      `shape(sparse_ids) = shape(sparse_weights) = [d0, d1, ..., dn]`
+
+    then
+
+      `shape(output) = [d0, d1, ... dn-1, p1, ..., pm]`.
+
+    For instance, if params is a 10x20 matrix, and sp_ids / sp_weights are
+
+      ```python
+      [0, 0]: id 1, weight 2.0
+      [0, 1]: id 3, weight 0.5
+      [1, 0]: id -1, weight 1.0
+      [2, 3]: id 1, weight 3.0
+      ```
+
+    `default_id` is 0.
+
+    with `combiner`="mean", then the output will be a 3x20 matrix where
+
+      ```python
+      output[0, :] = (params[1, :] * 2.0 + params[3, :] * 0.5) / (2.0 + 0.5)
+      output[1, :] = (params[0, :] * 1.0) / 1.0
+      output[2, :] = (params[1, :] * 3.0) / 3.0
+      ```
+
+  Raises:
+    ValueError: if `embedding_weights` is empty.
+  """
+  ragged_ids = sparse_ids
+  ragged_weights = sparse_weights
+  if embedding_weights is None:
+    raise ValueError(f"Missing embedding_weights {embedding_weights}.")
+  if isinstance(embedding_weights, variables.PartitionedVariable):
+    embedding_weights = list(embedding_weights)  # get underlying Variables.
+  if not isinstance(embedding_weights, list):
+    embedding_weights = [embedding_weights]
+  if len(embedding_weights) < 1:
+    raise ValueError(f"Missing embedding_weights {embedding_weights}.")
+
+  dtype = ragged_weights.dtype if ragged_weights is not None else None
+  embedding_weights = [
+      w
+      if (
+          resource_variable_ops.is_resource_variable(w)
+          and dtype in (None, w.dtype)
+      )
+      else ops.convert_to_tensor(w, dtype=dtype)
+      for w in embedding_weights
+  ]
+
+  with ops.name_scope(
+      name, "embedding_lookup", embedding_weights + [ragged_ids, ragged_weights]
+  ) as scope:
+    # Prune invalid ids and weights.
+    ragged_ids, ragged_weights = _prune_invalid_ids_ragged(
+        ragged_ids, ragged_weights
+    )
+    if combiner != "sum":
+      ragged_ids, ragged_weights = _prune_invalid_weights_ragged(
+          ragged_ids, ragged_weights
+      )
+    ragged_ids, is_row_empty = ragged_array_ops.fill_empty_rows(
+        ragged_ids, default_id or 0
+    )
+    if ragged_weights is not None:
+      ragged_weights, _ = ragged_array_ops.fill_empty_rows(ragged_weights, 1.0)
+
+    result = embedding_lookup_sparse(
+        embedding_weights,
+        ragged_ids,
+        ragged_weights,
+        combiner=combiner,
+        partition_strategy=partition_strategy,
+        name=None if default_id is None else scope,
+        max_norm=max_norm,
+        allow_fast_lookup=allow_fast_lookup,
+    )
+
+    if default_id is None:
+      # Broadcast is_row_empty to the same shape as embedding_lookup_result,
+      # for use in Select.
+      is_row_empty = array_ops.tile(
+          array_ops.reshape(is_row_empty, [-1, 1]),
+          array_ops_stack.stack([1, array_ops.shape(result)[1]]),
+      )
+
+      result = array_ops.where(
+          is_row_empty, array_ops.zeros_like(result), result, name=scope
+      )
+
+    return result
+
+
+def _prune_invalid_ids_ragged(ids, weights):
+  """Prune invalid IDs (< 0) from the input ids and weights."""
+  is_id_valid = math_ops.greater_equal(ids.values, 0)
+  nrows = ids.nrows()
+  # TODO(philipphack): Consider calling ragged_array_ops.boolean_mask once the
+  # resulting performance is comparable to array_ops.boolean_mask. Currently,
+  # ragged_array_ops.boolean_mask constructs the returned RaggedTensor by
+  # calling its from_row_splits method which does not set value_row_ids and
+  # requires it to be computed on demand.
+  pruned_values = array_ops.boolean_mask_v2(ids.values, is_id_valid)
+  pruned_value_rowids = array_ops.boolean_mask_v2(
+      ids.value_rowids(), is_id_valid
+  )
+  ids = ragged_tensor.RaggedTensor.from_value_rowids(
+      pruned_values, pruned_value_rowids, nrows=nrows, validate=False
+  )
+  if weights is not None:
+    pruned_weights_values = array_ops.boolean_mask_v2(
+        weights.values, is_id_valid
+    )
+    weights = ragged_tensor.RaggedTensor.from_value_rowids(
+        pruned_weights_values, pruned_value_rowids, nrows=nrows, validate=False
+    )
+
+  return ids, weights
+
+
+def _prune_invalid_weights_ragged(ids, weights):
+  """Prune invalid weights (< 0) from the input ids and weights."""
+  if weights is not None:
+    is_weights_valid = math_ops.greater(weights.values, 0)
+    nrows = ids.nrows()
+    # TODO(philipphack): Consider calling ragged_array_ops.boolean_mask once the
+    # resulting performance is comparable to array_ops.boolean_mask. Currently,
+    # ragged_array_ops.boolean_mask constructs the returned RaggedTensor by
+    # calling its from_row_splits method which does not set value_row_ids and
+    # requires it to be computed on demand.
+    pruned_values = array_ops.boolean_mask_v2(ids.values, is_weights_valid)
+    pruned_value_rowids = array_ops.boolean_mask_v2(
+        ids.value_rowids(), is_weights_valid
+    )
+    ids = ragged_tensor.RaggedTensor.from_value_rowids(
+        pruned_values, pruned_value_rowids, nrows=nrows, validate=False
+    )
+
+    pruned_weights_values = array_ops.boolean_mask_v2(
+        weights.values, is_weights_valid
+    )
+    weights = ragged_tensor.RaggedTensor.from_value_rowids(
+        pruned_weights_values, pruned_value_rowids, nrows=nrows, validate=False
+    )
+
+  return ids, weights

tensorflow/python/ops/ragged/ragged_expand_dims_op_test.py

@@ -0,0 +1,122 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.expand_dims."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedExpandDimsOpTest(test_util.TensorFlowTestCase,
+                             parameterized.TestCase):
+
+  # An example 4-d ragged tensor with shape [3, (D2), (D3), 2], and the
+  # expected result calling for expand_dims on each axis.  c.f. the table of
+  # expected result shapes in the ragged_array_ops.expand_dims docstring.
+  EXAMPLE4D = [[[[1, 1], [2, 2]], [[3, 3]]],
+               [],
+               [[], [[4, 4], [5, 5], [6, 6]]]]  # pyformat: disable
+  EXAMPLE4D_EXPAND_AXIS = {
+      0: [EXAMPLE4D],
+      1: [[d0] for d0 in EXAMPLE4D],
+      2: [[[d1] for d1 in d0] for d0 in EXAMPLE4D],
+      3: [[[[d2] for d2 in d1] for d1 in d0] for d0 in EXAMPLE4D],
+      4: [[[[[d3] for d3 in d2] for d2 in d1] for d1 in d0] for d0 in EXAMPLE4D]
+  }
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Docstring examples: 2D Ragged Inputs
+      dict(rt_input=[[1, 2], [3]],
+           axis=0,
+           expected=[[[1, 2], [3]]],
+           expected_shape=[1, 2, None]),
+      dict(rt_input=[[1, 2], [3]],
+           axis=1,
+           expected=[[[1, 2]], [[3]]],
+           expected_shape=[2, 1, None]),
+      dict(rt_input=[[1, 2], [3]],
+           axis=2,
+           expected=[[[1], [2]], [[3]]],
+           expected_shape=[2, None, 1]),
+
+      #=========================================================================
+      # 2D Tensor Inputs
+      dict(rt_input=[[1, 2], [3, 4], [5, 6]],
+           ragged_rank=0,
+           axis=0,
+           expected=[[[1, 2], [3, 4], [5, 6]]],
+           expected_shape=[1, 3, 2]),
+      dict(rt_input=[[1, 2], [3, 4], [5, 6]],
+           ragged_rank=0,
+           axis=1,
+           expected=[[[1, 2]], [[3, 4]], [[5, 6]]],
+           expected_shape=[3, 1, 2]),
+      dict(rt_input=[[1, 2], [3, 4], [5, 6]],
+           ragged_rank=0,
+           axis=2,
+           expected=[[[1], [2]], [[3], [4]], [[5], [6]]],
+           expected_shape=[3, 2, 1]),
+
+      #=========================================================================
+      # 4D Ragged Inputs: [3, (D2), (D3), 2]
+      # c.f. the table of expected result shapes in the expand_dims docstring.
+      dict(rt_input=EXAMPLE4D,
+           ragged_rank=2,
+           axis=0,
+           expected=EXAMPLE4D_EXPAND_AXIS[0],
+           expected_shape=[1, 3, None, None, 2]),
+      dict(rt_input=EXAMPLE4D,
+           ragged_rank=2,
+           axis=1,
+           expected=EXAMPLE4D_EXPAND_AXIS[1],
+           expected_shape=[3, 1, None, None, 2]),
+      dict(rt_input=EXAMPLE4D,
+           ragged_rank=2,
+           axis=2,
+           expected=EXAMPLE4D_EXPAND_AXIS[2],
+           expected_shape=[3, None, 1, None, 2]),
+      dict(rt_input=EXAMPLE4D,
+           ragged_rank=2,
+           axis=3,
+           expected=EXAMPLE4D_EXPAND_AXIS[3],
+           expected_shape=[3, None, None, 1, 2]),
+      dict(rt_input=EXAMPLE4D,
+           ragged_rank=2,
+           axis=4,
+           expected=EXAMPLE4D_EXPAND_AXIS[4],
+           expected_shape=[3, None, None, 2, 1]),
+  ])  # pyformat: disable
+  def testRaggedExpandDims(self,
+                           rt_input,
+                           axis,
+                           expected,
+                           ragged_rank=None,
+                           expected_shape=None):
+    rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
+    expanded = ragged_array_ops.expand_dims(rt, axis=axis)
+    self.assertEqual(expanded.shape.ndims, rt.shape.ndims + 1)
+    if expected_shape is not None:
+      self.assertEqual(expanded.shape.as_list(), expected_shape)
+
+    self.assertAllEqual(expanded, expected)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_factory_ops.py

@@ -0,0 +1,386 @@
+# 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.
+# ==============================================================================
+"""Operations for constructing RaggedTensors."""
+
+from typing import Union
+
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.numpy_compat import np_reshape
+from tensorflow.python.util.tf_export import tf_export
+
+
+# ===============================================================================
+# Op to construct a constant RaggedTensor from a nested Python list.
+# ===============================================================================
+@tf_export("ragged.constant")
+@dispatch.add_dispatch_support
+def constant(
+    pylist,
+    dtype=None,
+    ragged_rank=None,
+    inner_shape=None,
+    name=None,
+    row_splits_dtype=dtypes.int64,
+) -> Union[ragged_tensor.RaggedTensor, ops._EagerTensorBase, ops.Operation]:
+  """Constructs a constant RaggedTensor from a nested Python list.
+
+  Example:
+
+  >>> tf.ragged.constant([[1, 2], [3], [4, 5, 6]])
+  <tf.RaggedTensor [[1, 2], [3], [4, 5, 6]]>
+
+  All scalar values in `pylist` must have the same nesting depth `K`, and the
+  returned `RaggedTensor` will have rank `K`.  If `pylist` contains no scalar
+  values, then `K` is one greater than the maximum depth of empty lists in
+  `pylist`.  All scalar values in `pylist` must be compatible with `dtype`.
+
+  Args:
+    pylist: A nested `list`, `tuple` or `np.ndarray`.  Any nested element that
+      is not a `list`, `tuple` or `np.ndarray` must be a scalar value compatible
+      with `dtype`.
+    dtype: The type of elements for the returned `RaggedTensor`.  If not
+      specified, then a default is chosen based on the scalar values in
+      `pylist`. If there are no scalar values in `pylist`, then the default is
+      `tf.float32`.
+    ragged_rank: An integer specifying the ragged rank of the returned
+      `RaggedTensor`.  Must be nonnegative and less than `K`. Defaults to
+      `max(0, K - 1)` if `inner_shape` is not specified.  Defaults to `max(0, K
+      - 1 - len(inner_shape))` if `inner_shape` is specified.
+    inner_shape: A tuple of integers specifying the shape for individual inner
+      values in the returned `RaggedTensor`.  Defaults to `()` if `ragged_rank`
+      is not specified.  If `ragged_rank` is specified, then a default is chosen
+      based on the contents of `pylist`.
+    name: A name prefix for the returned tensor (optional).
+    row_splits_dtype: data type for the constructed `RaggedTensor`'s row_splits.
+      One of `tf.int32` or `tf.int64`.
+
+  Returns:
+    A potentially ragged tensor with rank `K` and the specified `ragged_rank`,
+    containing the values from `pylist`.
+
+  Raises:
+    ValueError: If the scalar values in `pylist` have inconsistent nesting
+      depth; or if ragged_rank or inner_shape are incompatible with `pylist`.
+  """
+  def ragged_factory(values, row_splits):
+    row_splits = constant_op.constant(row_splits, dtype=row_splits_dtype)
+    return ragged_tensor.RaggedTensor.from_row_splits(values, row_splits,
+                                                      validate=False)
+
+  with ops.name_scope(name, "RaggedConstant"):
+    return _constant_value(ragged_factory, constant_op.constant, pylist, dtype,
+                           ragged_rank, inner_shape)
+
+
+@tf_export(v1=["ragged.constant_value"])
+@dispatch.add_dispatch_support
+def constant_value(
+    pylist,
+    dtype=None,
+    ragged_rank=None,
+    inner_shape=None,
+    row_splits_dtype="int64",
+) -> Union[ragged_tensor_value.RaggedTensorValue, np.ndarray]:
+  """Constructs a RaggedTensorValue from a nested Python list.
+
+  Warning: This function returns a `RaggedTensorValue`, not a `RaggedTensor`.
+  If you wish to construct a constant `RaggedTensor`, use
+  [`ragged.constant(...)`](constant.md) instead.
+
+  Example:
+
+  >>> tf.compat.v1.ragged.constant_value([[1, 2], [3], [4, 5, 6]])
+  tf.RaggedTensorValue(values=array([1, 2, 3, 4, 5, 6]),
+                       row_splits=array([0, 2, 3, 6]))
+
+  All scalar values in `pylist` must have the same nesting depth `K`, and the
+  returned `RaggedTensorValue` will have rank `K`.  If `pylist` contains no
+  scalar values, then `K` is one greater than the maximum depth of empty lists
+  in `pylist`.  All scalar values in `pylist` must be compatible with `dtype`.
+
+  Args:
+    pylist: A nested `list`, `tuple` or `np.ndarray`.  Any nested element that
+      is not a `list` or `tuple` must be a scalar value compatible with `dtype`.
+    dtype: `numpy.dtype`.  The type of elements for the returned `RaggedTensor`.
+      If not specified, then a default is chosen based on the scalar values in
+      `pylist`.
+    ragged_rank: An integer specifying the ragged rank of the returned
+      `RaggedTensorValue`.  Must be nonnegative and less than `K`. Defaults to
+      `max(0, K - 1)` if `inner_shape` is not specified.  Defaults to `max(0, K
+      - 1 - len(inner_shape))` if `inner_shape` is specified.
+    inner_shape: A tuple of integers specifying the shape for individual inner
+      values in the returned `RaggedTensorValue`.  Defaults to `()` if
+      `ragged_rank` is not specified.  If `ragged_rank` is specified, then a
+      default is chosen based on the contents of `pylist`.
+    row_splits_dtype: data type for the constructed `RaggedTensorValue`'s
+      row_splits.  One of `numpy.int32` or `numpy.int64`.
+
+  Returns:
+    A `tf.RaggedTensorValue` or `numpy.array` with rank `K` and the specified
+    `ragged_rank`, containing the values from `pylist`.
+
+  Raises:
+    ValueError: If the scalar values in `pylist` have inconsistent nesting
+      depth; or if ragged_rank or inner_shape are incompatible with `pylist`.
+  """
+  if dtype is not None and isinstance(dtype, dtypes.DType):
+    dtype = dtype.as_numpy_dtype
+  row_splits_dtype = dtypes.as_dtype(row_splits_dtype).as_numpy_dtype
+  def _ragged_factory(values, row_splits):
+    row_splits = np.array(row_splits, dtype=row_splits_dtype)
+    return ragged_tensor_value.RaggedTensorValue(values, row_splits)
+
+  def _inner_factory(pylist, dtype, shape, name=None):  # pylint: disable=unused-argument
+    if dtype is object or dtype is None:
+      return np_reshape(np.array(pylist, dtype=dtype), shape)
+    else:
+      return np_reshape(np.array(pylist).astype(dtype), shape)
+
+  return _constant_value(
+      _ragged_factory, _inner_factory, pylist, dtype, ragged_rank, inner_shape
+  )
+
+
+def _constant_value(
+    ragged_factory, inner_factory, pylist, dtype, ragged_rank, inner_shape
+):
+  """Constructs a constant RaggedTensor or RaggedTensorValue.
+
+  Args:
+    ragged_factory: A factory function with the signature:
+      `ragged_factory(values, row_splits)`
+    inner_factory: A factory function with the signature: `inner_factory(pylist,
+      dtype, shape, name)`
+    pylist: A nested `list`, `tuple` or `np.ndarray`.
+    dtype: Data type for returned value.
+    ragged_rank: Ragged rank for returned value.
+    inner_shape: Inner value shape for returned value.
+
+  Returns:
+    A value returned by `ragged_factory` or `inner_factory`.
+
+  Raises:
+    ValueError: If the scalar values in `pylist` have inconsistent nesting
+      depth; or if ragged_rank or inner_shape are incompatible with `pylist`.
+  """
+  if ragged_tensor.is_ragged(pylist):
+    raise TypeError("pylist may not be a RaggedTensor or RaggedTensorValue.")
+  # np.ndim builds an array, so we short-circuit lists and tuples.
+  if not isinstance(pylist, (list, tuple)) and np.ndim(pylist) == 0:
+    # Scalar value
+    if ragged_rank is not None and ragged_rank != 0:
+      raise ValueError("Invalid pylist=%r: incompatible with ragged_rank=%d" %
+                       (pylist, ragged_rank))
+    if inner_shape is not None and inner_shape:
+      raise ValueError(
+          "Invalid pylist=%r: incompatible with dim(inner_shape)=%d" %
+          (pylist, len(inner_shape)))
+    return inner_factory(pylist, dtype, ())
+
+  if ragged_rank is not None and ragged_rank < 0:
+    raise ValueError(
+        "Invalid ragged_rank=%r: must be nonnegative" % ragged_rank)
+
+  # Find the depth of scalar values in `pylist`.
+  scalar_depth, max_depth = _find_scalar_and_max_depth(pylist)
+  if scalar_depth is not None:
+    if max_depth > scalar_depth:
+      raise ValueError("Invalid pylist=%r: empty list nesting is greater "
+                       "than scalar value nesting" % pylist)
+    if ragged_rank is not None and max_depth < ragged_rank:
+      raise ValueError(f"Invalid pylist={pylist}, max depth smaller than "
+                       f"ragged_rank={ragged_rank}")
+
+  # If both inner_shape and ragged_rank were specified, then check that
+  # they are compatible with pylist.
+  if inner_shape is not None and ragged_rank is not None:
+    expected_depth = ragged_rank + len(inner_shape) + 1
+    if ((scalar_depth is not None and expected_depth != scalar_depth) or
+        (scalar_depth is None and expected_depth < max_depth)):
+      raise ValueError(
+          "Invalid pylist=%r: incompatible with ragged_rank=%d "
+          "and dim(inner_shape)=%d" % (pylist, ragged_rank, len(inner_shape)))
+
+  # Check if the result is a `Tensor`.
+  if (ragged_rank == 0 or
+      (ragged_rank is None and
+       ((max_depth < 2) or
+        (inner_shape is not None and max_depth - len(inner_shape) < 2)))):
+    return inner_factory(pylist, dtype, inner_shape)
+
+  # Compute default value for inner_shape.
+  if inner_shape is None:
+    if ragged_rank is None:
+      inner_shape = ()
+    else:
+      inner_shape = _default_inner_shape_for_pylist(pylist, ragged_rank)
+
+  # Compute default value for ragged_rank.
+  if ragged_rank is None:
+    if scalar_depth is None:
+      ragged_rank = max(1, max_depth - 1)
+    else:
+      ragged_rank = max(1, scalar_depth - 1 - len(inner_shape))
+
+  # Build the splits for each ragged rank, and concatenate the inner values
+  # into a single list.
+  nested_splits = []
+  values = pylist
+  for dim in range(ragged_rank):
+    nested_splits.append([0])
+    concatenated_values = []
+    for row in values:
+      nested_splits[dim].append(nested_splits[dim][-1] + len(row))
+      concatenated_values.extend(row)
+    values = concatenated_values
+
+  values = inner_factory(
+      values, dtype=dtype, shape=(len(values),) + inner_shape, name="values")
+  for row_splits in reversed(nested_splits):
+    values = ragged_factory(values, row_splits)
+  return values
+
+
+def _find_scalar_and_max_depth(pylist):
+  """Finds nesting depth of scalar values in pylist.
+
+  Args:
+    pylist: A nested python `list` or `tuple`.
+
+  Returns:
+    A tuple `(scalar_depth, max_depth)`.  `scalar_depth` is the nesting
+    depth of scalar values in `pylist`, or `None` if `pylist` contains no
+    scalars.  `max_depth` is the maximum depth of `pylist` (including
+    empty lists).
+
+  Raises:
+    ValueError: If pylist has inconsistent nesting depths for scalars.
+  """
+  # Check if pylist is not scalar. np.ndim builds an array, so we
+  # short-circuit lists and tuples.
+  if isinstance(pylist, (list, tuple)) or np.ndim(pylist) != 0:
+    scalar_depth = None
+    max_depth = 1
+    for child in pylist:
+      child_scalar_depth, child_max_depth = _find_scalar_and_max_depth(child)
+      if child_scalar_depth is not None:
+        if scalar_depth is not None and scalar_depth != child_scalar_depth + 1:
+          raise ValueError("all scalar values must have the same nesting depth")
+        scalar_depth = child_scalar_depth + 1
+      max_depth = max(max_depth, child_max_depth + 1)
+    return (scalar_depth, max_depth)
+  return (0, 0)
+
+
+def _default_inner_shape_for_pylist(pylist, ragged_rank):
+  """Computes a default inner shape for the given python list."""
+
+  def get_inner_shape(item):
+    """Returns the inner shape for a python list `item`."""
+    if not isinstance(item, (list, tuple)) and np.ndim(item) == 0:
+      return ()
+    # Note that we need this check here in case `item` is not a Python list but
+    # fakes as being one (pylist). For a scenario of this, see test added in
+    # https://github.com/tensorflow/tensorflow/pull/48945
+    elif len(item) > 0:  # pylint: disable=g-explicit-length-test
+      return (len(item),) + get_inner_shape(item[0])
+    return (0,)
+
+  def check_inner_shape(item, shape):
+    """Checks that `item` has a consistent shape matching `shape`."""
+    is_nested = isinstance(item, (list, tuple)) or np.ndim(item) != 0
+    if is_nested != bool(shape):
+      raise ValueError("inner values have inconsistent shape")
+    if is_nested:
+      if shape[0] != len(item):
+        raise ValueError("inner values have inconsistent shape")
+      for child in item:
+        check_inner_shape(child, shape[1:])
+
+  # Collapse the ragged layers to get the list of inner values.
+  flat_values = pylist
+  for dim in range(ragged_rank):
+    if not all(
+        isinstance(v, (list, tuple)) or np.ndim(v) != 0 for v in flat_values):
+      raise ValueError("pylist has scalar values depth %d, but ragged_rank=%d "
+                       "requires scalar value depth greater than %d" %
+                       (dim + 1, ragged_rank, ragged_rank))
+    flat_values = sum((list(v) for v in flat_values), [])
+
+  # Compute the inner shape looking only at the leftmost elements; and then
+  # use check_inner_shape to verify that other elements have the same shape.
+  inner_shape = get_inner_shape(flat_values)
+  check_inner_shape(flat_values, inner_shape)
+  return inner_shape[1:]
+
+
+@tf_export(v1=["ragged.placeholder"])
+@dispatch.add_dispatch_support
+def placeholder(dtype, ragged_rank, value_shape=None, name=None):
+  """Creates a placeholder for a `tf.RaggedTensor` that will always be fed.
+
+  **Important**: This ragged tensor will produce an error if evaluated.
+  Its value must be fed using the `feed_dict` optional argument to
+  `Session.run()`, `Tensor.eval()`, or `Operation.run()`.
+
+
+  Args:
+    dtype: The data type for the `RaggedTensor`.
+    ragged_rank: The ragged rank for the `RaggedTensor`
+    value_shape: The shape for individual flat values in the `RaggedTensor`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `RaggedTensor` that may be used as a handle for feeding a value, but
+    not evaluated directly.
+
+  Raises:
+    RuntimeError: if eager execution is enabled
+
+  @compatibility(TF2)
+  This API is not compatible with eager execution and `tf.function`. To migrate
+  to TF2, rewrite the code to be compatible with eager execution. Check the
+  [migration
+  guide](https://www.tensorflow.org/guide/migrate#1_replace_v1sessionrun_calls)
+  on replacing `Session.run` calls. In TF2, you can just pass tensors directly
+  into ops and layers. If you want to explicitly set up your inputs, also see
+  [Keras functional API](https://www.tensorflow.org/guide/keras/functional) on
+  how to use `tf.keras.Input` to replace `tf.compat.v1.ragged.placeholder`.
+  `tf.function` arguments also do the job of `tf.compat.v1.ragged.placeholder`.
+  For more details please read [Better
+  performance with tf.function](https://www.tensorflow.org/guide/function).
+  @end_compatibility
+  """
+  if ragged_rank == 0:
+    return array_ops.placeholder(dtype, value_shape, name)
+
+  with ops.name_scope(name, "RaggedPlaceholder", []):
+    flat_shape = tensor_shape.TensorShape([None]).concatenate(value_shape)
+    result = array_ops.placeholder(dtype, flat_shape, "flat_values")
+    for i in reversed(range(ragged_rank)):
+      row_splits = array_ops.placeholder(dtypes.int64, [None],
+                                         "row_splits_%d" % i)
+      result = ragged_tensor.RaggedTensor.from_row_splits(result, row_splits,
+                                                          validate=False)
+    return result

tensorflow/python/ops/ragged/ragged_factory_ops_test.py

@@ -0,0 +1,358 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Tests that ragged tensors work with GPU, such as placement of int and string.
+
+Test using ragged tensors with map_fn and distributed dataset. Since GPU does
+not support strings, ragged tensors containing string should always be placed
+on CPU.
+"""
+
+from absl.testing import parameterized
+from tensorflow.python.data.ops import dataset_ops
+from tensorflow.python.data.ops import multi_device_iterator_ops
+from tensorflow.python.distribute import mirrored_strategy
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import map_fn
+from tensorflow.python.ops.math_ops import reduce_sum
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.string_ops import string_to_hash_bucket
+from tensorflow.python.platform import test
+
+
+def ragged_int64():
+  return ragged_factory_ops.constant(
+      [
+          [3, 1, 4, 1],
+          [],
+          [5, 9, 2],
+          [6],
+          [],
+          [3, 1, 4, 1],
+          [3, 1],
+          [2, 1, 4, 1],
+      ],
+      dtype=dtypes.int64,
+  )
+
+
+def ragged_str():
+  return ragged_factory_ops.constant([
+      ['3', '1', '4', '1'],
+      [],
+      ['5', '9', '2'],
+      ['6'],
+      [],
+      ['3', '1', '4', '1'],
+      ['3', '1'],
+      ['2', '1', '4', '1'],
+  ])
+
+
+def dense_str():
+  return constant_op.constant([
+      ['3', '1', '4', '1'],
+      ['', '', '', ''],
+      ['5', '9', '2', ''],
+      ['6', '', '', ''],
+      ['', '', '', ''],
+      ['3', '1', '4', '1'],
+      ['3', '1', '', ''],
+      ['2', '1', '4', '1'],
+  ])
+
+
+class RaggedFactoryOpsTest(test_util.TensorFlowTestCase,
+                           parameterized.TestCase):
+
+  @parameterized.named_parameters(
+      ('Int64', ragged_int64,),
+      ('Str', ragged_str,),
+  )
+  def testRaggedWithMapFn(self, ragged_factory):
+
+    @def_function.function
+    def map_fn_producer(inputs):
+      return map_fn.map_fn_v2(lambda x: x, inputs)
+
+    t = ragged_factory()
+    result = self.evaluate(map_fn_producer(t))
+    self.assertAllEqual(t.values, result.values)
+
+  # When drop_remainder=True, the IteratorGetNext op is used and the type
+  # of the iterator's output is the type of the dataset elements.
+  # When drop_remainder=False the IteratorGetNextAsOptional op is used and
+  # the type of the iterator's output is an Optional that contains
+  # the type of the dataset elements. Optionals are encoded using variants.
+  @parameterized.named_parameters(
+      ('Int64Drop', ragged_int64, True),
+      ('StrDrop', ragged_str, True),
+      ('Int64NoDrop', ragged_int64, False),
+      ('StrNoDrop', ragged_str, False),
+  )
+  def testRaggedWithMultiDeviceIterator(self, ragged_factory, drop_remainder):
+
+    @def_function.function
+    def dataset_producer(t):
+      ragged_ds = dataset_ops.Dataset.from_tensor_slices(t).batch(
+          2, drop_remainder)
+      it = multi_device_iterator_ops.MultiDeviceIterator(ragged_ds, ['GPU:0'])
+      with ops.device_v2('GPU:0'):
+        return it.get_next_as_optional()
+
+    t = ragged_factory()
+    result = dataset_producer(t)
+    self.assertAllEqual(
+        self.evaluate(t[0]), self.evaluate(result[0].get_value()[0]))
+
+  @parameterized.named_parameters(
+      ('Int65Drop', ragged_int64, True),
+      ('StrDrop', ragged_str, True),
+      ('Int65NoDrop', ragged_int64, False),
+      ('StrNoDrop', ragged_str, False),
+  )
+  @test_util.run_gpu_only
+  def testRaggedWithDistributedDatasetIter(self, ragged_factory,
+                                           drop_remainder):
+
+    @def_function.function
+    def distributed_dataset_producer(t):
+      strategy = mirrored_strategy.MirroredStrategy(['GPU:0', 'GPU:1'])
+      ragged_ds = dataset_ops.Dataset.from_tensor_slices(t).batch(2)
+      dist_dataset = strategy.experimental_distribute_dataset(ragged_ds)
+      ds = iter(dist_dataset)
+      result0 = strategy.experimental_local_results(next(ds))
+      result1 = strategy.experimental_local_results(next(ds))
+      result2 = strategy.experimental_local_results(next(ds))
+      result3 = strategy.experimental_local_results(next(ds))
+      # Reach the end of the iterator
+      for ignore in ds:  # pylint: disable=unused-variable
+        pass
+      return result0, result1, result2, result3
+
+    t = ragged_factory()
+
+    result0, result1, result2, result3 = distributed_dataset_producer(t)
+    self.assertAllEqual(self.evaluate(t[0]), self.evaluate(result0[0][0]))
+    self.assertAllEqual(self.evaluate(t[1]), self.evaluate(result0[1][0]))
+    self.assertAllEqual(self.evaluate(t[2]), self.evaluate(result1[0][0]))
+    self.assertAllEqual(self.evaluate(t[3]), self.evaluate(result1[1][0]))
+    self.assertAllEqual(self.evaluate(t[4]), self.evaluate(result2[0][0]))
+    self.assertAllEqual(self.evaluate(t[5]), self.evaluate(result2[1][0]))
+    self.assertAllEqual(self.evaluate(t[6]), self.evaluate(result3[0][0]))
+    self.assertAllEqual(self.evaluate(t[7]), self.evaluate(result3[1][0]))
+
+  @parameterized.named_parameters(
+      ('Int65Drop', ragged_int64, True),
+      ('StrDrop', ragged_str, True),
+      ('Int65NoDrop', ragged_int64, False),
+      ('StrNoDrop', ragged_str, False),
+  )
+  @test_util.run_v2_only
+  def testRaggedWithDistributedDatasetReplicaFn(self, ragged_factory,
+                                                drop_remainder):
+
+    def distributed_dataset_producer(t):
+      strategy = mirrored_strategy.MirroredStrategy(['GPU:0', 'GPU:1'])
+      ragged_ds = dataset_ops.Dataset.from_tensor_slices(t).batch(2)
+      dist_dataset = strategy.experimental_distribute_dataset(ragged_ds)
+
+      @def_function.function
+      def replica_fn(elem):
+        return elem
+
+      result = []
+      for x in dist_dataset:
+        result.append(strategy.run(replica_fn, args=(x,)))
+      return result
+
+    t = ragged_factory()
+    result = distributed_dataset_producer(t)
+
+    self.assertAllEqual(
+        self.evaluate(t[0]), self.evaluate(result[0].values[0][0]))
+    self.assertAllEqual(
+        self.evaluate(t[1]), self.evaluate(result[0].values[1][0]))
+    self.assertAllEqual(
+        self.evaluate(t[2]), self.evaluate(result[1].values[0][0]))
+    self.assertAllEqual(
+        self.evaluate(t[3]), self.evaluate(result[1].values[1][0]))
+    self.assertAllEqual(
+        self.evaluate(t[4]), self.evaluate(result[2].values[0][0]))
+    self.assertAllEqual(
+        self.evaluate(t[5]), self.evaluate(result[2].values[1][0]))
+    self.assertAllEqual(
+        self.evaluate(t[6]), self.evaluate(result[3].values[0][0]))
+    self.assertAllEqual(
+        self.evaluate(t[7]), self.evaluate(result[3].values[1][0]))
+
+  @parameterized.named_parameters(
+      ('DenseDrop', dense_str, True),
+      ('RaggedDrop', ragged_str, True),
+      ('DenseNoDrop', dense_str, False),
+      ('RaggedNoDrop', ragged_str, False),
+  )
+  @test_util.run_gpu_only
+  def testIntStringWithDistributedDataset(self, string_factory, drop_remainder):
+
+    @def_function.function
+    def distributed_dataset_producer(t):
+      strategy = mirrored_strategy.MirroredStrategy(['GPU:0', 'GPU:1'])
+      ragged_ds = dataset_ops.Dataset.from_tensor_slices(t).batch(
+          2, drop_remainder)
+      dist_dataset = strategy.experimental_distribute_dataset(ragged_ds)
+      ds = iter(dist_dataset)
+      result0 = strategy.experimental_local_results(next(ds))
+      result1 = strategy.experimental_local_results(next(ds))
+      result2 = strategy.experimental_local_results(next(ds))
+      result3 = strategy.experimental_local_results(next(ds))
+      # Reach the end of the iterator
+      for ignore in ds:  # pylint: disable=unused-variable
+        pass
+      return result0, result1, result2, result3
+
+    ds_dict = {'int': ragged_int64(), 'str': string_factory()}
+    result0, result1, result2, result3 = distributed_dataset_producer(ds_dict)
+    self.assertAllEqual(
+        self.evaluate(ds_dict['int'][0]), self.evaluate(result0[0]['int'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['str'][0]), self.evaluate(result0[0]['str'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['int'][1]), self.evaluate(result0[1]['int'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['str'][1]), self.evaluate(result0[1]['str'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['int'][2]), self.evaluate(result1[0]['int'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['str'][2]), self.evaluate(result1[0]['str'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['int'][3]), self.evaluate(result1[1]['int'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['str'][3]), self.evaluate(result1[1]['str'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['int'][4]), self.evaluate(result2[0]['int'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['str'][4]), self.evaluate(result2[0]['str'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['int'][5]), self.evaluate(result2[1]['int'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['str'][5]), self.evaluate(result2[1]['str'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['int'][6]), self.evaluate(result3[0]['int'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['str'][6]), self.evaluate(result3[0]['str'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['int'][7]), self.evaluate(result3[1]['int'][0]))
+    self.assertAllEqual(
+        self.evaluate(ds_dict['str'][7]), self.evaluate(result3[1]['str'][0]))
+
+  @parameterized.named_parameters(
+      ('DenseDrop', dense_str, True),
+      ('RaggedDrop', ragged_str, True),
+      ('DenseNoDrop', dense_str, False),
+      ('RaggedNoDrop', ragged_str, True),
+  )
+  @test_util.run_v2_only
+  def testOpsWithDistributedDataset(self, string_factory, drop_remainder):
+
+    def distributed_dataset_producer(t):
+      strategy = mirrored_strategy.MirroredStrategy(['GPU:0', 'GPU:1'])
+      ragged_ds = dataset_ops.Dataset.from_tensor_slices(t).batch(
+          2, drop_remainder)
+      dist_dataset = strategy.experimental_distribute_dataset(ragged_ds)
+
+      @def_function.function
+      def replica_fn(elem):
+        # Example of typical preprocessing of string to numeric feature
+        hashed = string_to_hash_bucket(elem['str'], 10)
+        return 1000 * hashed
+
+      result = []
+      for x in dist_dataset:
+        result.append(strategy.run(replica_fn, args=(x,)))
+      return result
+
+    ds_dict = {'str': string_factory()}
+    result = distributed_dataset_producer(ds_dict)
+    expect_length = [len(i) for i in ds_dict['str']]
+    self.assertAllEqual([[5000, 3000, 5000, 3000][:expect_length[0]]],
+                        self.evaluate(result[0].values[0]))
+    self.assertAllEqual([[9000, 9000, 9000, 9000][:expect_length[1]]],
+                        self.evaluate(result[0].values[1]))
+    self.assertAllEqual([[0, 3000, 2000, 9000][:expect_length[2]]],
+                        self.evaluate(result[1].values[0]))
+    self.assertAllEqual([[2000, 9000, 9000, 9000][:expect_length[3]]],
+                        self.evaluate(result[1].values[1]))
+    self.assertAllEqual([[9000, 9000, 9000, 9000][:expect_length[4]]],
+                        self.evaluate(result[2].values[0]))
+    self.assertAllEqual([[5000, 3000, 5000, 3000][:expect_length[5]]],
+                        self.evaluate(result[2].values[1]))
+    self.assertAllEqual([[5000, 3000, 9000, 9000][:expect_length[6]]],
+                        self.evaluate(result[3].values[0]))
+    self.assertAllEqual([[2000, 3000, 5000, 3000][:expect_length[7]]],
+                        self.evaluate(result[3].values[1]))
+
+  @parameterized.named_parameters(
+      ('DenseDrop', dense_str, True),
+      ('RaggedDrop', ragged_str, True),
+      ('DenseNoDrop', dense_str, False),
+      ('RaggedNoDrop', ragged_str, False),
+  )
+  @test_util.run_v2_only
+  def testIntStringOpsWithDistributedDataset(self, string_factory,
+                                             drop_remainder):
+
+    ri = ragged_int64()
+    # To ease testing both dense and ragged strings, pass in the ragged sizes
+    # so the dense strings can be sliced to match.
+    element_sizes = [len(i) for i in ri]
+    ds_dict = {'int': ri, 'str': string_factory(), 'size': element_sizes}
+
+    def distributed_dataset_producer(t):
+      strategy = mirrored_strategy.MirroredStrategy(['GPU:0', 'GPU:1'])
+      ragged_ds = dataset_ops.Dataset.from_tensor_slices(t).batch(
+          2, drop_remainder)
+      dist_dataset = strategy.experimental_distribute_dataset(ragged_ds)
+
+      @def_function.function
+      def replica_fn(elem):
+        # Example of typical preprocessing of string to numeric feature
+        hashed = string_to_hash_bucket(elem['str'], 10)
+        # For dense string case, slice it to size of ragged int
+        hashed_sliced = hashed[:, :elem['size'][0]]
+        # Computation with both feature from string and numeric dataset output
+        return reduce_sum(hashed_sliced) + 100 * reduce_sum(elem['int'])
+
+      result = []
+      for x in dist_dataset:
+        result.append(strategy.run(replica_fn, args=(x,)))
+      return result
+
+    result = distributed_dataset_producer(ds_dict)
+    self.assertAllEqual(916, self.evaluate(result[0].values[0]))
+    self.assertAllEqual(0, self.evaluate(result[0].values[1]))
+    self.assertAllEqual(1605, self.evaluate(result[1].values[0]))
+    self.assertAllEqual(602, self.evaluate(result[1].values[1]))
+    self.assertAllEqual(0, self.evaluate(result[2].values[0]))
+    self.assertAllEqual(916, self.evaluate(result[2].values[1]))
+    self.assertAllEqual(408, self.evaluate(result[3].values[0]))
+    self.assertAllEqual(813, self.evaluate(result[3].values[1]))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/ops/ragged/ragged_fill_empty_rows_op_test.py

@@ -0,0 +1,147 @@
+# Copyright 2022 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.
+# ==============================================================================
+"""Tests tf.ragged.fill_empty_rows."""
+
+import numpy as np
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import gradient_checker
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+class RaggedFillEmptyRowsTest(test_util.TensorFlowTestCase):
+
+  def testFillInt(self):
+    with test_util.use_gpu():
+      default_value = constant_op.constant(-1, dtype=dtypes.int32)
+      ragged_input = ragged_factory_ops.constant(
+          [[], [1, 3, 5, 7], [], [2, 4, 6, 8], []], dtype=dtypes.int32
+      )
+      ragged_output, empty_row_indicator = ragged_array_ops.fill_empty_rows(
+          ragged_input, default_value
+      )
+
+      self.assertAllEqual(
+          ragged_output, [[-1], [1, 3, 5, 7], [-1], [2, 4, 6, 8], [-1]]
+      )
+      self.assertAllEqual(ragged_output.row_lengths(), [1, 4, 1, 4, 1])
+      self.assertAllEqual(
+          ragged_output.values, [-1, 1, 3, 5, 7, -1, 2, 4, 6, 8, -1]
+      )
+      self.assertAllEqual(empty_row_indicator, [True, False, True, False, True])
+
+  @test_util.run_deprecated_v1
+  def testFillFloat(self):
+    with self.session():
+      values = constant_op.constant(
+          [1.0, 3.0, 5.0, 7.0, 2.0, 4.0, 6.0, 8.0], dtype=dtypes.float64
+      )
+      default_value = constant_op.constant(-1.0, dtype=dtypes.float64)
+      ragged_input = ragged_tensor.RaggedTensor.from_row_lengths(
+          values=values, row_lengths=[0, 4, 0, 4, 0]
+      )
+      ragged_output, empty_row_indicator = ragged_array_ops.fill_empty_rows(
+          ragged_input, default_value
+      )
+
+      self.assertAllEqual(ragged_output.row_lengths(), [1, 4, 1, 4, 1])
+      self.assertAllEqual(
+          ragged_output.values,
+          [-1.0, 1.0, 3.0, 5.0, 7.0, -1.0, 2.0, 4.0, 6.0, 8.0, -1.0],
+      )
+      self.assertAllEqual(empty_row_indicator, [True, False, True, False, True])
+
+      values_grad_err = gradient_checker.compute_gradient_error(
+          values, values.shape.as_list(), ragged_output.values, [11], delta=1e-8
+      )
+      self.assertGreater(values_grad_err, 0)
+      self.assertLess(values_grad_err, 1e-8)
+
+      default_value_grad_err = gradient_checker.compute_gradient_error(
+          default_value,
+          default_value.shape.as_list(),
+          ragged_output.values,
+          [11],
+          delta=1e-8,
+      )
+      self.assertGreater(default_value_grad_err, 0)
+      self.assertLess(default_value_grad_err, 1e-8)
+
+  def testFillFloatFunction(self):
+    @def_function.function
+    def func(ragged_input):
+      default_value = constant_op.constant(-1.0, dtype=dtypes.float64)
+      return ragged_array_ops.fill_empty_rows(ragged_input, default_value)
+
+    ragged_input = ragged_factory_ops.constant(
+        [[], [1.0, 3.0, 5.0, 7.0], [], [2.0, 4.0, 6.0, 8.0], []],
+        dtype=dtypes.float64,
+    )
+    ragged_output, empty_row_indicator = func(ragged_input)
+    self.assertAllEqual(ragged_output.row_lengths(), [1, 4, 1, 4, 1])
+    self.assertAllEqual(
+        ragged_output.values,
+        [-1.0, 1.0, 3.0, 5.0, 7.0, -1.0, 2.0, 4.0, 6.0, 8.0, -1.0],
+    )
+    self.assertAllEqual(empty_row_indicator, [True, False, True, False, True])
+
+  def testFillString(self):
+    with test_util.force_cpu():
+      values = constant_op.constant(
+          ["a", "c", "e", "g", "b", "d", "f", "h"], dtype=dtypes.string
+      )
+      default_value = constant_op.constant("x", dtype=dtypes.string)
+      ragged_input = ragged_tensor.RaggedTensor.from_row_lengths(
+          values=values, row_lengths=[0, 4, 0, 4, 0]
+      )
+      ragged_output, empty_row_indicator = ragged_array_ops.fill_empty_rows(
+          ragged_input, default_value
+      )
+
+      self.assertAllEqual(ragged_output.row_lengths(), [1, 4, 1, 4, 1])
+      self.assertAllEqual(
+          ragged_output.values,
+          [b"x", b"a", b"c", b"e", b"g", b"x", b"b", b"d", b"f", b"h", b"x"],
+      )
+      self.assertAllEqual(
+          empty_row_indicator, np.array([1, 0, 1, 0, 1]).astype(np.bool_)
+      )
+
+  def testNoEmptyRows(self):
+    with test_util.use_gpu():
+      values = constant_op.constant(
+          [1, 3, 5, 7, 2, 4, 6, 8], dtype=dtypes.int32
+      )
+      default_value = constant_op.constant(-1, dtype=dtypes.int32)
+      ragged_input = ragged_tensor.RaggedTensor.from_row_lengths(
+          values=values, row_lengths=[4, 4]
+      )
+      ragged_output, empty_row_indicator = ragged_array_ops.fill_empty_rows(
+          ragged_input, default_value
+      )
+
+      self.assertAllEqual(ragged_output.row_lengths(), [4, 4])
+      self.assertAllEqual(ragged_output.values, values)
+      self.assertAllEqual(empty_row_indicator, np.zeros(2).astype(np.bool_))
+
+
+if __name__ == "__main__":
+  googletest.main()

tensorflow/python/ops/ragged/ragged_from_sparse_op_test.py

@@ -0,0 +1,108 @@
+# 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.
+# ==============================================================================
+"""Tests for RaggedTensor.from_sparse."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorToSparseOpTest(test_util.TensorFlowTestCase):
+
+  def testDocStringExample(self):
+    st = sparse_tensor.SparseTensor(
+        indices=[[0, 0], [0, 1], [0, 2], [1, 0], [3, 0]],
+        values=[1, 2, 3, 4, 5],
+        dense_shape=[4, 3])
+    rt = RaggedTensor.from_sparse(st)
+
+    self.assertAllEqual(rt, [[1, 2, 3], [4], [], [5]])
+
+  def testEmpty(self):
+    st = sparse_tensor.SparseTensor(
+        indices=array_ops.zeros([0, 2], dtype=dtypes.int64),
+        values=[],
+        dense_shape=[4, 3])
+    rt = RaggedTensor.from_sparse(st)
+
+    self.assertAllEqual(rt, [[], [], [], []])
+
+  def testBadSparseTensorRank(self):
+    st1 = sparse_tensor.SparseTensor(indices=[[0]], values=[0], dense_shape=[3])
+    self.assertRaisesRegex(ValueError, r'rank\(st_input\) must be 2',
+                           RaggedTensor.from_sparse, st1)
+
+    st2 = sparse_tensor.SparseTensor(
+        indices=[[0, 0, 0]], values=[0], dense_shape=[3, 3, 3])
+    self.assertRaisesRegex(ValueError, r'rank\(st_input\) must be 2',
+                           RaggedTensor.from_sparse, st2)
+
+    if not context.executing_eagerly():
+      st3 = sparse_tensor.SparseTensor(
+          indices=array_ops.placeholder(dtypes.int64),
+          values=[0],
+          dense_shape=array_ops.placeholder(dtypes.int64))
+      self.assertRaisesRegex(ValueError, r'rank\(st_input\) must be 2',
+                             RaggedTensor.from_sparse, st3)
+
+  def testGoodPartialSparseTensorRank(self):
+    if not context.executing_eagerly():
+      st1 = sparse_tensor.SparseTensor(
+          indices=[[0, 0]],
+          values=[0],
+          dense_shape=array_ops.placeholder(dtypes.int64))
+      st2 = sparse_tensor.SparseTensor(
+          indices=array_ops.placeholder(dtypes.int64),
+          values=[0],
+          dense_shape=[4, 3])
+
+      # Shouldn't throw ValueError
+      RaggedTensor.from_sparse(st1)
+      RaggedTensor.from_sparse(st2)
+
+  def testNonRaggedSparseTensor(self):
+    # "index_suffix" means the value of the innermost dimension of the index
+    # (i.e., indices[i][-1]).
+    # See comments in _assert_sparse_indices_are_ragged_right() for more
+    # details/background.
+
+    # index_suffix of first index is not zero.
+    st1 = sparse_tensor.SparseTensor(
+        indices=[[0, 1], [0, 2], [2, 0]], values=[1, 2, 3], dense_shape=[3, 3])
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r'.*SparseTensor is not right-ragged'):
+      self.evaluate(RaggedTensor.from_sparse(st1))
+    # index_suffix of an index that starts a new row is not zero.
+    st2 = sparse_tensor.SparseTensor(
+        indices=[[0, 0], [0, 1], [2, 1]], values=[1, 2, 3], dense_shape=[3, 3])
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r'.*SparseTensor is not right-ragged'):
+      self.evaluate(RaggedTensor.from_sparse(st2))
+    # index_suffix of an index that continues a row skips a cell.
+    st3 = sparse_tensor.SparseTensor(
+        indices=[[0, 1], [0, 1], [0, 3]], values=[1, 2, 3], dense_shape=[3, 3])
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r'.*SparseTensor is not right-ragged'):
+      self.evaluate(RaggedTensor.from_sparse(st3))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_from_tensor_op_test.py

@@ -0,0 +1,648 @@
+# 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.
+# ==============================================================================
+"""Tests for RaggedTensor.from_tensor."""
+
+from absl.testing import parameterized
+
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorFromTensorOpTest(test_util.TensorFlowTestCase,
+                                   parameterized.TestCase):
+
+  def testDocStringExamples(self):
+    # The examples from RaggedTensor.from_tensor.__doc__.
+    dt = constant_op.constant([[5, 7, 0], [0, 3, 0], [6, 0, 0]])
+    self.assertAllEqual(
+        RaggedTensor.from_tensor(dt), [[5, 7, 0], [0, 3, 0], [6, 0, 0]])
+
+    self.assertAllEqual(
+        RaggedTensor.from_tensor(dt, lengths=[1, 0, 3]), [[5], [], [6, 0, 0]])
+
+    self.assertAllEqual(
+        RaggedTensor.from_tensor(dt, padding=0), [[5, 7], [0, 3], [6]])
+
+    dt_3d = constant_op.constant([[[5, 0], [7, 0], [0, 0]],
+                                  [[0, 0], [3, 0], [0, 0]],
+                                  [[6, 0], [0, 0], [0, 0]]])
+    self.assertAllEqual(
+        RaggedTensor.from_tensor(dt_3d, lengths=([2, 0, 3], [1, 1, 2, 0, 1])),
+        [[[5], [7]], [], [[6, 0], [], [0]]])
+
+  @parameterized.parameters(
+      # 2D test cases, no length or padding.
+      {
+          'tensor': [[]],
+          'expected': [[]],
+          'expected_shape': [1, 0],
+      },
+      {
+          'tensor': [[1]],
+          'expected': [[1]],
+          'expected_shape': [1, 1],
+      },
+      {
+          'tensor': [[1, 2]],
+          'expected': [[1, 2]],
+          'expected_shape': [1, 2],
+      },
+      {
+          'tensor': [[1], [2], [3]],
+          'expected': [[1], [2], [3]],
+          'expected_shape': [3, 1],
+      },
+      {
+          'tensor': [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+          'expected': [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+          'expected_shape': [3, 3],
+      },
+      # 3D test cases, no length or padding
+      {
+          'tensor': [[[]]],
+          'expected': [[[]]],
+          'expected_shape': [1, 1, 0],
+      },
+      {
+          'tensor': [[[]]],
+          'expected': [[[]]],
+          'ragged_rank': 1,
+          'expected_shape': [1, 1, 0],
+      },
+      {
+          'tensor': [[[1]]],
+          'expected': [[[1]]],
+          'expected_shape': [1, 1, 1],
+      },
+      {
+          'tensor': [[[1, 2]]],
+          'expected': [[[1, 2]]],
+          'expected_shape': [1, 1, 2],
+      },
+      {
+          'tensor': [[[1, 2], [3, 4]]],
+          'expected': [[[1, 2], [3, 4]]],
+          'expected_shape': [1, 2, 2],
+      },
+      {
+          'tensor': [[[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]]],
+          'expected': [[[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]]],
+          'expected_shape': [4, 1, 2],
+      },
+      {
+          'tensor': [[[1], [2]], [[3], [4]], [[5], [6]], [[7], [8]]],
+          'expected': [[[1], [2]], [[3], [4]], [[5], [6]], [[7], [8]]],
+          'expected_shape': [4, 2, 1],
+      },
+      # 2D test cases, with length
+      {
+          'tensor': [[1]],
+          'lengths': [1],
+          'expected': [[1]],
+          'expected_shape': [1, None],
+      },
+      {
+          'tensor': [[1]],
+          'lengths': [0],
+          'expected': [[]],
+          'expected_shape': [1, None],
+      },
+      {
+          'tensor': [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+          'lengths': [0, 1, 2],
+          'expected': [[], [4], [7, 8]],
+          'expected_shape': [3, None],
+      },
+      {
+          'tensor': [[1, 2, 3], [4, 5, 6], [7, 8, 9]],
+          'lengths': [0, 0, 0],
+          'expected': [[], [], []],
+          'expected_shape': [3, None],
+      },
+      {
+          'tensor': [[1, 2], [3, 4]],
+          'lengths': [2, 2],
+          'expected': [[1, 2], [3, 4]],
+          'expected_shape': [2, None],
+      },
+      {
+          'tensor': [[1, 2], [3, 4]],
+          'lengths': [7, 8],  # lengths > ncols: truncated to ncols
+          'expected': [[1, 2], [3, 4]],
+          'expected_shape': [2, None],
+      },
+      {
+          'tensor': [[1, 2], [3, 4]],
+          'lengths': [-2, -1],  # lengths < 0: treated as zero
+          'expected': [[], []],
+          'expected_shape': [2, None],
+      },
+      # 3D test cases, with length
+      {
+          'tensor': [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+          'lengths': [0, 0],
+          'expected': [[], []],
+          'expected_shape': [2, None, 2],
+      },
+      {
+          'tensor': [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+          'lengths': [1, 2],
+          'expected': [[[1, 2]], [[5, 6], [7, 8]]],
+          'expected_shape': [2, None, 2],
+      },
+      {
+          'tensor': [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+          'lengths': [2, 2],
+          'expected': [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+          'expected_shape': [2, None, 2],
+      },
+      # 2D test cases, with padding
+      {
+          'tensor': [[1]],
+          'padding': 0,
+          'expected': [[1]],
+          'expected_shape': [1, None],
+      },
+      {
+          'tensor': [[0]],
+          'padding': 0,
+          'expected': [[]],
+          'expected_shape': [1, None],
+      },
+      {
+          'tensor': [[0, 1]],
+          'padding': 0,
+          'expected': [[0, 1]],
+          'expected_shape': [1, None],
+      },
+      {
+          'tensor': [[1, 0]],
+          'padding': 0,
+          'expected': [[1]],
+          'expected_shape': [1, None],
+      },
+      {
+          'tensor': [[1, 0, 1, 0, 0, 1, 0, 0]],
+          'padding': 0,
+          'expected': [[1, 0, 1, 0, 0, 1]],
+          'expected_shape': [1, None],
+      },
+      {
+          'tensor': [[3, 7, 0, 0], [2, 0, 0, 0], [5, 0, 0, 0]],
+          'padding': 0,
+          'expected': [[3, 7], [2], [5]],
+          'expected_shape': [3, None],
+      },
+      # 3D test cases, with padding
+      {
+          'tensor': [[[1]]],
+          'padding': [0],
+          'expected': [[[1]]],
+          'expected_shape': [1, None, 1],
+      },
+      {
+          'tensor': [[[0]]],
+          'padding': [0],
+          'expected': [[]],
+          'expected_shape': [1, None, 1],
+      },
+      {
+          'tensor': [[[0, 0], [1, 2]], [[3, 4], [0, 0]]],
+          'padding': [0, 0],
+          'expected': [[[0, 0], [1, 2]], [[3, 4]]],
+          'expected_shape': [2, None, 2],
+      },
+      # 4D test cases, with padding
+      {
+          'tensor': [
+              [[[1, 2], [3, 4]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]],
+              [[[0, 0], [0, 0]], [[5, 6], [7, 8]], [[0, 0], [0, 0]]],
+              [[[0, 0], [0, 0]], [[0, 0], [0, 0]], [[0, 0], [0, 0]]]
+          ],
+          'padding': [[0, 0], [0, 0]],
+          'expected': [
+              [[[1, 2], [3, 4]]],
+              [[[0, 0], [0, 0]], [[5, 6], [7, 8]]],
+              []
+          ],
+          'expected_shape': [3, None, 2, 2],
+      },
+      # 3D test cases, with ragged_rank=2.
+      {
+          'tensor': [[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+          'ragged_rank': 2,
+          'expected': [[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+          'expected_shape': [2, 2, 2],
+      },
+      {
+          'tensor': [[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+          'ragged_rank': 2,
+          'lengths': [2, 0, 2, 1],
+          'expected': [[[1, 2], []], [[5, 6], [7]]],
+          'expected_shape': [2, 2, None],
+      },
+      {
+          'tensor': [[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+          'ragged_rank': 2,
+          'padding': 0,
+          'expected': [[[1], [2, 3]], [[], [4]]],
+          'expected_shape': [2, 2, None],
+      },
+      # 4D test cases, with ragged_rank>1
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'ragged_rank': 2,
+          'expected': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                       [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'expected_shape': [2, 2, 2, 2],
+      },
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'ragged_rank': 3,
+          'expected': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                       [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'expected_shape': [2, 2, 2, 2],
+      },
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'ragged_rank': 2,
+          'padding': [0, 0],
+          'expected': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                       [[[5, 6], [7, 0]], [[0, 8]]]],
+          'expected_shape': [2, 2, None, 2],
+      },
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'lengths': ([2, 2], [1, 2, 2, 1]),
+          'expected': [[[[1, 0]], [[0, 0], [4, 0]]],
+                       [[[5, 6], [7, 0]], [[0, 8]]]],
+          'ragged_rank': 2,
+          'use_ragged_rank': False,  # lengths contains nested_row_lengths.
+          'expected_shape': [2, None, None, 2],
+      },
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'lengths': [[2, 2], [1, 2, 2, 1]],
+          'expected': [[[[1, 0]], [[0, 0], [4, 0]]],
+                       [[[5, 6], [7, 0]], [[0, 8]]]],
+          'ragged_rank': 2,
+          'use_ragged_rank': False,  # lengths contains nested_row_lengths.
+          'expected_shape': [2, None, None, 2],
+      },
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'ragged_rank': 3,
+          'padding': 0,
+          'expected': [[[[1], [2, 3]], [[], [4]]],
+                       [[[5, 6], [7]], [[0, 8], []]]],
+          'expected_shape': [2, 2, 2, None],
+      },
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'lengths': ([2, 2], [2, 2, 2, 2], [1, 2, 0, 1, 2, 1, 2, 0]),
+          'expected': [[[[1], [2, 3]], [[], [4]]],
+                       [[[5, 6], [7]], [[0, 8], []]]],
+          'ragged_rank': 3,
+          'use_ragged_rank': False,  # lengths contains nested_row_lengths.
+          'expected_shape': [2, None, None, None],
+      },
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'lengths': [[2, 2], [2, 2, 2, 2], [1, 2, 0, 1, 2, 1, 2, 0]],
+          'expected': [[[[1], [2, 3]], [[], [4]]],
+                       [[[5, 6], [7]], [[0, 8], []]]],
+          'ragged_rank': 3,
+          'use_ragged_rank': False,  # lengths contains nested_row_lengths.
+          'expected_shape': [2, None, None, None],
+      },
+  )  # pyformat: disable
+  def disabled_testRaggedFromTensor(self,
+                                    tensor,
+                                    expected,
+                                    lengths=None,
+                                    padding=None,
+                                    ragged_rank=1,
+                                    use_ragged_rank=True,
+                                    expected_shape=None):
+    dt = constant_op.constant(tensor)
+    if use_ragged_rank:
+      rt = RaggedTensor.from_tensor(dt, lengths, padding, ragged_rank)
+    else:
+      rt = RaggedTensor.from_tensor(dt, lengths, padding)
+    self.assertEqual(type(rt), RaggedTensor)
+    self.assertEqual(rt.ragged_rank, ragged_rank)
+    self.assertTrue(
+        dt.shape.is_compatible_with(rt.shape),
+        '%s is incompatible with %s' % (dt.shape, rt.shape))
+    if expected_shape is not None:
+      self.assertEqual(rt.shape.as_list(), expected_shape)
+    self.assertAllEqual(rt, expected)
+    self.assertAllEqual(rt, RaggedTensor.from_nested_row_splits(
+        rt.flat_values, rt.nested_row_splits, validate=True))
+
+  def testHighDimensions(self):
+    # Use distinct prime numbers for all dimension shapes in this test, so
+    # we can see any errors that are caused by mixing up dimension sizes.
+    dt = array_ops.reshape(
+        math_ops.range(3 * 5 * 7 * 11 * 13 * 17), [3, 5, 7, 11, 13, 17])
+    for ragged_rank in range(1, 4):
+      rt = RaggedTensor.from_tensor(dt, ragged_rank=ragged_rank)
+      self.assertEqual(type(rt), RaggedTensor)
+      self.assertEqual(rt.ragged_rank, ragged_rank)
+      self.assertTrue(
+          dt.shape.is_compatible_with(rt.shape),
+          '%s is incompatible with %s' % (dt.shape, rt.shape))
+      self.assertAllEqual(rt, self.evaluate(dt).tolist())
+      self.assertAllEqual(rt, RaggedTensor.from_nested_row_splits(
+          rt.flat_values, rt.nested_row_splits, validate=True))
+
+  @parameterized.parameters(
+      # With no padding or lengths
+      {
+          'dt_shape': [0, 0],
+          'expected': []
+      },
+      {
+          'dt_shape': [0, 3],
+          'expected': []
+      },
+      {
+          'dt_shape': [3, 0],
+          'expected': [[], [], []]
+      },
+      {
+          'dt_shape': [0, 2, 3],
+          'expected': []
+      },
+      {
+          'dt_shape': [1, 0, 0],
+          'expected': [[]]
+      },
+      {
+          'dt_shape': [2, 0, 3],
+          'expected': [[], []]
+      },
+      {
+          'dt_shape': [2, 3, 0],
+          'expected': [[[], [], []], [[], [], []]]
+      },
+      {
+          'dt_shape': [2, 3, 0, 1],
+          'expected': [[[], [], []], [[], [], []]]
+      },
+      {
+          'dt_shape': [2, 3, 1, 0],
+          'expected': [[[[]], [[]], [[]]], [[[]], [[]], [[]]]]
+      },
+      # With padding
+      {
+          'dt_shape': [0, 0],
+          'padding': 0,
+          'expected': []
+      },
+      {
+          'dt_shape': [0, 3],
+          'padding': 0,
+          'expected': []
+      },
+      {
+          'dt_shape': [3, 0],
+          'padding': 0,
+          'expected': [[], [], []]
+      },
+      {
+          'dt_shape': [0, 2, 3],
+          'padding': [0, 0, 0],
+          'expected': []
+      },
+      {
+          'dt_shape': [2, 0, 3],
+          'padding': [0, 0, 0],
+          'expected': [[], []]
+      },
+      {
+          'dt_shape': [2, 3, 0],
+          'padding': [],
+          'expected': [[], []]
+      },
+      # With lengths
+      {
+          'dt_shape': [0, 0],
+          'lengths': [],
+          'expected': []
+      },
+      {
+          'dt_shape': [0, 3],
+          'lengths': [],
+          'expected': []
+      },
+      {
+          'dt_shape': [3, 0],
+          'lengths': [0, 0, 0],
+          'expected': [[], [], []]
+      },
+      {
+          'dt_shape': [3, 0],
+          'lengths': [2, 3, 4],  # lengths > ncols: truncated to ncols
+          'expected': [[], [], []]
+      },
+      {
+          'dt_shape': [0, 2, 3],
+          'lengths': [],
+          'expected': []
+      },
+      {
+          'dt_shape': [2, 0, 3],
+          'lengths': [0, 0],
+          'expected': [[], []]
+      },
+      {
+          'dt_shape': [2, 3, 0],
+          'lengths': [0, 0],
+          'expected': [[], []]
+      },
+  )
+  def testEmpty(self, dt_shape, expected, lengths=None, padding=None):
+    dt = array_ops.zeros(dt_shape)
+    for ragged_rank in range(1, len(dt_shape) - 1):
+      rt = RaggedTensor.from_tensor(dt, lengths, padding, ragged_rank)
+      self.assertEqual(type(rt), RaggedTensor)
+      self.assertEqual(rt.ragged_rank, ragged_rank)
+      self.assertTrue(dt.shape.is_compatible_with(rt.shape))
+      self.assertAllEqual(rt, expected)
+      self.assertAllEqual(rt, RaggedTensor.from_nested_row_splits(
+          rt.flat_values, rt.nested_row_splits, validate=True))
+
+  @parameterized.named_parameters([
+      {
+          'testcase_name': '2D_UnknownRank',
+          'tensor': [[1, 2], [3, 4]],
+          'tensor_shape': None,
+      },
+      {
+          'testcase_name': '2D_Shape_None_None',
+          'tensor': [[1, 2], [3, 4]],
+          'tensor_shape': [None, None],
+      },
+      {
+          'testcase_name': '2D_Shape_2_None',
+          'tensor': [[1, 2], [3, 4]],
+          'tensor_shape': [2, None],
+      },
+      {
+          'testcase_name': '2D_Shape_None_2',
+          'tensor': [[1, 2], [3, 4]],
+          'tensor_shape': [None, 2],
+      },
+      {
+          'testcase_name': '4D_UnknownRank',
+          'tensor': np.ones([4, 3, 2, 1]),
+          'tensor_shape': None,
+      },
+      {
+          'testcase_name': '4D_Shape_None_None_None_None',
+          'tensor': np.ones([4, 3, 2, 1]),
+          'tensor_shape': [None, None, None, None],
+      },
+      {
+          'tensor': np.ones([4, 3, 2, 1]),
+          'tensor_shape': [4, None, None, 1],
+          'testcase_name': '4D_Shape_4_None_None_1',
+      },
+  ])
+  def testPartialShapes(self, tensor, tensor_shape, shape=None,
+                        expected=None):
+    if expected is None:
+      expected = tensor
+
+    if context.executing_eagerly():
+      return  # static shapes are always fully defined in eager mode.
+
+    dt = constant_op.constant(tensor)
+    for ragged_rank in range(1, len(dt.shape) - 1):
+      dt_placeholder = array_ops.placeholder_with_default(tensor, tensor_shape)
+      rt = RaggedTensor.from_tensor(dt_placeholder, ragged_rank=ragged_rank)
+      self.assertIsInstance(rt, RaggedTensor)
+      self.assertEqual(rt.ragged_rank, ragged_rank)
+      self.assertTrue(
+          dt.shape.is_compatible_with(rt.shape),
+          '%s is incompatible with %s' % (dt.shape, rt.shape))
+      if shape is not None:
+        self.assertEqual(rt.shape.as_list(), shape)
+      self.assertAllEqual(rt, expected.tolist())
+      self.assertAllEqual(rt, RaggedTensor.from_nested_row_splits(
+          rt.flat_values, rt.nested_row_splits, validate=True))
+
+  @parameterized.parameters(
+      {
+          'tensor': [[1]],
+          'lengths': [0],
+          'padding':
+              0,
+          'error': (ValueError,
+                    'Specify argument `lengths` or `padding`, but not both.')
+      },
+      {
+          'tensor': [[1]],
+          'lengths': [0.5],
+          'error': (
+              TypeError,
+              r'Argument `tensor` \(name\: lengths\) must be of type integer.*')
+      },
+      {
+          'tensor': [[1, 2, 3]],
+          'lengths': [[1], [1]],
+          'error': (ValueError, r'Shape \(1, 3\) must have rank at least 3')
+      },
+      {
+          'tensor': [[1]],
+          'padding': 'a',
+          'error': (TypeError, '.*')
+      },
+      {
+          'tensor': [[1]],
+          'padding': [1],
+          'error': (ValueError, r'Shapes \(1,\) and \(\) are incompatible')
+      },
+      {
+          'tensor': [[[1]]],
+          'padding': 1,
+          'error': (ValueError, r'Shapes \(\) and \(1,\) are incompatible')
+      },
+      {
+          'tensor': [[1]],
+          'ragged_rank':
+              'bad',
+          'error': (TypeError,
+                    r'Argument `ragged_rank` must be an int. Received bad.')
+      },
+      {
+          'tensor': [[1]],
+          'ragged_rank':
+              0,
+          'error':
+              (ValueError,
+               r'Argument `ragged_rank` must be greater than 0. Received 0.')
+      },
+      {
+          'tensor': [[1]],
+          'ragged_rank':
+              -1,
+          'error':
+              (ValueError,
+               r'Argument `ragged_rank` must be greater than 0. Received -1.')
+      },
+      {
+          'tensor': [[[[1, 0], [2, 3]], [[0, 0], [4, 0]]],
+                     [[[5, 6], [7, 0]], [[0, 8], [0, 0]]]],
+          'lengths': ([2, 2], [2, 2, 2, 2]),
+          'ragged_rank':
+              3,
+          'error':
+              (ValueError,
+               r'If Argument `lengths` is a tuple of row_lengths, argument '
+               r'`ragged_rank` must be len\(lengths\): 2. Received '
+               r'ragged_rank: 3.')
+      },
+  )
+  def testErrors(self,
+                 tensor,
+                 lengths=None,
+                 padding=None,
+                 ragged_rank=1,
+                 error=None):
+    dt = constant_op.constant(tensor)
+    self.assertRaisesRegex(error[0], error[1], RaggedTensor.from_tensor, dt,
+                           lengths, padding, ragged_rank)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_functional_ops.py

@@ -0,0 +1,200 @@
+# 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 tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops.ragged import ragged_config
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("ragged.map_flat_values")
+@dispatch.add_dispatch_support
+def map_flat_values(op, *args, **kwargs):
+  """Applies `op` to the `flat_values` of one or more RaggedTensors.
+
+  Replaces any `RaggedTensor` in `args` or `kwargs` with its `flat_values`
+  tensor (which collapses all ragged dimensions), and then calls `op`.  Returns
+  a `RaggedTensor` that is constructed from the input `RaggedTensor`s'
+  `nested_row_splits` and the value returned by the `op`.
+
+  If the input arguments contain multiple `RaggedTensor`s, then they must have
+  identical `nested_row_splits`.
+
+  This operation is generally used to apply elementwise operations to each value
+  in a `RaggedTensor`.
+
+  Warning: `tf.ragged.map_flat_values` does *not* apply `op` to each row of a
+  ragged tensor.  This difference is important for non-elementwise operations,
+  such as `tf.reduce_sum`.  If you wish to apply a non-elementwise operation to
+  each row of a ragged tensor, use `tf.map_fn` instead.  (You may need to
+  specify an `output_signature` when using `tf.map_fn` with ragged tensors.)
+
+  Examples:
+
+  >>> rt = tf.ragged.constant([[1, 2, 3], [], [4, 5], [6]])
+  >>> tf.ragged.map_flat_values(tf.ones_like, rt)
+  <tf.RaggedTensor [[1, 1, 1], [], [1, 1], [1]]>
+  >>> tf.ragged.map_flat_values(tf.multiply, rt, rt)
+  <tf.RaggedTensor [[1, 4, 9], [], [16, 25], [36]]>
+  >>> tf.ragged.map_flat_values(tf.add, rt, 5)
+  <tf.RaggedTensor [[6, 7, 8], [], [9, 10], [11]]>
+
+  Example with a non-elementwise operation (note that `map_flat_values` and
+  `map_fn` return different results):
+
+  >>> rt = tf.ragged.constant([[1.0, 3.0], [], [3.0, 6.0, 3.0]])
+  >>> def normalized(x):
+  ...   return x / tf.reduce_sum(x)
+  >>> tf.ragged.map_flat_values(normalized, rt)
+  <tf.RaggedTensor [[0.0625, 0.1875], [], [0.1875, 0.375, 0.1875]]>
+  >>> tf.map_fn(normalized, rt)
+  <tf.RaggedTensor [[0.25, 0.75], [], [0.25, 0.5, 0.25]]>
+
+  Args:
+    op: The operation that should be applied to the RaggedTensor `flat_values`.
+      `op` is typically an element-wise operation (such as math_ops.add), but
+      any operation that preserves the size of the outermost dimension can be
+      used.  I.e., `shape[0]` of the value returned by `op` must match
+      `shape[0]` of the `RaggedTensor`s' `flat_values` tensors.
+    *args: Arguments for `op`.
+    **kwargs: Keyword arguments for `op`.
+
+  Returns:
+    A `RaggedTensor` whose `ragged_rank` matches the `ragged_rank` of all
+    input `RaggedTensor`s.
+  Raises:
+    ValueError: If args contains no `RaggedTensors`, or if the `nested_splits`
+      of the input `RaggedTensor`s are not identical.
+  """
+  # Replace RaggedTensors with their values; and collect the partitions tensors
+  # from each RaggedTensor.
+  partition_lists = []
+  flat_values_nrows = []
+  inner_args = _replace_ragged_with_flat_values(args, partition_lists,
+                                                flat_values_nrows)
+  inner_kwargs = _replace_ragged_with_flat_values(kwargs, partition_lists,
+                                                  flat_values_nrows)
+  if not partition_lists:
+    return op(*args, **kwargs)
+
+  # If we can statically determine that the inputs are incompatible, then raise
+  # an error.  (We can't guarantee full compatibility statically, so we need to
+  # perform some runtime checks too; but this allows us to fail sooner in some
+  # cases.)
+  if flat_values_nrows:
+    flat_values_nrows = set(flat_values_nrows)
+    if len(flat_values_nrows) != 1:
+      raise ValueError("Input RaggedTensors' flat_values must all have the "
+                       "same outer-dimension size.  Got sizes: %s" %
+                       flat_values_nrows)
+    flat_values_nrows = flat_values_nrows.pop()  # Get the single element
+  else:
+    flat_values_nrows = None
+
+  partition_dtypes = set(p[0].dtype for p in partition_lists)
+  if len(partition_dtypes) > 1:
+    if not ragged_config.auto_cast_partition_dtype():
+      raise ValueError("Input RaggedTensors have mismatched row partition "
+                       "dtypes; use RaggedTensor.with_row_splits_dtype() to "
+                       "convert them to compatible dtypes.")
+
+    partition_lists = [
+        [p.with_dtype(dtypes.int64)
+         for p in partition_list]  # pylint: disable=g-complex-comprehension
+        for partition_list in partition_lists
+    ]
+
+  # Delegate to `op`
+  op_output = op(*inner_args, **inner_kwargs)
+  # Check that the result has the expected shape (if known).
+  if flat_values_nrows is not None:
+    if not op_output.shape[:1].is_compatible_with([flat_values_nrows]):
+      raise ValueError(
+          "tf.ragged.map_flat_values requires that the output of `op` have "
+          "the same outer-dimension size as flat_values of any ragged "
+          "inputs. (output shape: %s; expected outer dimension size: %s)" %
+          (op_output.shape, flat_values_nrows))
+  # Compose the result from the transformed values and the partitions.
+  return ragged_tensor.RaggedTensor._from_nested_row_partitions(  # pylint: disable=protected-access
+      op_output,
+      _merge_partition_lists(partition_lists),
+      validate=False)
+
+
+def _replace_ragged_with_flat_values(value, partition_lists, flat_values_nrows):
+  """Replace RaggedTensors with their flat_values, and record their partitions.
+
+  Returns a copy of `value`, with any nested `RaggedTensor`s replaced by their
+  `flat_values` tensor.  Looks inside lists, tuples, and dicts.
+
+  Appends each `RaggedTensor`'s `RowPartition`s to `partition_lists`.
+
+  Args:
+    value: The value that should be transformed by replacing `RaggedTensors`.
+    partition_lists: An output parameter used to record the row partitions
+      for any `RaggedTensors` that were replaced.
+    flat_values_nrows: An output parameter used to record the outer dimension
+      size for each replacement `flat_values` (when known).  Contains a list of
+      int.
+
+  Returns:
+    A copy of `value` with nested `RaggedTensors` replaced by their `values`.
+  """
+  # Base case
+  if ragged_tensor.is_ragged(value):
+    value = ragged_tensor.convert_to_tensor_or_ragged_tensor(value)
+    partition_lists.append(value._nested_row_partitions)  # pylint: disable=protected-access
+    nrows = tensor_shape.dimension_at_index(value.flat_values.shape, 0).value
+    if nrows is not None:
+      flat_values_nrows.append(nrows)
+    return value.flat_values
+
+  # Recursion cases
+  def recurse(v):
+    return _replace_ragged_with_flat_values(v, partition_lists,
+                                            flat_values_nrows)
+
+  if isinstance(value, list):
+    return [recurse(v) for v in value]
+  elif isinstance(value, tuple):
+    return tuple(recurse(v) for v in value)
+  elif isinstance(value, dict):
+    return dict((k, recurse(v)) for (k, v) in value.items())
+  else:
+    return value
+
+
+def _merge_partition_lists(partition_lists):
+  """Merges the given list of lists of RowPartitions.
+
+  Args:
+    partition_lists: A list of lists of RowPartition.
+
+  Returns:
+    A list of RowPartitions, where `result[i]` is formed by merging
+    `partition_lists[j][i]` for all `j`, using
+    `RowPartition._merge_precomputed_encodings`.
+  """
+  dst = list(partition_lists[0])
+  for src in partition_lists[1:]:
+    if len(src) != len(dst):
+      raise ValueError("All ragged inputs must have the same ragged_rank.")
+    for i in range(len(dst)):
+      # pylint: disable=protected-access
+      dst[i] = dst[i]._merge_precomputed_encodings(src[i])
+  return dst

tensorflow/python/ops/ragged/ragged_gather_nd_op_test.py

@@ -0,0 +1,266 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_gather_ops.gather_nd."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_gather_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedGatherNdOpTest(test_util.TensorFlowTestCase,
+                           parameterized.TestCase):
+
+  DOCSTRING_PARAMS = [[['000', '001'], ['010']],
+                      [['100'], ['110', '111', '112'], ['120']],
+                      [[], ['210']]]  # pyformat: disable
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Docstring Examples
+      #=========================================================================
+      dict(
+          descr='Docstring example 1',
+          params=ragged_factory_ops.constant_value(DOCSTRING_PARAMS),
+          indices=[[2], [0]],
+          expected=ragged_factory_ops.constant_value(
+              [[[], [b'210']], [[b'000', b'001'], [b'010']]])),
+      dict(
+          descr='Docstring example 2',
+          params=ragged_factory_ops.constant_value(DOCSTRING_PARAMS),
+          indices=[[2, 1], [0, 0]],
+          expected=ragged_factory_ops.constant_value(
+              [[b'210'], [b'000', b'001']])),
+      dict(
+          descr='Docstring example 3',
+          params=ragged_factory_ops.constant_value(DOCSTRING_PARAMS),
+          indices=[[0, 0, 1], [1, 1, 2]],
+          expected=[b'001', b'112']),
+      #=========================================================================
+      # Indices with 0 values (selects the entire params)
+      #=========================================================================
+      dict(
+          descr='params: [B1, (B2)], indices: [0], result: [B1, (B2)]',
+          params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
+          indices=np.zeros([0], dtype=np.int32),
+          expected=ragged_factory_ops.constant_value(
+              [[b'a', b'b', b'c'], [b'd']])),
+      dict(
+          descr='params: [B1, (B2)], indices: [A1, 0], result: [A1, B1, (B2)]',
+          params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
+          indices=np.zeros([3, 0], dtype=np.int32),
+          expected=ragged_factory_ops.constant_value(
+              [[[b'a', b'b', b'c'], [b'd']],
+               [[b'a', b'b', b'c'], [b'd']],
+               [[b'a', b'b', b'c'], [b'd']]])),
+      dict(
+          descr=('params: [B1, (B2)], indices: [A1, A2, 0], '
+                 'result: [A1, A2, B1, (B2)]'),
+          params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
+          indices=np.zeros([1, 3, 0], dtype=np.int32),
+          expected=ragged_factory_ops.constant_value(
+              [[[[b'a', b'b', b'c'], [b'd']],
+                [[b'a', b'b', b'c'], [b'd']],
+                [[b'a', b'b', b'c'], [b'd']]]])),
+      dict(
+          descr='params: [B1], indices: [A1, (A2), 0], result: [A1, (A2), B1]',
+          params=['a'],
+          indices=ragged_factory_ops.constant_value(
+              [[[], []], [[]]],
+              ragged_rank=1,
+              dtype=np.int32),
+          expected=ragged_factory_ops.constant_value(
+              [[[b'a'], [b'a']], [[b'a']]],
+              ragged_rank=1)),
+      #=========================================================================
+      # Indices with 1 value (selects row from params)
+      #=========================================================================
+      dict(
+          descr='params: [B1, (B2)], indices: [A1, 1], result: [A1, (B2)]',
+          params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
+          indices=[[1], [0]],
+          expected=ragged_factory_ops.constant_value(
+              [[b'd'], [b'a', b'b', b'c']])),
+      dict(
+          descr=('params: [B1, (B2), (B3)], indices: [A1, 1], '
+                 'result: [A1, (B2), (B3)]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
+          indices=[[1], [1]],
+          expected=ragged_factory_ops.constant_value(
+              [[[b'e', b'f']], [[b'e', b'f']]])),
+      dict(
+          descr=('params: [B1, B2, B3], indices: [A1, (A2), 1], '
+                 'result: [A1, (A2), B2, B3]'),
+          params=[[['a']], [['b']]],
+          indices=ragged_factory_ops.constant_value([[[0]]], ragged_rank=1),
+          expected=ragged_factory_ops.constant_value(
+              [[[[b'a']]]], ragged_rank=1)),
+      #=========================================================================
+      # Indices with 2 values (selects row & col from params)
+      #=========================================================================
+      dict(
+          descr='params: [B1, (B2)], indices: [A1, 2], result: [A1]',
+          params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
+          indices=[[1, 0], [0, 0], [0, 2]],
+          expected=ragged_factory_ops.constant_value([b'd', b'a', b'c'])),
+      dict(
+          descr=('params: [B1, (B2), (B3)], indices: [A1, 2], '
+                 'result: [A1, (B3)]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
+          indices=[[1, 0], [0, 1], [0, 0]],
+          expected=ragged_factory_ops.constant_value(
+              [[b'e', b'f'], [b'd'], [b'a', b'b', b'c']])),
+      dict(
+          descr=('params: [B1, (B2), (B3)], indices: [A1, A2, 2], '
+                 'result: [A1, (A2), (B3)]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
+          indices=[[[1, 0], [0, 1], [0, 0]]],
+          expected=ragged_factory_ops.constant_value(
+              [[[b'e', b'f'], [b'd'], [b'a', b'b', b'c']]])),
+      dict(
+          descr=('params: [B1, (B2), B3], indices: [A1, A2, 2], '
+                 'result: [A1, A2, B3]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b'], ['c', 'd']],
+               [['e', 'f']]],
+              ragged_rank=1),
+          indices=[[[1, 0], [0, 1], [0, 0]]],
+          expected=[[[b'e', b'f'], [b'c', b'd'], [b'a', b'b']]]),
+      dict(
+          descr=('params: [B1, (B2), B3], indices: [A1, A2, A3, 2], '
+                 'result: [A1, A2, A3, B3]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b'], ['c', 'd']],
+               [['e', 'f']]],
+              ragged_rank=1),
+          indices=[[[[1, 0], [0, 1], [0, 0]]]],
+          expected=[[[[b'e', b'f'], [b'c', b'd'], [b'a', b'b']]]]),
+      dict(
+          descr=('params: [B1, (B2), (B3)], indices: [A1, (A2), 2], '
+                 'result: [A1, (A2), (B3)]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
+          indices=ragged_factory_ops.constant_value(
+              [[[1, 0], [0, 1]], [[0, 0]]],
+              ragged_rank=1),
+          expected=ragged_factory_ops.constant_value(
+              [[[b'e', b'f'], [b'd']], [[b'a', b'b', b'c']]])),
+      #=========================================================================
+      # Indices with 3 values
+      #=========================================================================
+      dict(
+          descr=('params: [B1, (B2), (B3)], indices: [A1, 3], '
+                 'result: [A1]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d']], [['e', 'f']]]),
+          indices=[[1, 0, 1], [0, 0, 0], [0, 1, 0]],
+          expected=[b'f', b'a', b'd']),
+      dict(
+          descr=('params: [B1, (B2), B3], indices: [A1, 3], '
+                 'result: [A1]'),
+          params=ragged_factory_ops.constant_value(
+              [[['a', 'b'], ['c', 'd']], [['e', 'f']]],
+              ragged_rank=1),
+          indices=[[1, 0, 1], [0, 0, 0], [0, 1, 1]],
+          expected=[b'f', b'a', b'd']),
+      dict(
+          descr=('params: [B1, (B2), (B3), B4], indices: [A1, 3], '
+                 'result: [A1, B4]'),
+          params=ragged_factory_ops.constant_value(
+              [[[['a', 'b'], ['c', 'd']], [['e', 'f']]]],
+              ragged_rank=2),
+          indices=[[0, 0, 1], [0, 0, 0], [0, 1, 0]],
+          expected=[[b'c', b'd'], [b'a', b'b'], [b'e', b'f']]),
+      dict(
+          descr=('Pass through bad_indices_policy for non ragged params+indices'),
+          params=[1, 3, 5, 7],
+          indices=[[3], [999], [1], [0]],
+          expected=[7, 0, 3, 1],
+          bad_indices_policy='IGNORE'),
+  ])  # pyformat: disable
+  def testRaggedGatherNd(
+      self, descr, params, indices, expected, bad_indices_policy=''
+  ):
+    result = ragged_gather_ops.gather_nd(
+        params, indices, bad_indices_policy=bad_indices_policy
+    )
+    self.assertAllEqual(result, expected)
+
+  def testRaggedGatherNdUnknownRankError(self):
+    if context.executing_eagerly():
+      return
+    params = ragged_factory_ops.constant([['a', 'b'], ['c', 'd']])
+    indices1 = array_ops.placeholder(dtypes.int32, shape=None)
+    indices2 = array_ops.placeholder(dtypes.int32, shape=[None])
+
+    with self.assertRaisesRegex(ValueError,
+                                'indices.rank be statically known.'):
+      ragged_gather_ops.gather_nd(params, indices1)
+    with self.assertRaisesRegex(
+        ValueError, r'indices.shape\[-1\] must be statically known.'):
+      ragged_gather_ops.gather_nd(params, indices2)
+
+  @parameterized.parameters([
+      dict(
+          params=['a'],
+          indices=0,
+          error=(ValueError, errors.InvalidArgumentError),
+      ),
+      dict(
+          params=ragged_factory_ops.constant_value([['a']]),
+          indices=0,
+          message='indices.rank must be at least 1.',
+      ),
+      dict(
+          params=['a', 'b', 'c'],
+          indices=ragged_factory_ops.constant_value([[0]]),
+          message='The innermost dimension of indices may not be ragged',
+      ),
+      dict(
+          params=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d']]),
+          indices=np.zeros([1, 3, 0], dtype=np.int32),
+          bad_indices_policy='IGNORE',
+          message=(
+              'non-default bad_indices_policy not supported for ragged gather'
+          ),
+      ),
+  ])
+  def testRaggedGatherNdStaticError(
+      self,
+      params,
+      indices,
+      bad_indices_policy='',
+      message=None,
+      error=ValueError,
+  ):
+    with self.assertRaisesRegex(error, message):
+      ragged_gather_ops.gather_nd(
+          params, indices, bad_indices_policy=bad_indices_policy
+      )
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_gather_op_test.py

@@ -0,0 +1,433 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.gather."""
+
+from absl.testing import parameterized
+
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_gather_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedGatherOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters([
+      # Basic gather (axis=0 and batch_dims=0)
+      dict(testcase_name='Params1DTensor_Indices1DTensor',
+           params=['a', 'b', 'c', 'd', 'e'],
+           indices=[2, 0, 2, 1],
+           expected=['c', 'a', 'c', 'b']),
+      dict(testcase_name='Params1DTensor_Indices2DRagged',
+           params=['a', 'b', 'c', 'd', 'e'],
+           indices=[[3, 1, 2], [1], [], [0]],
+           expected=[['d', 'b', 'c'], ['b'], [], ['a']]),
+      dict(testcase_name='Params2DRagged_Indices0DTensor',
+           params=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           indices=1,
+           expected=['c', 'd', 'e']),
+      dict(testcase_name='Params2DRagged_Indices1DTensor',
+           params=[['a', 'b', 'c'], ['d'], [], ['e']],
+           indices=[3, 1, 2, 1, 0],
+           expected=[
+               ['e'], ['d'], [], ['d'], ['a', 'b', 'c']]),
+      dict(testcase_name='Params2DRagged_Indices2DRagged',
+           params=[['a', 'b', 'c'], ['d'], [], ['e']],
+           indices=[[3, 1, 2], [1], [], [0]],
+           expected=[
+               [['e'], ['d'], []], [['d']], [], [['a', 'b', 'c']]]),
+      dict(testcase_name='Params3DRagged_Indices2DTensor',
+           params=[
+               [['a', 'b'], []], [['c', 'd'], ['e'], ['f']], [['g']]],
+           indices=[[1, 2], [0, 1], [2, 2]],
+           indices_ragged_rank=0,
+           expected=[
+               [[['c', 'd'], ['e'], ['f']], [['g']]],
+               [[['a', 'b'], []], [['c', 'd'], ['e'], ['f']]],
+               [[['g']], [['g']]]]),
+      dict(testcase_name='Params3DRagged_Indices3DTensor',
+           params=[[['a', 'b'], []],
+                   [['c', 'd'], ['e'], ['f']],
+                   [['g']]],
+           indices=[[[1, 2], [0, 1], [2, 2]], [[0, 0], [1, 2], [0, 1]]],
+           indices_ragged_rank=0,
+           expected=[
+               [[[['c', 'd'], ['e'], ['f']], [['g']]],
+                [[['a', 'b'], []], [['c', 'd'], ['e'], ['f']]],
+                [[['g']], [['g']]]],
+               [[[['a', 'b'], []], [['a', 'b'], []]],
+                [[['c', 'd'], ['e'], ['f']], [['g']]],
+                [[['a', 'b'], []], [['c', 'd'], ['e'], ['f']]]]]),
+      dict(testcase_name='Params1DTensor_Indices4DRaggedRank2',
+           params=['a', 'b', 'c', 'd', 'e', 'f', 'g'],
+           indices=[[[[3, 4], [0, 6]], []],
+                    [[[2, 1], [1, 0]], [[2, 5]], [[2, 3]]],
+                    [[[1, 0]]]],
+           indices_ragged_rank=2,
+           expected=[
+               [[['d', 'e'], ['a', 'g']], []],
+               [[['c', 'b'], ['b', 'a']], [['c', 'f']], [['c', 'd']]],
+               [[['b', 'a']]]]),
+      # Batch gather (batch_dims=1)
+      dict(testcase_name='Batch1D_Params2DRagged_Indices1DTensor',
+           params=[['a', 'b'], ['c'], ['d', 'e', 'f', 'g'], ['h']],
+           indices=[1, 0, 3, 0],
+           batch_dims=1,
+           expected=['b', 'c', 'g', 'h']),
+      dict(testcase_name='Batch1D_Params2DRagged_Indices2DTensor',
+           params=[['a', 'b'], ['c'], ['d', 'e', 'f', 'g'], ['h']],
+           indices=[[1, 0], [0, 0], [3, 1], [0, 0]],
+           indices_ragged_rank=0,
+           batch_dims=1,
+           expected=[['b', 'a'], ['c', 'c'], ['g', 'e'], ['h', 'h']]),
+      dict(testcase_name='Batch1D_Params2DRagged_Indices2DRagged',
+           params=[['a', 'b'], ['c'], ['d', 'e', 'f', 'g'], ['h']],
+           indices=[[1, 0], [], [3, 2, 1], [0]],
+           batch_dims=1,
+           expected=[['b', 'a'], [], ['g', 'f', 'e'], ['h']]),
+      dict(testcase_name='Batch1D_Params3DRagged_Indices3DRagged',
+           params=[[['a'], ['b', 'c']],
+                   [],
+                   [['d', 'e', 'f'], ['g'], ['h', 'i'], ['j']],
+                   [['k']]],
+           indices=[[[1, 0], []], [], [[3, 2, 1], [0]], [[0]]],
+           batch_dims=1,
+           expected=[[[['b', 'c'], ['a']], []],
+                     [],
+                     [[['j'], ['h', 'i'], ['g']], [['d', 'e', 'f']]],
+                     [[['k']]]]),
+      # Batch gather (batch_dims=2)
+      dict(testcase_name='Batch2D_Params3DRagged_Indices2DRagged',
+           params=[[['a', 'b', 'c'], ['d', 'e'], ['f']],
+                   [['g'], ['h', 'i']]],
+           indices=[[0, 1, 0], [0, 1]],
+           batch_dims=2,
+           expected=[['a', 'e', 'f'], ['g', 'i']]),
+      dict(testcase_name='Batch2D_Params3DRagged_Indices3DRagged',
+           params=[[['a', 'b', 'c'], ['d', 'e'], ['f']],
+                   [['g'], ['h', 'i']]],
+           indices=[[[2, 1, 0], [1, 1], [0]], [[0], []]],
+           batch_dims=2,
+           expected=[[['c', 'b', 'a'], ['e', 'e'], ['f']], [['g'], []]]),
+      # Batch gather (batch_dims=3)
+      dict(testcase_name='Batch3D_Params4DRagged_Indices3DRagged',
+           params=[[[['a', 'b', 'c'], ['d', 'e'], ['f']],
+                    [['g'], ['h', 'i']]], [[['j']]]],
+           indices=[[[0, 1, 0], [0, 1]], [[0]]],
+           batch_dims=3,
+           expected=[[['a', 'e', 'f'], ['g', 'i']], [['j']]]),
+      # Axis gather (axis=1)
+      dict(testcase_name='Params2DRagged_Indices0DTensor_axis_1',
+           params=[['a', 'b'], ['c', 'd', 'e'], ['f', 'g'], ['h', 'i', 'j'],
+                   ['k', 'l']],
+           indices=1,
+           axis=1,
+           expected=['b', 'd', 'g', 'i', 'l']),
+      dict(testcase_name='Params2DRagged_Indices1DTensor_axis_1',
+           params=[['a', 'b'], ['c', 'd', 'e'], ['f', 'g'], ['h', 'i', 'j'],
+                   ['k', 'l']],
+           indices=[1, 0],
+           axis=1,
+           expected=[['b', 'a'], ['d', 'c'], ['g', 'f'], ['i', 'h'],
+                     ['l', 'k']]),
+      dict(testcase_name='Params3DRagged_Indices0DTensor_axis_1',
+           params=[[['a', 'b'], ['c', 'd', 'e']],
+                   [['f', 'g'], ['h', 'i', 'j'], ['k', 'l']]],
+           indices=1,
+           axis=1,
+           expected=[['c', 'd', 'e'], ['h', 'i', 'j']]),
+      dict(testcase_name='Params3DRagged_Indices1DTensor_axis_1',
+           params=[[['a', 'b'], ['c', 'd', 'e']],
+                   [['f', 'g'], ['h', 'i', 'j'], ['k', 'l']]],
+           indices=[1, 0],
+           axis=1,
+           expected=[[['c', 'd', 'e'], ['a', 'b']],
+                     [['h', 'i', 'j'], ['f', 'g']]]),
+      # Batch/axis gather, batch = 1, axis > batch
+      dict(testcase_name='Params3DRagged_Indices1DTensor_batch_1_axis_2',
+           params=[[['a', 'b'], ['c', 'd', 'e']],
+                   [['f', 'g'], ['h', 'i', 'j'], ['k', 'l']]],
+           indices=[1, 0],
+           axis=2,
+           batch_dims=1,
+           expected=[['b', 'd'], ['f', 'h', 'k']]),
+      dict(testcase_name='Params4DRagged_Indices1DTensor_batch_1_axis_2',
+           params=[[[['a', 'b'], ['c', 'd', 'e']]],
+                   [[['f', 'g']], [['h', 'i', 'j'], ['k', 'l']]]],
+           indices=[0, 1],
+           axis=2,
+           batch_dims=1,
+           expected=[[['a', 'b']],
+                     [['h', 'i', 'j'], ['k', 'l']]]),
+  ])  # pyformat: disable
+  def testRaggedGather(self,
+                       params,
+                       indices,
+                       expected,
+                       axis=None,
+                       batch_dims=0,
+                       params_ragged_rank=None,
+                       indices_ragged_rank=None):
+    params = ragged_factory_ops.constant(params, ragged_rank=params_ragged_rank)
+    indices = ragged_factory_ops.constant(
+        indices, ragged_rank=indices_ragged_rank)
+    actual = ragged_gather_ops.gather(
+        params, indices, axis=axis, batch_dims=batch_dims)
+    self.assertAllEqual(actual, self._str_to_bytes(expected))
+
+  def _str_to_bytes(self, x):
+    if isinstance(x, list):
+      return [self._str_to_bytes(v) for v in x]
+    elif isinstance(x, str) and bytes is not str:
+      return bytes(x, 'utf-8')
+    else:
+      return x
+
+  def testOutOfBoundsError(self):
+    tensor_params = ['a', 'b', 'c']
+    tensor_indices = [0, 1, 2]
+    ragged_params = ragged_factory_ops.constant([['a', 'b'], ['c']])
+    ragged_indices = ragged_factory_ops.constant([[0, 3]])
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r'indices\[1\] = 3 is not in \[0, 3\)'):
+      self.evaluate(ragged_gather_ops.gather(tensor_params, ragged_indices))
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r'indices\[2\] = 2 is not in \[0, 2\)'):
+      self.evaluate(ragged_gather_ops.gather(ragged_params, tensor_indices))
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r'indices\[1\] = 3 is not in \[0, 2\)'):
+      self.evaluate(ragged_gather_ops.gather(ragged_params, ragged_indices))
+
+  def testUnknownIndicesRankError(self):
+    if context.executing_eagerly():
+      return
+    params = ragged_factory_ops.constant([], ragged_rank=1)
+    indices = constant_op.constant([0], dtype=dtypes.int64)
+    indices = array_ops.placeholder_with_default(indices, None)
+    self.assertRaisesRegex(ValueError,
+                           r'rank\(indices\) must be known statically',
+                           ragged_gather_ops.gather, params, indices)
+
+  # pylint: disable=bad-whitespace
+  @parameterized.parameters([
+      # params.shape=[2, None]; indices.shape=[3]
+      dict(
+          params        = [[1.0, 2.0], [3.0, 4.0, 5.0]],
+          indices       = [0, 0, 1],
+          expected_out  = [[1.0, 2.0], [1.0, 2.0], [3.0, 4.0, 5.0]],
+          out_grad      = [[0.1, 0.2], [0.3, 0.4], [0.5, 0.6, 0.7]],
+          expected_grad = [[0.4, 0.6], [0.5, 0.6, 0.7]]),
+      # params.shape=[2, None]; indices.shape=[0]
+      dict(
+          params        = [[1, 2], [3, 4, 5]],
+          indices       = [],
+          expected_out  = [],
+          out_grad      = [],
+          expected_grad = [[0, 0], [0, 0, 0]]),
+      # params.shape=[2, None]; indices.shape=[2, 2]
+      dict(
+          params        = [[1.0, 2.0], [3.0, 4.0, 5.0]],
+          indices       = [[0, 0], [1, 0]],
+          expected_out  = [[[1.0, 2.0], [1.0, 2.0]],
+                           [[3.0, 4.0, 5.0], [1.0, 2.0]]],
+          out_grad      = [[[0.1, 0.2], [0.3, 0.4]],
+                           [[0.5, 0.6, 0.7], [0.8, 0.9]]],
+          expected_grad = [[1.2, 1.5], [0.5, 0.6, 0.7]]),
+      # params.shape=[3, None, None]; indices.shape=[3]
+      dict(
+          params        = [[[1, 2], [3, 4, 5]], [[6.0]], [[7.0, 8.0]]],
+          indices       = [2, 1, 2],
+          expected_out  = [[[7.0, 8.0]], [[6.0]], [[7.0, 8.0]]],
+          out_grad      = [[[0.1, 0.2]], [[0.3]], [[0.4, 0.5]]],
+          expected_grad = [[[0, 0], [0, 0, 0]], [[0.3]], [[0.5, 0.7]]]),
+      # params.shape=[3, None, None]; indices.shape=[0]
+      dict(
+          params        = [[[1, 2], [3, 4, 5]], [[6.0]], [[7.0, 8.0]]],
+          indices       = [2, 1, 2],
+          expected_out  = [[[7.0, 8.0]], [[6.0]], [[7.0, 8.0]]],
+          out_grad      = [[[0.1, 0.2]], [[0.3]], [[0.4, 0.5]]],
+          expected_grad = [[[0, 0], [0, 0, 0]], [[0.3]], [[0.5, 0.7]]]),
+      # params.shape=[0, None]; indices.shape=[0]
+      dict(
+          params        = [],
+          indices       = [],
+          expected_out  = [],
+          out_grad      = [],
+          expected_grad = [],
+          params_ragged_rank = 1),
+      # params.shape=[2, None, 2]; indices.shape=[3]
+      dict(
+          params        = [[[1, 2], [3, 4]], [], [[5, 6]]],
+          indices       = [1, 1, 2, 0, 2],
+          expected_out  = [[], [], [[5, 6]], [[1, 2], [3, 4]], [[5, 6]]],
+          out_grad      = [[], [], [[1, 2]], [[3, 4], [5, 6]], [[7, 7]]],
+          expected_grad = [[[3, 4], [5, 6]], [], [[8, 9]]],
+          params_ragged_rank = 1),
+  ])  # pyformat: disable
+  @test_util.run_deprecated_v1
+  def testGradient(self,
+                   params,
+                   indices,
+                   expected_out,
+                   out_grad,
+                   expected_grad,
+                   params_ragged_rank=None):
+    """Tests that ragged_gather generates the right gradient.
+
+    Args:
+      params: The `params` that should be passed to `gather`.
+      indices: The `indices` that should be passed to `gather`.
+      expected_out: The expected value of `gather(params, indices)`.
+        `expected_out.shape = indices.shape + params.shape[1:]`.
+      out_grad: The value that should be fed in as the gradient for `out` when
+        testing the gradient of `ragged_gather`.  Must have the same shape as
+        `expected_out`.
+      expected_grad: The expected gradient for that should be returned for
+        `params`.  Must have the same shape as `params`.
+      params_ragged_rank: The ragged_rank of `params`.
+    """
+    if context.executing_eagerly():
+      return
+
+    params = ragged_factory_ops.constant(
+        params, dtype=dtypes.float32, ragged_rank=params_ragged_rank)
+    indices = constant_op.constant(indices, dtype=dtypes.int32)
+    out_ragged_rank = params.ragged_rank + indices.shape.ndims - 1
+    out_grad = ragged_factory_ops.constant(
+        out_grad, dtype=dtypes.float32, ragged_rank=out_ragged_rank)
+    expected_out = ragged_factory_ops.constant(
+        expected_out, dtype=dtypes.float32, ragged_rank=out_ragged_rank)
+    expected_grad = ragged_factory_ops.constant(
+        expected_grad,
+        dtype=dtypes.float32,
+        ragged_rank=params.ragged_rank)
+
+    out = ragged_gather_ops.gather(params, indices)
+    self.assertAllClose(out, expected_out)
+
+    grads = gradients_impl.gradients(
+        out.flat_values,
+        (params.nested_row_splits + (params.flat_values, indices,)),
+        out_grad.flat_values)
+    param_nested_splits_grads = grads[:-2]
+    params_flat_values_grad = grads[-2]
+    indices_grad = grads[-1]
+    self.assertEqual(indices_grad, None)
+    for splits_grad in param_nested_splits_grads:
+      self.assertEqual(splits_grad, None)
+
+    # The gradient generates an IndexedSlices; convert back to a normal Tensor.
+    self.assertIsInstance(params_flat_values_grad, indexed_slices.IndexedSlices)
+    params_flat_values_grad = ops.convert_to_tensor(params_flat_values_grad)
+
+    params_grad = params.with_flat_values(params_flat_values_grad)
+    self.assertAllClose(params_grad, expected_grad, atol=2e-6, rtol=2e-6)
+
+  @parameterized.parameters([
+      # Basic gather (batch_dims == 0, axis == 0)
+      dict(params_shape=[3, 4], indices_shape=[], axis=0),
+      dict(params_shape=[3, 4], indices_shape=[5], axis=0),
+      dict(params_shape=[3, 4], indices_shape=[2, 5], axis=0),
+      # Gather over axis (axis > 0)
+      dict(params_shape=[3, 4], indices_shape=[], axis=1),
+      dict(params_shape=[3, 4], indices_shape=[2], axis=1),
+      dict(params_shape=[3, 4], indices_shape=[2, 5], axis=1),
+      dict(params_shape=[7, 3, 1], indices_shape=[2, 4], axis=1),
+      dict(params_shape=[3, 4, 5, 6], indices_shape=[2, 1, 7], axis=1),
+      dict(params_shape=[7, 3, 5], indices_shape=[], axis=2),
+      dict(params_shape=[7, 3, 5], indices_shape=[2], axis=2),
+      dict(params_shape=[7, 3, 5], indices_shape=[4, 2], axis=2),
+      dict(params_shape=[7, 3, 5, 6], indices_shape=[4, 2], axis=2),
+      dict(params_shape=[7, 3, 5, 6], indices_shape=[], axis=3),
+      dict(params_shape=[7, 3, 5, 6], indices_shape=[4], axis=3),
+      dict(params_shape=[7, 3, 5, 6], indices_shape=[8, 4], axis=3),
+      dict(params_shape=[7, 3, 5, 6], indices_shape=[2, 3, 2, 3], axis=3),
+      # Batched gather (batch_dims > 0)
+      dict(params_shape=[7, 3], indices_shape=[7], batch_dims=1),
+      dict(params_shape=[7, 3], indices_shape=[7, 5], batch_dims=1),
+      dict(params_shape=[5, 3], indices_shape=[5, 7, 4, 2], batch_dims=1),
+      dict(params_shape=[2, 3, 6], indices_shape=[2], batch_dims=1),
+      dict(params_shape=[7, 3, 6], indices_shape=[7, 5, 4, 2], batch_dims=1),
+      dict(params_shape=[7, 3, 5], indices_shape=[7, 3], batch_dims=2),
+      dict(params_shape=[7, 3, 5], indices_shape=[7, 3, 2], batch_dims=2),
+      dict(params_shape=[7, 3, 5, 6], indices_shape=[7, 3, 5], batch_dims=3),
+      dict(params_shape=[2, 3, 5, 6], indices_shape=[2, 3, 5, 7], batch_dims=3),
+      # Batched gather with axis (axis > batch_dims > 0)
+      dict(params_shape=[2, 3, 6], indices_shape=[2], axis=2, batch_dims=1),
+      dict(params_shape=[2, 3, 6], indices_shape=[2, 4], axis=2, batch_dims=1),
+      dict(
+          params_shape=[3, 1, 6, 7], indices_shape=[3, 4], axis=3,
+          batch_dims=1),
+      dict(
+          params_shape=[3, 2, 6, 7], indices_shape=[3, 4], axis=3,
+          batch_dims=1),
+      dict(
+          params_shape=[2, 3, 6, 7], indices_shape=[2, 3], axis=3,
+          batch_dims=2),
+  ])
+  def testMatchesDenseGather(self,
+                             params_shape,
+                             indices_shape,
+                             axis=None,
+                             batch_dims=0):
+    # Build random params & indices matrics w/ the expected shapes.
+    if axis is None:
+      axis = batch_dims
+    params = np.random.randint(100, size=params_shape, dtype=np.int32)
+    indices = np.random.randint(
+        params_shape[axis], size=indices_shape, dtype=np.int32)
+
+    # Use array_ops.gather to get the expected value.
+    expected = array_ops.gather(
+        params, indices, axis=axis, batch_dims=batch_dims)
+
+    # Build ragged tensors with varying ragged_ranks from params & axis.
+    params_tensors = [params] + [
+        ragged_tensor.RaggedTensor.from_tensor(params, ragged_rank=i)
+        for i in range(1, len(params_shape))
+    ]
+    indices_tensors = [indices] + [
+        ragged_tensor.RaggedTensor.from_tensor(indices, ragged_rank=i)
+        for i in range(1, len(indices_shape))
+    ]
+
+    # For each combination of params & axis tensors, check that
+    # ragged_gather_ops.gather matches array_ops.gather.
+    for params_tensor in params_tensors:
+      for indices_tensor in indices_tensors:
+        actual = ragged_gather_ops.gather(
+            params_tensor, indices_tensor, axis=axis, batch_dims=batch_dims)
+        if isinstance(actual, ragged_tensor.RaggedTensor):
+          actual = actual.to_tensor()
+        self.assertAllEqual(
+            expected, actual, 'params.ragged_rank=%s, indices.ragged_rank=%s' %
+            (getattr(params_tensor, 'ragged_rank',
+                     0), getattr(indices_tensor, 'ragged_rank', 0)))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_gather_ops.py

@@ -0,0 +1,538 @@
+# 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.
+# ==============================================================================
+"""Gather operations for RaggedTensors."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import indexed_slices
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gen_ragged_array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import dispatch
+
+
+# ===============================================================================
+# ragged_gather
+# ===============================================================================
+@dispatch.dispatch_for_api(array_ops.gather_v2)
+def gather(params: ragged_tensor.RaggedOrDense,
+           indices: ragged_tensor.RaggedOrDense,
+           validate_indices=None,
+           axis=None,
+           batch_dims=0,
+           name=None):
+  """Gathers ragged slices from `params` axis `0` according to `indices`.
+
+  See `tf.gather` for full documentation.  (This version has the same API
+  as `tf.gather`, but supports ragged `params` and `indices`.)
+
+  Examples:
+
+  >>> params = tf.constant(['a', 'b', 'c', 'd', 'e'])
+  >>> indices = tf.constant([3, 1, 2, 1, 0])
+  >>> ragged_params = tf.ragged.constant([['a', 'b', 'c'], ['d'], [], ['e']])
+  >>> ragged_indices = tf.ragged.constant([[3, 1, 2], [1], [], [0]])
+
+  >>> tf.gather(params, ragged_indices)
+  <tf.RaggedTensor [[b'd', b'b', b'c'], [b'b'], [], [b'a']]>
+
+  >>> tf.gather(ragged_params, indices)
+  <tf.RaggedTensor [[b'e'], [b'd'], [], [b'd'], [b'a', b'b', b'c']]>
+
+  >>> tf.gather(ragged_params, ragged_indices)
+  <tf.RaggedTensor [[[b'e'], [b'd'], []], [[b'd']], [], [[b'a', b'b', b'c']]]>
+
+  Args:
+    params: The potentially ragged tensor from which to gather values. Must be
+      at least rank 1.
+    indices: The potentially ragged tensor indicating which values to gather.
+      Must have dtype `int32` or `int64`.  Values must be in the range `[0,
+      params.shape[0]]`.
+    validate_indices: Ignored.
+    axis: The axis in `params` to gather `indices` from.
+    batch_dims: The number of batch dimensions.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `RaggedTensor`, where `output.dtype=params.dtype` and
+    `output.shape=indices.shape + params.shape[1:]` and
+    `output.ragged_rank=indices.shape.ndims + params.ragged_rank`.
+
+  Raises:
+    ValueError: If indices.shape.ndims is not known statically.
+  """
+  del validate_indices
+
+  with ops.name_scope(name, 'RaggedGather', [params, indices]):
+    params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        params, name='params')
+    indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        indices, name='indices')
+    params, indices = ragged_tensor.match_row_splits_dtypes(params, indices)
+
+    if batch_dims != indices.shape.rank:
+      batch_dims = array_ops.get_positive_axis(
+          batch_dims,
+          indices.shape.rank,
+          axis_name='batch_dims',
+          ndims_name='rank(indices)')
+    if params.shape.rank is not None and batch_dims >= params.shape.rank:
+      raise ValueError('batch_dims must be less than rank(params)')
+    if axis is None:
+      axis = batch_dims
+    axis = array_ops.get_positive_axis(
+        axis, params.shape.rank, ndims_name='rank(params)')
+    if axis < batch_dims:
+      raise ValueError('axis must be greater than or equal to batch_dims')
+    if indices.shape.rank is not None:
+      if not 0 <= batch_dims <= indices.shape.rank:
+        raise ValueError(
+            'batch_dims=%s must be between 0 and rank(indices)=%s' %
+            (batch_dims, indices.shape.rank))
+
+    return _gather(params, indices, axis, batch_dims)
+
+
+def _gather(params, indices, axis, batch_dims):
+  """Helper that implements the body for ragged gather().
+
+  Assumes that `params` and `indices` have been converted to tensors or
+  ragged tensors, and that `axis` and `batch_dims` have been normalized to
+  be positive.  (So these conversions & normalizations can be skipped in
+  recursive calls to _gather).
+
+  Args:
+    params: The tensor from which to gather values.
+    indices: The indices of values to gather.
+    axis: The axis in `params` to gather `indices` from.
+    batch_dims: The number of batch dimensions.
+
+  Returns:
+    A potentially ragged tensor.
+  """
+  params_is_ragged = ragged_tensor.is_ragged(params)
+  indices_is_ragged = ragged_tensor.is_ragged(indices)
+
+  if not (params_is_ragged or indices_is_ragged):
+    return array_ops.gather(params, indices, axis=axis, batch_dims=batch_dims)
+
+  if batch_dims > 0:
+    return _batch_gather(params, indices, axis, batch_dims)
+
+  if axis > 0:
+    return _axis_gather(params, indices, axis)
+
+  if indices_is_ragged:
+    return indices.with_values(_gather(params, indices.values, 0, 0))
+
+  if indices.shape.ndims is None:
+    raise ValueError('rank(indices) must be known statically')
+
+  out_ragged_rank = indices.shape.ndims + len(params.nested_row_splits) - 1
+  result = gen_ragged_array_ops.ragged_gather(
+      indices=indices,
+      params_dense_values=params.flat_values,
+      params_nested_splits=params.nested_row_splits,
+      OUTPUT_RAGGED_RANK=out_ragged_rank)
+
+  result = ragged_tensor.RaggedTensor.from_nested_row_splits(
+      result.output_dense_values, result.output_nested_splits, validate=False)
+
+  # Inject uniform_row_lengths into the result RaggedTensors for dimensions
+  # corresponding to dense outer dimensions of `indices`.
+  # TODO(edloper): Change this to construct the result using RowPartition
+  # objects instead, so we don't need to modify private variables.
+  if indices.shape.ndims > 1:
+    target = result
+    indices_shape = array_ops.shape(indices, out_type=params.row_splits.dtype)
+    shape_cumprod = math_ops.cumprod(indices_shape)
+    for dim in range(indices.shape.ndims - 1):
+      # pylint: disable=protected-access
+      target._cached_nrows = shape_cumprod[dim]
+      target._uniform_row_length = indices_shape[dim + 1]
+      target = target.values
+
+  return result
+
+
+def _batch_gather(params, indices, axis, batch_dims):
+  """Helper that implements the body for ragged gather() when batch_dims>0.
+
+  Args:
+    params: The tensor from which to gather values.
+    indices: The indices of values to gather.
+    axis: The axis in `params` to gather `indices` from.
+    batch_dims: The number of batch dimensions.
+
+  Returns:
+    A potentially ragged tensor.
+  """
+  # Perform static checks that `params` and `indices` have compatible batch
+  # dimensions.  Note: we do not perform *runtime* checks that `params` and
+  # `indices` actually have the same row-splits (because we wish to avoid the
+  # runtime cost of those checks).  If `params` and `indices` are
+  # incompatible, the resulting `RaggedTensor` may be nonsensical.
+  if not params.shape[:batch_dims].is_compatible_with(
+      indices.shape[:batch_dims]):
+    raise ValueError('batch shape from indices %s does not match params '
+                     'shape %s' % (indices.shape[:batch_dims], params.shape))
+
+  if batch_dims > 1:
+    # Convert params & indices to ragged tensors.
+    if not isinstance(params, ragged_tensor.RaggedTensor):
+      if indices.uniform_row_length is None:
+        raise ValueError(
+            'batch shape from indices does not match params shape: ragged '
+            'indices dimension corresponds to uniform params dimension')
+      params = ragged_tensor.RaggedTensor.from_tensor(
+          params, ragged_rank=1, row_splits_dtype=indices.row_splits.dtype)
+    if not isinstance(indices, ragged_tensor.RaggedTensor):
+      if params.uniform_row_length is None:
+        raise ValueError(
+            'batch shape from indices does not match params shape: ragged '
+            'params dimension corresponds to uniform indices dimension')
+      indices = ragged_tensor.RaggedTensor.from_tensor(
+          indices, ragged_rank=1, row_splits_dtype=params.row_splits.dtype)
+    # Flatten the two outer batch dimensions into a single batch dimension,
+    # and recurse.
+    return params.with_values(
+        _gather(params.values, indices.values, axis - 1, batch_dims - 1))
+
+  if axis > 1:
+    # Convert an axis dimension into a batch dimension, by adding a dimension
+    # to `indices`, and tiling it to match `params`.  E.g., if `params`
+    # had shape `[B, P1, P2]`, and `indices` had shape `[B, I1, I2]`, then we
+    # tile `indices` to have shape `[B, P1, I1, I2]`.  That way, we can treat
+    # the `P1` dimension as a batch dimension.
+    if not isinstance(indices, ragged_tensor.RaggedTensor):
+      adjusted_indices = params.with_values(
+          array_ops.repeat(indices, params.row_lengths(), 0))
+    else:
+      if not isinstance(params, ragged_tensor.RaggedTensor):
+        params = ragged_tensor.RaggedTensor.from_tensor(
+            params, ragged_rank=1, row_splits_dtype=indices.row_splits.dtype)
+      adjusted_indices = _gather(
+          indices,
+          params.with_values(
+              array_ops.repeat(
+                  math_ops.range(params.nrows()), params.row_lengths())), 0, 0)
+    return _batch_gather(params, adjusted_indices, axis, batch_dims + 1)
+
+  if indices.shape.rank is None:
+    raise ValueError('rank(indices) must be known statically')
+
+  assert batch_dims == 1
+  # If params.shape=[B, P1...PN] and indices.shape=[B, I1...IM], then:
+  #
+  #     output[b,        i1...im,      p2...pn] =
+  #     params[b, indices[b, i1...im], p2...pn]
+  #
+  # We construct `output` by flattening `params`, adjusting the `indices` to
+  # point into that flattened list, and recursively calling `gather`.
+  flat_params = _flatten_dims_0_and_1(params)
+  adjustments = _row_starts(params, indices.dtype)  # offset for each batch
+  # increase adjustments's rank so it broadcasts w/ the outer dim of indices
+  adjustments = _increase_rank_to(adjustments, indices.shape.ndims)
+  adjusted_indices = indices + adjustments
+  return _gather(flat_params, adjusted_indices, axis - 1, 0)
+
+
+def _axis_gather(params, indices, axis):
+  """Helper that implements ragged gather when axis>0 and batch_dims==0.
+
+  Args:
+    params: The tensor from which to gather values.
+    indices: The indices of values to gather.
+    axis: The axis in `params` to gather `indices` from.
+
+  Returns:
+    A potentially ragged tensor.
+  """
+  if axis > 1:
+    if not isinstance(params, ragged_tensor.RaggedTensor):
+      params = ragged_tensor.RaggedTensor.from_tensor(
+          params, ragged_rank=1, row_splits_dtype=indices.row_splits.dtype)
+    # Recurse, using the flattened params (but do not flatten indices).
+    return params.with_values(_gather(params.values, indices, axis - 1, 0))
+
+  if indices.shape.rank is None:
+    raise ValueError('rank(indices) must be known statically')
+
+  # Note: there is no checking of indices. If there is some index
+  # out of bounds, the results may be nonsensical.
+
+  assert axis == 1
+  # If params.shape=[P1...PN] and indices.shape=[I1...IM], then:
+  #
+  #     output[p1,      i1...im,     p3...pn] =
+  #     params[p1, indices[i1...im], p3...pn]
+  #
+  # We construct `output` by flattening `params`, adjusting the `indices` to
+  # have one additional dimension, and to point into that flattened list, and
+  # recursively calling `gather`.
+  flat_params = _flatten_dims_0_and_1(params)
+  adjustments = _row_starts(params, indices.dtype)  # offset for each batch
+  adjustments = _increase_rank_to(adjustments, indices.shape.ndims + 1)
+  adjusted_indices = indices + adjustments
+  return _gather(flat_params, adjusted_indices, axis - 1, 0)
+
+
+def _flatten_dims_0_and_1(t):
+  """Returns a copy of `t` with the outer two dimensions merged."""
+  if isinstance(t, ragged_tensor.RaggedTensor):
+    return t.values
+  else:
+    t_shape = array_ops.shape(t)
+    return array_ops.reshape(t, array_ops.concat([[-1], t_shape[2:]], axis=0))
+
+
+def _row_starts(t, dtype):
+  """Returns the start indices for the rows in `t`."""
+  if isinstance(t, ragged_tensor.RaggedTensor):
+    return math_ops.cast(t.row_starts(), dtype)
+  else:
+    t_shape = array_ops.shape(t, out_type=dtype)
+    return math_ops.range(t_shape[0]) * t_shape[1]
+
+
+def _increase_rank_to(t, rank):
+  """Adds *trailing* size-1 dimensions to `t` until it has the given rank."""
+  if isinstance(t, ragged_tensor.RaggedTensor):
+    return t.with_values(_increase_rank_to(t, rank - 1))
+  else:
+    old_dims = array_ops.shape(t)
+    new_dims = array_ops.ones([rank - array_ops.rank(t)], old_dims.dtype)
+    new_shape = array_ops.concat([old_dims, new_dims], axis=0)
+    return array_ops.reshape(t, new_shape)
+
+
+@dispatch.dispatch_for_api(array_ops.gather)
+def _ragged_gather_v1(params: ragged_tensor.RaggedOrDense,
+                      indices: ragged_tensor.RaggedOrDense,
+                      validate_indices=None,
+                      name=None,
+                      axis=0,
+                      batch_dims=0):
+  return gather(params, indices, validate_indices, axis, batch_dims, name)
+
+
+# ===============================================================================
+# ragged.gather_nd
+# ===============================================================================
+@dispatch.dispatch_for_api(array_ops.gather_nd_v2)
+def gather_nd(
+    params: ragged_tensor.RaggedOrDense,
+    indices: ragged_tensor.RaggedOrDense,
+    batch_dims=0,
+    name=None,
+    bad_indices_policy='',
+):
+  """Gather slices from `params` using `n`-dimensional indices.
+
+  This operation is similar to `gather`, but it uses the innermost dimension
+  of `indices` to define a slice into `params`.  In particular, if:
+
+  * `indices` has shape `[A1...AN, I]`
+  * `params` has shape `[B1...BM]`
+
+  Then:
+
+  * `result` has shape `[A1...AN, B_{I+1}...BM]`.
+  * `result[a1...aN] = params[indices[a1...aN, :]]`
+
+  Args:
+    params: A potentially ragged tensor with shape `[A1...AN, I]`.
+    indices: A potentially ragged tensor with shape `[B1...BM]`.
+    batch_dims: Must be zero.
+    name: A name for the operation (optional).
+    bad_indices_policy: A string. If `""` or `"DEFAULT"`, the default behavior
+      is used (error on CPU and ignore on GPU). If `"IGNORE"`, the bad indices
+      are ignored and 0 is stored in the
+
+  Returns:
+    A potentially ragged tensor with shape `[A1...AN, B_{I+1}...BM]`.
+
+  #### Examples:
+
+  >>> params = tf.ragged.constant(
+  ...     [ [ ['000', '001'], ['010'              ]          ],
+  ...       [ ['100'       ], ['110', '111', '112'], ['120'] ],
+  ...       [ [            ], ['210'              ]          ] ])
+
+  >>> # Gather 2D slices from a 3D tensor
+  >>> tf.gather_nd(params, [[2], [0]])
+  <tf.RaggedTensor [[[], [b'210']], [[b'000', b'001'], [b'010']]]>
+
+  >>> # Gather 1D slices from a 3D tensor
+  >>> tf.gather_nd(params, [[2, 1], [0, 0]])
+  <tf.RaggedTensor [[b'210'], [b'000', b'001']]>
+
+  >>> # Gather scalars from a 3D tensor
+  >>> tf.gather_nd(params, [[0, 0, 1], [1, 1, 2]]).numpy()
+  array([b'001', b'112'], dtype=object)
+  """
+  if not isinstance(batch_dims, int) or batch_dims != 0:
+    raise ValueError('batch_dims != 0 is not supported for ragged gather yet.')
+  if not (ragged_tensor.is_ragged(params) or ragged_tensor.is_ragged(indices)):
+    return array_ops.gather_nd(
+        params, indices, name=name, bad_indices_policy=bad_indices_policy
+    )
+
+  if bad_indices_policy not in ('', 'DEFAULT'):
+    raise ValueError(
+        'non-default bad_indices_policy not supported for ragged gather'
+    )
+
+  with ops.name_scope(name, 'RaggedGatherNd', [params, indices]):
+
+    params = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        params, name='params')
+    indices = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        indices, name='indices')
+    params, indices = ragged_tensor.match_row_splits_dtypes(params, indices)
+    indices_shape = indices.shape
+    indices_ndims = indices_shape.ndims
+    if indices_ndims is None:
+      raise ValueError('indices.rank be statically known.')
+    if indices_ndims == 0:
+      raise ValueError('indices.rank must be at least 1.')
+    if (ragged_tensor.is_ragged(indices) and
+        indices_ndims == indices.ragged_rank + 1):
+      raise ValueError('The innermost dimension of indices may not be ragged')
+
+    # `index_size` is the "n" in "gather_nd" -- i.e., the number of dimensions
+    # that each index slices into.
+    index_size = tensor_shape.dimension_value(indices_shape[-1])
+    if index_size is None:
+      raise ValueError('indices.shape[-1] must be statically known.')
+
+    # If `indices` has more than 2 dimensions, then recurse.  If `indices` is
+    # dense, then we convert it to ragged before recursing, and then convert
+    # the result back to `dense` if appropriate.
+    if indices_ndims > 2:
+      indices_is_dense = not ragged_tensor.is_ragged(indices)
+      if indices_is_dense:
+        indices = ragged_tensor.RaggedTensor.from_tensor(
+            indices, ragged_rank=indices_ndims - 2,
+            row_splits_dtype=params.row_splits.dtype)
+      result = indices.with_flat_values(gather_nd(params, indices.flat_values))
+      if (indices_is_dense and ragged_tensor.is_ragged(result) and
+          result.ragged_rank == indices_ndims - 2):
+        result = ragged_tensor.RaggedTensor.to_tensor(result)
+      return result
+
+    # indices_ndims <= 2, and the innermost dimension of indices may not be
+    # ragged, so `indices` must not be ragged.
+    assert not ragged_tensor.is_ragged(indices)
+    assert ragged_tensor.is_ragged(params)
+
+    # Handle corner case: An empty index tuple selects the entire `params`
+    # value.  So if `index_size` is zero, then tile `params`.
+    if index_size == 0:
+      params_ndims = params.ragged_rank + array_ops.rank(params.flat_values)
+      for dim in range(indices_ndims - 1):
+        params = ragged_array_ops.expand_dims(params, axis=0)
+      multiples = array_ops.concat([
+          array_ops.shape(indices)[:-1],
+          array_ops.ones([params_ndims], dtypes.int32)
+      ],
+                                   axis=0)
+      return ragged_array_ops.tile(params, multiples)
+
+    # When index_size=1, we can just flatten the index tuples and use gather.
+    elif index_size == 1:
+      flattened_index_tuples = array_ops.reshape(indices, [-1])
+      return gather(params, flattened_index_tuples)
+
+    # Otherwise, params is a RaggedTensor, and indices is a 1D or 2D Tensor.
+    # Flatten both the index tuples and the params, such that the flattened
+    # index tuples point to the correct values in the flattened params; and
+    # then use ragged.gather on the flattened index tuples & params.
+    else:
+      indices = math_ops.cast(indices, params.row_splits.dtype)
+
+      # Flatten the outermost 2 dimensions of the index tuples & params.
+      flattened_index_tuples = array_ops.gather(params.row_splits,
+                                                indices[..., 0])
+      flattened_index_tuples += indices[..., 1]
+      flattened_params = params.values
+
+      # Flatten any remaining dimensions.
+      for dim in range(2, index_size):
+        if not ragged_tensor.is_ragged(flattened_params):
+          flattened_index_tuples = array_ops.expand_dims(
+              flattened_index_tuples, axis=1)
+          flattened_index_tuples = array_ops.concat(
+              [flattened_index_tuples, indices[..., dim:]], axis=1)
+          return array_ops.gather_nd(flattened_params, flattened_index_tuples)
+
+        flattened_index_tuples = array_ops.gather(
+            flattened_params.row_starts(), flattened_index_tuples)
+        flattened_index_tuples += indices[..., dim]
+        flattened_params = flattened_params.values
+
+      # Gather using the flattened index tuples and params.
+      return gather(flattened_params, flattened_index_tuples)
+
+
+@dispatch.dispatch_for_api(array_ops.gather_nd)
+def _ragged_gather_nd_v1(
+    params: ragged_tensor.RaggedOrDense,
+    indices: ragged_tensor.RaggedOrDense,
+    name=None,
+    batch_dims=0,
+    bad_indices_policy='',
+):
+  return gather_nd(
+      params, indices, batch_dims, name, bad_indices_policy=bad_indices_policy
+  )
+
+
+# ===============================================================================
+# Gradient for the RaggedGather kernel
+# ===============================================================================
+@ops.RegisterGradient('RaggedGather')
+def _ragged_gather_grad(op, *grads):
+  """Gradient for RaggedGather op."""
+  param_nested_splits = op.inputs[:-2]
+  param_inner_values = op.inputs[-2]
+  indices = op.inputs[-1]
+  grad_inner_values = grads[-1]
+
+  # For each row in `params`, find the range of values in `params.inner_values`
+  # that is covered by that row.  In particular, the values in row `i` are
+  # `param_inner_values[combined_splits[i]:combined_splits[i+1]`.
+  combined_splits = param_nested_splits[0]
+  for row_splits in param_nested_splits[1:]:
+    combined_splits = array_ops.gather(row_splits, combined_splits)
+
+  # The outer dimensions of `indices` correspond 1:1 with the outer dimensions
+  # of `ragged_grad` that are encoded by `grad_nested_splits`.  Thus, the
+  # flattened `indices` correspond 1:1 with `grad_inner_values`.
+  flat_indices = array_ops.reshape(indices, [-1])
+
+  # Build an IndexedSlices where the values are taken from `flat_grad`.
+  grad_indices = ragged_math_ops.range(
+      array_ops.gather(combined_splits, flat_indices),
+      array_ops.gather(combined_splits[1:], flat_indices)).values
+
+  param_inner_values_grad = indexed_slices.IndexedSlices(
+      values=grad_inner_values, indices=grad_indices,
+      dense_shape=array_ops.shape(param_inner_values))
+  return [None for _ in param_nested_splits] + [param_inner_values_grad, None]

tensorflow/python/ops/ragged/ragged_getitem.py

@@ -0,0 +1,477 @@
+# 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.
+# ==============================================================================
+"""Python-style indexing and slicing for RaggedTensors."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_shape
+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 cond
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_gather_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("__operators__.ragged_getitem", v1=[])
+@dispatch.add_dispatch_support
+def ragged_tensor_getitem(rt_input, key):
+  """Returns the specified piece of this RaggedTensor.
+
+  Supports multidimensional indexing and slicing, with one restriction:
+  indexing into a ragged inner dimension is not allowed.  This case is
+  problematic because the indicated value may exist in some rows but not
+  others.  In such cases, it's not obvious whether we should (1) report an
+  IndexError; (2) use a default value; or (3) skip that value and return a
+  tensor with fewer rows than we started with.  Following the guiding
+  principles of Python ("In the face of ambiguity, refuse the temptation to
+  guess"), we simply disallow this operation.
+
+  Args:
+    rt_input: The RaggedTensor to slice.
+    key: Indicates which piece of the RaggedTensor to return, using standard
+      Python semantics (e.g., negative values index from the end).  `key`
+      may have any of the following types:
+
+      * `int` constant
+      * Scalar integer `Tensor`
+      * `slice` containing integer constants and/or scalar integer
+        `Tensor`s
+      * `Ellipsis`
+      * `tf.newaxis`
+      * `tuple` containing any of the above (for multidimensional indexing)
+
+  Returns:
+    A `Tensor` or `RaggedTensor` object.  Values that include at least one
+    ragged dimension are returned as `RaggedTensor`.  Values that include no
+    ragged dimensions are returned as `Tensor`.  See above for examples of
+    expressions that return `Tensor`s vs `RaggedTensor`s.
+
+  Raises:
+    ValueError: If `key` is out of bounds.
+    ValueError: If `key` is not supported.
+    TypeError: If the indices in `key` have an unsupported type.
+
+  Examples:
+
+  >>> # A 2-D ragged tensor with 1 ragged dimension.
+  >>> rt = tf.ragged.constant([['a', 'b', 'c'], ['d', 'e'], ['f'], ['g']])
+  >>> rt[0].numpy()                 # First row (1-D `Tensor`)
+  array([b'a', b'b', b'c'], dtype=object)
+  >>> rt[:3].to_list()              # First three rows (2-D RaggedTensor)
+  [[b'a', b'b', b'c'], [b'd', b'e'], [b'f']]
+  >>> rt[3, 0].numpy()              # 1st element of 4th row (scalar)
+  b'g'
+
+  >>> # A 3-D ragged tensor with 2 ragged dimensions.
+  >>> rt = tf.ragged.constant([[[1, 2, 3], [4]],
+  ...                          [[5], [], [6]],
+  ...                          [[7]],
+  ...                          [[8, 9], [10]]])
+  >>> rt[1].to_list()               # Second row (2-D RaggedTensor)
+  [[5], [], [6]]
+  >>> rt[3, 0].numpy()              # First element of fourth row (1-D Tensor)
+  array([8, 9], dtype=int32)
+  >>> rt[:, 1:3].to_list()          # Items 1-3 of each row (3-D RaggedTensor)
+  [[[4]], [[], [6]], [], [[10]]]
+  >>> rt[:, -1:].to_list()          # Last item of each row (3-D RaggedTensor)
+  [[[4]], [[6]], [[7]], [[10]]]
+  """
+  if not isinstance(rt_input, ragged_tensor.RaggedTensor):
+    raise TypeError("Ragged __getitem__ expects a ragged_tensor.")
+  scope_tensors = [rt_input] + list(_tensors_in_key_list(key))
+  if isinstance(key, (list, tuple)):
+    key = list(key)
+  else:
+    key = [key]
+  with ops.name_scope(None, "RaggedGetItem", scope_tensors):
+    return _ragged_getitem(rt_input, key)
+
+
+def _ragged_getitem(rt_input, key_list):
+  """Helper for indexing and slicing ragged tensors with __getitem__().
+
+  Extracts the specified piece of the `rt_input`.  See
+  `RaggedTensor.__getitem__` for examples and restrictions.
+
+  Args:
+    rt_input: The `RaggedTensor` from which a piece should be returned.
+    key_list: The list of keys specifying which piece to return. Each key
+      corresponds with a separate dimension.
+
+  Returns:
+    The indicated piece of rt_input.
+
+  Raises:
+    ValueError: If `key_list` is not supported.
+    TypeError: If any keys in `key_list` have an unsupported type.
+  """
+  if not key_list:
+    return rt_input
+  row_key = key_list[0]
+  inner_keys = key_list[1:]
+
+  if row_key is Ellipsis:
+    expanded_key_list = _expand_ellipsis(key_list, rt_input.shape.ndims)
+    return _ragged_getitem(rt_input, expanded_key_list)
+
+  # Adding a new axis: Get rt_input[inner_keys], and wrap it in a RaggedTensor
+  # that puts all values in a single row.
+  if row_key is array_ops.newaxis:
+    inner_rt = _ragged_getitem(rt_input, inner_keys)
+    nsplits = tensor_shape.dimension_at_index(inner_rt.row_splits.shape, 0)
+    if nsplits.value is not None:
+      nsplits = nsplits.value
+    else:
+      nsplits = array_ops.shape(inner_rt.row_splits,
+                                out_type=inner_rt.row_splits.dtype)[0]
+    return ragged_tensor.RaggedTensor.from_uniform_row_length(
+        inner_rt, nsplits - 1, nrows=1, validate=False)
+
+  # Slicing a range of rows: first slice the outer dimension, and then
+  # call `_ragged_getitem_inner_dimensions` to handle the inner keys.
+  if isinstance(row_key, slice):
+    sliced_rt_input = _slice_ragged_row_dimension(rt_input, row_key)
+    if rt_input.uniform_row_length is not None:
+      # If the inner dimension has uniform_row_length, then preserve it (by
+      # re-wrapping the values in a new RaggedTensor).  Note that the row
+      # length won't have changed, since we're slicing a range of rows (and not
+      # slicing the rows themselves).
+      sliced_rt_input = ragged_tensor.RaggedTensor.from_uniform_row_length(
+          sliced_rt_input.values, rt_input.uniform_row_length,
+          nrows=sliced_rt_input.nrows())
+    return _ragged_getitem_inner_dimensions(sliced_rt_input, inner_keys)
+
+  # Indexing a single row: slice values to get the indicated row, and then
+  # use a recursive call to __getitem__ to handle the inner keys.
+  else:
+    starts = rt_input.row_splits[:-1]
+    limits = rt_input.row_splits[1:]
+    if context.executing_eagerly():
+      # In python, __getitem__ should throw IndexError for out of bound
+      # indices. This will allow iteration run correctly as python will
+      # translate IndexError into StopIteration for next()/__next__().
+      # Below is an example:
+      #    import tensorflow as tf
+      #    r = tf.ragged.constant([[1., 2.], [3., 4., 5.], [6.]])
+      #    for elem in r:
+      #      print(elem)
+      # In non eager mode, the exception is thrown when session runs
+      # so we don't know if out of bound happens before.
+      # In eager mode, however, it is possible to find out when to
+      # throw out of bound IndexError.
+      # In the following row_key >= len(starts) is checked. In case of
+      # TypeError which happens when row_key is not an integer, the exception
+      # will simply be ignored as it will be processed later anyway.
+      try:
+        if int(row_key) >= len(starts):
+          raise IndexError("Row key {} out of bounds".format(row_key))
+      except (TypeError, ValueError):
+        pass
+    row = rt_input.values[starts[row_key]:limits[row_key]]
+    return row.__getitem__(inner_keys)
+
+
+def _slice_ragged_row_dimension(rt_input, row_key):
+  """Slice the outer dimension of `rt_input` according to the given `slice`.
+
+  Args:
+    rt_input: The `RaggedTensor` to slice.
+    row_key: The `slice` object that should be used to slice `rt_input`.
+
+  Returns:
+    A `RaggedTensor` containing the indicated slice of `rt_input`.
+  """
+  if row_key.start is None and row_key.stop is None and row_key.step is None:
+    return rt_input
+
+  # Use row_key to slice the starts & limits.
+  new_starts = rt_input.row_splits[:-1][row_key]
+  new_limits = rt_input.row_splits[1:][row_key]
+  zero_pad = array_ops.zeros([1], rt_input.row_splits.dtype)
+
+  # If there's no slice step, then we can just select a single continuous
+  # span of `ragged.values(rt_input)`.
+  if row_key.step is None or row_key.step == 1:
+    # Construct the new splits.  If new_starts and new_limits are empty,
+    # then this reduces to [0].  Otherwise, this reduces to:
+    #   concat([[new_starts[0]], new_limits])
+    new_splits = array_ops.concat(
+        [zero_pad[array_ops.size(new_starts):], new_starts[:1], new_limits],
+        axis=0)
+    values_start = new_splits[0]
+    values_limit = new_splits[-1]
+    return ragged_tensor.RaggedTensor.from_row_splits(
+        rt_input.values[values_start:values_limit], new_splits - values_start,
+        validate=False)
+
+  # If there is a slice step (aka a strided slice), then use ragged_gather to
+  # collect the necessary elements of `ragged.values(rt_input)`.
+  else:
+    return _build_ragged_tensor_from_value_ranges(new_starts, new_limits, 1,
+                                                  rt_input.values)
+
+
+def _ragged_getitem_inner_dimensions(rt_input, key_list):
+  """Retrieve inner dimensions, keeping outermost dimension unchanged.
+
+  Args:
+    rt_input: The `RaggedTensor` or `Tensor` from which a piece should be
+      extracted.
+    key_list: The __getitem__ keys for slicing the inner dimensions.
+
+  Returns:
+    A `RaggedTensor`.
+
+  Raises:
+    ValueError: If key_list is not supported.
+  """
+  if not key_list:
+    return rt_input
+
+  if not isinstance(rt_input, ragged_tensor.RaggedTensor):
+    return rt_input.__getitem__([slice(None, None, None)] + key_list)
+
+  column_key = key_list[0]
+  if column_key is Ellipsis:
+    expanded_key_list = _expand_ellipsis(key_list, rt_input.values.shape.ndims)
+    return _ragged_getitem_inner_dimensions(rt_input, expanded_key_list)
+
+  # Adding a new axis to a ragged inner dimension: recursively get the inner
+  # dimensions of rt_input with key_list[1:], and then wrap the result in a
+  # RaggedTensor that puts each value in its own row.
+  if column_key is array_ops.newaxis:
+    inner_rt = _ragged_getitem_inner_dimensions(rt_input, key_list[1:])
+    nsplits = tensor_shape.dimension_at_index(inner_rt.row_splits.shape, 0)
+    if nsplits.value is not None:
+      nsplits = nsplits.value
+    else:
+      nsplits = array_ops.shape(
+          inner_rt.row_splits, out_type=inner_rt.row_splits.dtype
+      )[0]
+    return ragged_tensor.RaggedTensor.from_uniform_row_length(
+        inner_rt, 1, nrows=nsplits - 1, validate=False)
+
+  # Slicing a range of columns in a ragged inner dimension.  We use a
+  # recursive call to process the values, and then assemble a RaggedTensor
+  # with those values.
+  if isinstance(column_key, slice):
+    if (column_key.start is None and column_key.stop is None and
+        column_key.step is None):
+      # Trivial slice: recursively process all values, & splits is unchanged.
+      return rt_input.with_values(
+          _ragged_getitem_inner_dimensions(rt_input.values, key_list[1:]))
+    else:
+      if not (
+          isinstance(column_key.start, (tensor_lib.Tensor, int, type(None)))
+          and isinstance(column_key.stop, (tensor_lib.Tensor, int, type(None)))
+      ):
+        raise TypeError("slice offsets must be integers or None")
+
+      # Nontrivial slice: use ragged_gather to extract the indicated slice as
+      # a new RaggedTensor (inner_rt), and then recursively process its values.
+      starts = rt_input.row_splits[:-1]
+      limits = rt_input.row_splits[1:]
+      step = 1 if column_key.step is None else column_key.step
+      lower_bound = _if_ge_zero(step, lambda: starts, lambda: starts - 1)
+      upper_bound = _if_ge_zero(step, lambda: limits, lambda: limits - 1)
+      # inner_rt_starts[i] = index to start gathering for row i.
+      if column_key.start is None:
+        inner_rt_starts = _if_ge_zero(step, lambda: starts, lambda: limits - 1)
+      else:
+        start_offset = math_ops.cast(column_key.start, starts.dtype)
+        inner_rt_starts = _if_ge_zero(
+            column_key.start,
+            lambda: math_ops.minimum(starts + start_offset, upper_bound),
+            lambda: math_ops.maximum(limits + start_offset, lower_bound))
+      # inner_rt_limits[i] = index to stop gathering for row i.
+      if column_key.stop is None:
+        inner_rt_limits = _if_ge_zero(step, lambda: limits, lambda: starts - 1)
+      else:
+        stop_offset = math_ops.cast(column_key.stop, starts.dtype)
+        inner_rt_limits = _if_ge_zero(
+            column_key.stop,
+            lambda: math_ops.minimum(starts + stop_offset, upper_bound),
+            lambda: math_ops.maximum(limits + stop_offset, lower_bound))
+      inner_rt = _build_ragged_tensor_from_value_ranges(
+          inner_rt_starts, inner_rt_limits, column_key.step, rt_input.values)
+      # If the row dimension is uniform, then calculate the new
+      # uniform_row_length, and rebuild inner_rt using that uniform_row_lengths.
+      if rt_input.uniform_row_length is not None:
+        new_row_length = _slice_length(rt_input.uniform_row_length, column_key)
+        inner_rt = ragged_tensor.RaggedTensor.from_uniform_row_length(
+            inner_rt.values, new_row_length, rt_input.nrows())
+      return inner_rt.with_values(
+          _ragged_getitem_inner_dimensions(inner_rt.values, key_list[1:]))
+
+  # Indexing a single column in a ragged inner dimension: raise an Exception.
+  # See RaggedTensor.__getitem__.__doc__ for an explanation of why indexing
+  # into a ragged inner dimension is problematic.
+  if rt_input.uniform_row_length is None:
+    raise ValueError("Cannot index into an inner ragged dimension.")
+
+  # Indexing a single column in a uniform inner dimension: check that the
+  # given index is in-bounds, and then use a strided slice over rt_input.values
+  # to take the indicated element from each row.
+  row_length = rt_input.uniform_row_length
+  column_key = math_ops.cast(column_key, row_length.dtype)
+  oob_err_msg = "Index out of bounds when indexing into a ragged tensor"
+  oob_checks = [
+      check_ops.assert_greater_equal(
+          column_key, -row_length, message=oob_err_msg),
+      check_ops.assert_less(column_key, row_length, message=oob_err_msg),
+  ]
+  with ops.control_dependencies(oob_checks):
+    offset = _if_ge_zero(column_key, lambda: column_key,
+                         lambda: row_length + column_key)
+    sliced_rt = rt_input.values[offset::row_length]
+    return _ragged_getitem_inner_dimensions(sliced_rt, key_list[1:])
+
+
+def _slice_length(value_length, slice_key):
+  """Computes the number of elements in a slice of a value with a given length.
+
+  Returns the equivalent of: `len(range(value_length)[slice_key])`
+
+  Args:
+    value_length: Scalar int `Tensor`: the length of the value being sliced.
+    slice_key: A `slice` object used to slice elements from the value.
+
+  Returns:
+    The number of elements in the sliced value.
+  """
+  # Note: we could compute the slice length without creating a zeros tensor
+  # with some variant of (stop-start)//step, but doing so would require more
+  # ops (for checking bounds, handling negative indices, negative step sizes,
+  # etc); and we expect this to be an uncommon operation, so we use this
+  # simpler implementation.
+  zeros = array_ops.zeros(value_length, dtype=dtypes.bool)
+  return array_ops.size(zeros[slice_key], out_type=value_length.dtype)
+
+
+def _expand_ellipsis(key_list, num_remaining_dims):
+  """Expands the ellipsis at the start of `key_list`.
+
+  Assumes that the first element of `key_list` is Ellipsis.  This will either
+  remove the Ellipsis (if it corresponds to zero indices) or prepend a new
+  `slice(None, None, None)` (if it corresponds to more than zero indices).
+
+  Args:
+    key_list: The arguments to `__getitem__()`.
+    num_remaining_dims: The number of dimensions remaining.
+
+  Returns:
+    A copy of `key_list` with he ellipsis expanded.
+  Raises:
+    ValueError: If ragged_rank.shape.ndims is None
+    IndexError: If there are too many elements in `key_list`.
+  """
+  if num_remaining_dims is None:
+    raise ValueError("Ellipsis not supported for unknown shape RaggedTensors")
+  num_indices = sum(1 for idx in key_list if idx is not array_ops.newaxis)
+  if num_indices > num_remaining_dims + 1:
+    raise IndexError("Too many indices for RaggedTensor")
+  elif num_indices == num_remaining_dims + 1:
+    return key_list[1:]
+  else:
+    return [slice(None, None, None)] + key_list
+
+
+def _tensors_in_key_list(key_list):
+  """Generates all Tensors in the given slice spec."""
+  if isinstance(key_list, tensor_lib.Tensor):
+    yield key_list
+  if isinstance(key_list, (list, tuple)):
+    for v in key_list:
+      for tensor in _tensors_in_key_list(v):
+        yield tensor
+  if isinstance(key_list, slice):
+    for tensor in _tensors_in_key_list(key_list.start):
+      yield tensor
+    for tensor in _tensors_in_key_list(key_list.stop):
+      yield tensor
+    for tensor in _tensors_in_key_list(key_list.step):
+      yield tensor
+
+
+def _build_ragged_tensor_from_value_ranges(starts, limits, step, values):
+  """Returns a `RaggedTensor` containing the specified sequences of values.
+
+  Returns a RaggedTensor `output` where:
+
+  ```python
+  output.shape[0] = starts.shape[0]
+  output[i] = values[starts[i]:limits[i]:step]
+  ```
+
+  Requires that `starts.shape == limits.shape` and
+  `0 <= starts[i] <= limits[i] <= values.shape[0]`.
+
+  Args:
+    starts: 1D integer Tensor specifying the start indices for the sequences of
+      values to include.
+    limits: 1D integer Tensor specifying the limit indices for the sequences of
+      values to include.
+    step: Integer value specifying the step size for strided slices.
+    values: The set of values to select from.
+
+  Returns:
+    A `RaggedTensor`.
+
+  Raises:
+    ValueError: Until the prerequisite ops are checked in.
+  """
+  # Use `ragged_range` to get the index of each value we should include.
+  if step is None:
+    step = 1
+  step = ops.convert_to_tensor(step, name="step")
+  if step.dtype.is_integer:
+    step = math_ops.cast(step, starts.dtype)
+  else:
+    raise TypeError("slice strides must be integers or None")
+  value_indices = ragged_math_ops.range(starts, limits, step,
+                                        row_splits_dtype=starts.dtype)
+
+  # Use `ragged_gather` or `array_ops.gather` to collect the values.
+  if isinstance(values, ragged_tensor.RaggedTensor):
+    gathered_values = ragged_gather_ops.gather(
+        params=values, indices=value_indices.values)
+  else:
+    gathered_values = array_ops.gather(
+        params=values, indices=value_indices.values)
+
+  # Assemble the RaggedTensor from splits & values.
+  return value_indices.with_values(gathered_values)
+
+
+def _if_ge_zero(value, true_fn, false_fn):
+  """Returns `true_fn() if value >= 0 else false_fn()`."""
+  # If `value` is statically known, then don't use a control flow op.
+  if isinstance(value, tensor_lib.Tensor):
+    const_value = tensor_util.constant_value(value)
+    if const_value is None:
+      return cond.cond(value >= 0, true_fn, false_fn)
+    else:
+      value = const_value
+  if value >= 0:
+    return true_fn()
+  else:
+    return false_fn()

tensorflow/python/ops/ragged/ragged_getitem_test.py

@@ -0,0 +1,589 @@
+# 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.
+# ==============================================================================
+"""Tests for third_party.tensorflow.python.ops.ragged_tensor."""
+
+import re
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import tensor_getitem_override
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
+
+from tensorflow.python.platform import googletest
+
+
+class _SliceBuilder:
+  """Helper to construct arguments for __getitem__.
+
+  Usage: _SliceBuilder()[<expr>] slice_spec Python generates for <expr>.
+  """
+
+  def __getitem__(self, slice_spec):
+    return slice_spec
+
+
+SLICE_BUILDER = _SliceBuilder()
+
+
+def _make_tensor_slice_spec(slice_spec, use_constant=True):
+  """Wraps all integers in an extended slice spec w/ a tensor.
+
+  This function is used to help test slicing when the slice spec contains
+  tensors, rather than integers.
+
+  Args:
+    slice_spec: The extended slice spec.
+    use_constant: If true, then wrap each integer with a tf.constant.  If false,
+      then wrap each integer with a tf.placeholder.
+
+  Returns:
+    A copy of slice_spec, but with each integer i replaced with tf.constant(i).
+  """
+
+  def make_piece_scalar(piece):
+    if isinstance(piece, int):
+      scalar = constant_op.constant(piece)
+      if use_constant:
+        return scalar
+      else:
+        return array_ops.placeholder_with_default(scalar, [])
+    elif isinstance(piece, slice):
+      return slice(
+          make_piece_scalar(piece.start), make_piece_scalar(piece.stop),
+          make_piece_scalar(piece.step))
+    else:
+      return piece
+
+  if isinstance(slice_spec, tuple):
+    return tuple(make_piece_scalar(piece) for piece in slice_spec)
+  else:
+    return make_piece_scalar(slice_spec)
+
+
+# Example 2D ragged tensor value with one ragged dimension and with scalar
+# values, expressed as nested python lists and as splits+values.
+EXAMPLE_RAGGED_TENSOR_2D = [[b'a', b'b'], [b'c', b'd', b'e'], [b'f'], [],
+                            [b'g']]
+EXAMPLE_RAGGED_TENSOR_2D_SPLITS = [0, 2, 5, 6, 6, 7]
+EXAMPLE_RAGGED_TENSOR_2D_VALUES = ['a', 'b', 'c', 'd', 'e', 'f', 'g']
+
+# Example 4D ragged tensor value, with two ragged dimensions and with values
+# whose shape is [2], expressed as nested python lists and as splits+values.
+EXAMPLE_RAGGED_TENSOR_4D = [
+    [                                       # rt[0]
+        [[1, 2], [3, 4], [5, 6]],           # rt[0][0]
+        [[7, 8], [9, 10], [11, 12]]],       # rt[0][1]
+    [],                                     # rt[1]
+    [                                       # rt[2]
+        [[13, 14], [15, 16], [17, 18]]],    # rt[2][0]
+    [                                       # rt[3]
+        [[19, 20]]]                         # rt[3][0]
+]  # pyformat: disable
+EXAMPLE_RAGGED_TENSOR_4D_SPLITS1 = [0, 2, 2, 3, 4]
+EXAMPLE_RAGGED_TENSOR_4D_SPLITS2 = [0, 3, 6, 9, 10]
+EXAMPLE_RAGGED_TENSOR_4D_VALUES = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10],
+                                   [11, 12], [13, 14], [15, 16], [17, 18],
+                                   [19, 20]]
+
+# Example 3D ragged tensor with uniform_row_lengths.
+EXAMPLE_RAGGED_TENSOR_3D = [[[1, 2, 3], [4], [5, 6]], [[], [7, 8, 9], []]]
+EXAMPLE_RAGGED_TENSOR_3D_ROWLEN = 3
+EXAMPLE_RAGGED_TENSOR_3D_SPLITS = [0, 3, 4, 6, 6, 9, 9]
+EXAMPLE_RAGGED_TENSOR_3D_VALUES = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedGetItemTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+  longMessage = True  # Property in unittest.Testcase. pylint: disable=invalid-name
+
+  #=============================================================================
+  # RaggedTensor.__getitem__
+  #=============================================================================
+
+  def _TestGetItem(self, rt, slice_spec, expected, expected_shape=None):
+    """Helper function for testing RaggedTensor.__getitem__.
+
+    Checks that calling `rt.__getitem__(slice_spec) returns the expected value.
+    Checks three different configurations for each slice spec:
+
+      * Call __getitem__ with the slice spec as-is (with int values)
+      * Call __getitem__ with int values in the slice spec wrapped in
+        `tf.constant()`.
+      * Call __getitem__ with int values in the slice spec wrapped in
+        `tf.compat.v1.placeholder()` (so value is not known at graph
+        construction time).
+
+    Args:
+      rt: The RaggedTensor to test.
+      slice_spec: The slice spec.
+      expected: The expected value of rt.__getitem__(slice_spec), as a python
+        list; or an exception class.
+      expected_shape: The expected shape for `rt.__getitem__(slice_spec)`.
+    """
+    tensor_slice_spec1 = _make_tensor_slice_spec(slice_spec, True)
+    tensor_slice_spec2 = _make_tensor_slice_spec(slice_spec, False)
+    value1 = rt.__getitem__(slice_spec)
+    value2 = rt.__getitem__(tensor_slice_spec1)
+    value3 = rt.__getitem__(tensor_slice_spec2)
+    self.assertAllEqual(value1, expected, 'slice_spec=%s' % (slice_spec,))
+    self.assertAllEqual(value2, expected, 'slice_spec=%s' % (slice_spec,))
+    self.assertAllEqual(value3, expected, 'slice_spec=%s' % (slice_spec,))
+    if expected_shape is not None:
+      value1.shape.assert_is_compatible_with(expected_shape)
+      value2.shape.assert_is_compatible_with(expected_shape)
+      value3.shape.assert_is_compatible_with(expected_shape)
+
+  def _TestGetItemException(self, rt, slice_spec, expected, message):
+    """Helper function for testing RaggedTensor.__getitem__ exceptions."""
+    tensor_slice_spec = _make_tensor_slice_spec(slice_spec, True)
+    with self.assertRaisesRegex(expected, message):
+      self.evaluate(rt.__getitem__(slice_spec))
+    with self.assertRaisesRegex(expected, message):
+      self.evaluate(rt.__getitem__(tensor_slice_spec))
+
+  @parameterized.parameters(
+      # Tests for rt[i]
+      (SLICE_BUILDER[-5], EXAMPLE_RAGGED_TENSOR_2D[-5]),
+      (SLICE_BUILDER[-4], EXAMPLE_RAGGED_TENSOR_2D[-4]),
+      (SLICE_BUILDER[-1], EXAMPLE_RAGGED_TENSOR_2D[-1]),
+      (SLICE_BUILDER[0], EXAMPLE_RAGGED_TENSOR_2D[0]),
+      (SLICE_BUILDER[1], EXAMPLE_RAGGED_TENSOR_2D[1]),
+      (SLICE_BUILDER[4], EXAMPLE_RAGGED_TENSOR_2D[4]),
+
+      # Tests for rt[i:]
+      (SLICE_BUILDER[-6:], EXAMPLE_RAGGED_TENSOR_2D[-6:]),
+      (SLICE_BUILDER[-3:], EXAMPLE_RAGGED_TENSOR_2D[-3:]),
+      (SLICE_BUILDER[-1:], EXAMPLE_RAGGED_TENSOR_2D[-1:]),
+      (SLICE_BUILDER[0:], EXAMPLE_RAGGED_TENSOR_2D[0:]),
+      (SLICE_BUILDER[3:], EXAMPLE_RAGGED_TENSOR_2D[3:]),
+      (SLICE_BUILDER[5:], EXAMPLE_RAGGED_TENSOR_2D[5:]),
+
+      # Tests for rt[:j]
+      (SLICE_BUILDER[:-6], EXAMPLE_RAGGED_TENSOR_2D[:-6]),
+      (SLICE_BUILDER[:-3], EXAMPLE_RAGGED_TENSOR_2D[:-3]),
+      (SLICE_BUILDER[:-1], EXAMPLE_RAGGED_TENSOR_2D[:-1]),
+      (SLICE_BUILDER[:0], EXAMPLE_RAGGED_TENSOR_2D[:0]),
+      (SLICE_BUILDER[:3], EXAMPLE_RAGGED_TENSOR_2D[:3]),
+      (SLICE_BUILDER[:5], EXAMPLE_RAGGED_TENSOR_2D[:5]),
+
+      # Tests for rt[i:j]
+      (SLICE_BUILDER[0:3], EXAMPLE_RAGGED_TENSOR_2D[0:3]),
+      (SLICE_BUILDER[3:5], EXAMPLE_RAGGED_TENSOR_2D[3:5]),
+      (SLICE_BUILDER[-5:3], EXAMPLE_RAGGED_TENSOR_2D[-5:3]),
+      (SLICE_BUILDER[3:1], EXAMPLE_RAGGED_TENSOR_2D[3:1]),
+      (SLICE_BUILDER[-1:1], EXAMPLE_RAGGED_TENSOR_2D[-1:1]),
+      (SLICE_BUILDER[1:-1], EXAMPLE_RAGGED_TENSOR_2D[1:-1]),
+
+      # Tests for rt[i, j]
+      (SLICE_BUILDER[0, 1], EXAMPLE_RAGGED_TENSOR_2D[0][1]),
+      (SLICE_BUILDER[1, 2], EXAMPLE_RAGGED_TENSOR_2D[1][2]),
+      (SLICE_BUILDER[-1, 0], EXAMPLE_RAGGED_TENSOR_2D[-1][0]),
+      (SLICE_BUILDER[-3, 0], EXAMPLE_RAGGED_TENSOR_2D[-3][0]),
+      (SLICE_BUILDER[:], EXAMPLE_RAGGED_TENSOR_2D),
+      (SLICE_BUILDER[:, :], EXAMPLE_RAGGED_TENSOR_2D),
+
+      # Empty slice spec.
+      ([], EXAMPLE_RAGGED_TENSOR_2D),
+
+      # Test for ellipsis
+      (SLICE_BUILDER[...], EXAMPLE_RAGGED_TENSOR_2D),
+      (SLICE_BUILDER[2, ...], EXAMPLE_RAGGED_TENSOR_2D[2]),
+      (SLICE_BUILDER[..., :], EXAMPLE_RAGGED_TENSOR_2D),
+      (SLICE_BUILDER[..., 2, 0], EXAMPLE_RAGGED_TENSOR_2D[2][0]),
+      (SLICE_BUILDER[2, ..., 0], EXAMPLE_RAGGED_TENSOR_2D[2][0]),
+      (SLICE_BUILDER[2, 0, ...], EXAMPLE_RAGGED_TENSOR_2D[2][0]),
+
+      # Test for array_ops.newaxis
+      (SLICE_BUILDER[array_ops.newaxis, :], [EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[:, array_ops.newaxis],
+       [[row] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+
+      # Slicing inner ragged dimensions.
+      (SLICE_BUILDER[-1:,
+                     1:4], [row[1:4] for row in EXAMPLE_RAGGED_TENSOR_2D[-1:]]),
+      (SLICE_BUILDER[:, 1:4], [row[1:4] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[:, -2:], [row[-2:] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+
+      # Strided slices
+      (SLICE_BUILDER[::2], EXAMPLE_RAGGED_TENSOR_2D[::2]),
+      (SLICE_BUILDER[::-1], EXAMPLE_RAGGED_TENSOR_2D[::-1]),
+      (SLICE_BUILDER[::-2], EXAMPLE_RAGGED_TENSOR_2D[::-2]),
+      (SLICE_BUILDER[::-3], EXAMPLE_RAGGED_TENSOR_2D[::-3]),
+      (SLICE_BUILDER[:, ::2], [row[::2] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[:, ::-1], [row[::-1] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[:, ::-2], [row[::-2] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[:, ::-3], [row[::-3] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[:, 2::-1],
+       [row[2::-1] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[:, -1::-1],
+       [row[-1::-1] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[..., -1::-1],
+       [row[-1::-1] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[:, 2::-2],
+       [row[2::-2] for row in EXAMPLE_RAGGED_TENSOR_2D]),
+      (SLICE_BUILDER[::-1, ::-1],
+       [row[::-1] for row in EXAMPLE_RAGGED_TENSOR_2D[::-1]]),
+  )  # pyformat: disable
+  def testWithRaggedRank1(self, slice_spec, expected):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    # Ragged tensor
+    rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES,
+                                      EXAMPLE_RAGGED_TENSOR_2D_SPLITS)
+
+    self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_2D)
+    self._TestGetItem(rt, slice_spec, expected)
+
+  # pylint: disable=g-complex-comprehension
+  @parameterized.parameters([(start, stop)
+                             for start in [-2, -1, None, 0, 1, 2]
+                             for stop in [-2, -1, None, 0, 1, 2]])
+  def testWithStridedSlices(self, start, stop):
+    test_value = [[1, 2, 3, 4, 5], [6, 7], [8, 9, 10], [], [9],
+                  [1, 2, 3, 4, 5, 6, 7, 8]]
+    rt = ragged_factory_ops.constant(test_value)
+    for step in [-3, -2, -1, 1, 2, 3]:
+      # Slice outer dimension
+      self.assertAllEqual(rt[start:stop:step], test_value[start:stop:step],
+                          'slice=%s:%s:%s' % (start, stop, step))
+      # Slice inner dimension
+      self.assertAllEqual(rt[:, start:stop:step],
+                          [row[start:stop:step] for row in test_value],
+                          'slice=%s:%s:%s' % (start, stop, step))
+
+  # pylint: disable=invalid-slice-index
+  @parameterized.parameters(
+      # Tests for out-of-bound errors
+      (SLICE_BUILDER[5], (IndexError, ValueError, errors.InvalidArgumentError),
+       '.*out of bounds.*'),
+      (SLICE_BUILDER[-6], (IndexError, ValueError, errors.InvalidArgumentError),
+       '.*out of bounds.*'),
+      (SLICE_BUILDER[0, 2], (IndexError, ValueError,
+                             errors.InvalidArgumentError), '.*out of bounds.*'),
+      (SLICE_BUILDER[3, 0], (IndexError, ValueError,
+                             errors.InvalidArgumentError), '.*out of bounds.*'),
+
+      # Indexing into an inner ragged dimension
+      (SLICE_BUILDER[:, 3], ValueError,
+       'Cannot index into an inner ragged dimension'),
+      (SLICE_BUILDER[:1, 3], ValueError,
+       'Cannot index into an inner ragged dimension'),
+      (SLICE_BUILDER[..., 3], ValueError,
+       'Cannot index into an inner ragged dimension'),
+
+      # Tests for type errors
+      (SLICE_BUILDER[0.5], TypeError, re.escape(
+          tensor_getitem_override._SLICE_TYPE_ERROR)),
+      (SLICE_BUILDER[1:3:0.5], TypeError, re.escape(
+          tensor_getitem_override._SLICE_TYPE_ERROR)),
+      (SLICE_BUILDER[:, 1:3:0.5], TypeError,
+       'slice strides must be integers or None'),
+      (SLICE_BUILDER[:, 0.5:1.5], TypeError,
+       'slice offsets must be integers or None'),
+      (SLICE_BUILDER['foo'], TypeError, re.escape(
+          tensor_getitem_override._SLICE_TYPE_ERROR)),
+      (SLICE_BUILDER[:, 'foo':'foo'], TypeError,
+       'slice offsets must be integers or None'),
+
+      # Tests for other errors
+      (SLICE_BUILDER[..., 0, 0,
+                     0], IndexError, 'Too many indices for RaggedTensor'),
+  )
+  def testErrorsWithRaggedRank1(self, slice_spec, expected, message):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    # Ragged tensor
+    rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES,
+                                      EXAMPLE_RAGGED_TENSOR_2D_SPLITS)
+
+    self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_2D)
+    self._TestGetItemException(rt, slice_spec, expected, message)
+
+  @parameterized.parameters(
+      # Tests for rt[index, index, ...]
+      (SLICE_BUILDER[2, 0], EXAMPLE_RAGGED_TENSOR_4D[2][0]),
+      (SLICE_BUILDER[2, 0, 1], EXAMPLE_RAGGED_TENSOR_4D[2][0][1]),
+      (SLICE_BUILDER[2, 0, 1, 1], EXAMPLE_RAGGED_TENSOR_4D[2][0][1][1]),
+      (SLICE_BUILDER[2, 0, 1:], EXAMPLE_RAGGED_TENSOR_4D[2][0][1:]),
+      (SLICE_BUILDER[2, 0, 1:, 1:], [[16], [18]]),
+      (SLICE_BUILDER[2, 0, :, 1], [14, 16, 18]),
+      (SLICE_BUILDER[2, 0, 1, :], EXAMPLE_RAGGED_TENSOR_4D[2][0][1]),
+
+      # Tests for rt[index, slice, ...]
+      (SLICE_BUILDER[0, :], EXAMPLE_RAGGED_TENSOR_4D[0]),
+      (SLICE_BUILDER[1, :], EXAMPLE_RAGGED_TENSOR_4D[1]),
+      (SLICE_BUILDER[0, :, :, 1], [[2, 4, 6], [8, 10, 12]]),
+      (SLICE_BUILDER[1, :, :, 1], []),
+      (SLICE_BUILDER[2, :, :, 1], [[14, 16, 18]]),
+      (SLICE_BUILDER[3, :, :, 1], [[20]]),
+
+      # Tests for rt[slice, slice, ...]
+      (SLICE_BUILDER[:, :], EXAMPLE_RAGGED_TENSOR_4D),
+      (SLICE_BUILDER[:, :, :, 1], [[[2, 4, 6], [8, 10, 12]], [], [[14, 16, 18]],
+                                   [[20]]]),
+      (SLICE_BUILDER[1:, :, :, 1], [[], [[14, 16, 18]], [[20]]]),
+      (SLICE_BUILDER[-3:, :, :, 1], [[], [[14, 16, 18]], [[20]]]),
+
+      # Test for ellipsis
+      (SLICE_BUILDER[...], EXAMPLE_RAGGED_TENSOR_4D),
+      (SLICE_BUILDER[2, ...], EXAMPLE_RAGGED_TENSOR_4D[2]),
+      (SLICE_BUILDER[2, 0, ...], EXAMPLE_RAGGED_TENSOR_4D[2][0]),
+      (SLICE_BUILDER[..., 0], [[[1, 3, 5], [7, 9, 11]], [], [[13, 15, 17]],
+                               [[19]]]),
+      (SLICE_BUILDER[2, ..., 0], [[13, 15, 17]]),
+      (SLICE_BUILDER[2, 0, ..., 0], [13, 15, 17]),
+
+      # Test for array_ops.newaxis
+      (SLICE_BUILDER[array_ops.newaxis, :], [EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, array_ops.newaxis],
+       [[row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+
+      # Empty slice spec.
+      ([], EXAMPLE_RAGGED_TENSOR_4D),
+
+      # Slicing inner ragged dimensions.
+      (SLICE_BUILDER[:, 1:4], [row[1:4] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, -2:], [row[-2:] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, :, :-1],
+       [[v[:-1] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, :, 1:2],
+       [[v[1:2] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[1:, 1:3, 1:2],
+       [[v[1:2] for v in row[1:3]] for row in EXAMPLE_RAGGED_TENSOR_4D[1:]]),
+
+      # Strided slices
+      (SLICE_BUILDER[::2], EXAMPLE_RAGGED_TENSOR_4D[::2]),
+      (SLICE_BUILDER[::-1], EXAMPLE_RAGGED_TENSOR_4D[::-1]),
+      (SLICE_BUILDER[::-2], EXAMPLE_RAGGED_TENSOR_4D[::-2]),
+      (SLICE_BUILDER[1::2], EXAMPLE_RAGGED_TENSOR_4D[1::2]),
+      (SLICE_BUILDER[:, ::2], [row[::2] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, 1::2], [row[1::2] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, :, ::2],
+       [[v[::2] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, :, 1::2],
+       [[v[1::2] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, :, ::-1],
+       [[v[::-1] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[:, :, ::-2],
+       [[v[::-2] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[..., ::-1, :],
+       [[v[::-1] for v in row] for row in EXAMPLE_RAGGED_TENSOR_4D]),
+      (SLICE_BUILDER[..., ::-1], [[[v[::-1] for v in col] for col in row]
+                                  for row in EXAMPLE_RAGGED_TENSOR_4D]),
+  )  # pyformat: disable
+  def testWithRaggedRank2(self, slice_spec, expected):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    rt = RaggedTensor.from_nested_row_splits(
+        EXAMPLE_RAGGED_TENSOR_4D_VALUES,
+        [EXAMPLE_RAGGED_TENSOR_4D_SPLITS1, EXAMPLE_RAGGED_TENSOR_4D_SPLITS2])
+    self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_4D)
+    self._TestGetItem(rt, slice_spec, expected)
+
+  @parameterized.parameters(
+      # Test for errors in unsupported cases
+      (SLICE_BUILDER[:, 0], ValueError,
+       'Cannot index into an inner ragged dimension.'),
+      (SLICE_BUILDER[:, :, 0], ValueError,
+       'Cannot index into an inner ragged dimension.'),
+
+      # Test for out-of-bounds errors.
+      (SLICE_BUILDER[1, 0], (IndexError, ValueError,
+                             errors.InvalidArgumentError), '.*out of bounds.*'),
+      (SLICE_BUILDER[0, 0, 3],
+       (IndexError, ValueError,
+        errors.InvalidArgumentError), '.*out of bounds.*'),
+      (SLICE_BUILDER[5], (IndexError, ValueError, errors.InvalidArgumentError),
+       '.*out of bounds.*'),
+      (SLICE_BUILDER[0, 5], (IndexError, ValueError,
+                             errors.InvalidArgumentError), '.*out of bounds.*'),
+  )
+  def testErrorsWithRaggedRank2(self, slice_spec, expected, message):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    rt = RaggedTensor.from_nested_row_splits(
+        EXAMPLE_RAGGED_TENSOR_4D_VALUES,
+        [EXAMPLE_RAGGED_TENSOR_4D_SPLITS1, EXAMPLE_RAGGED_TENSOR_4D_SPLITS2])
+    self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_4D)
+    self._TestGetItemException(rt, slice_spec, expected, message)
+
+  @parameterized.parameters(
+      (SLICE_BUILDER[:], []),
+      (SLICE_BUILDER[2:], []),
+      (SLICE_BUILDER[:-3], []),
+  )
+  def testWithEmptyTensor(self, slice_spec, expected):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    rt = RaggedTensor.from_row_splits([], [0])
+    self._TestGetItem(rt, slice_spec, expected)
+
+  @parameterized.parameters(
+      (SLICE_BUILDER[0], (IndexError, ValueError, errors.InvalidArgumentError),
+       '.*out of bounds.*'),
+      (SLICE_BUILDER[-1], (IndexError, ValueError, errors.InvalidArgumentError),
+       '.*out of bounds.*'),
+  )
+  def testErrorsWithEmptyTensor(self, slice_spec, expected, message):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    rt = RaggedTensor.from_row_splits([], [0])
+    self._TestGetItemException(rt, slice_spec, expected, message)
+
+  @parameterized.parameters(
+      (SLICE_BUILDER[-4], EXAMPLE_RAGGED_TENSOR_2D[-4]),
+      (SLICE_BUILDER[0], EXAMPLE_RAGGED_TENSOR_2D[0]),
+      (SLICE_BUILDER[-3:], EXAMPLE_RAGGED_TENSOR_2D[-3:]),
+      (SLICE_BUILDER[:3], EXAMPLE_RAGGED_TENSOR_2D[:3]),
+      (SLICE_BUILDER[3:5], EXAMPLE_RAGGED_TENSOR_2D[3:5]),
+      (SLICE_BUILDER[0, 1], EXAMPLE_RAGGED_TENSOR_2D[0][1]),
+      (SLICE_BUILDER[-3, 0], EXAMPLE_RAGGED_TENSOR_2D[-3][0]),
+  )
+  def testWithPlaceholderShapes(self, slice_spec, expected):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    # Intentionally use an unknown shape for `splits`, to force the code path
+    # that deals with having nrows unknown at graph construction time.
+    splits = constant_op.constant(
+        EXAMPLE_RAGGED_TENSOR_2D_SPLITS, dtype=dtypes.int64)
+    splits = array_ops.placeholder_with_default(splits, None)
+    rt = RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_2D_VALUES, splits)
+    self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_2D)
+    self._TestGetItem(rt, slice_spec, expected)
+
+  @parameterized.parameters(
+      (SLICE_BUILDER[..., 2], ValueError,
+       'Ellipsis not supported for unknown shape RaggedTensors'),)
+  def testErrorsWithPlaceholderShapes(self, slice_spec, expected, message):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    if not context.executing_eagerly():
+      # Intentionally use an unknown shape for `values`.
+      values = array_ops.placeholder_with_default([0], None)
+      rt = RaggedTensor.from_row_splits(values, [0, 1])
+      self._TestGetItemException(rt, slice_spec, expected, message)
+
+  def testNewAxis(self):
+    # rt: [[[['a', 'b'], ['c', 'd']], [], [['e', 'f']]], []]
+    splits1 = [0, 3, 3]
+    splits2 = [0, 2, 2, 3]
+    values = constant_op.constant([['a', 'b'], ['c', 'd'], ['e', 'f']])
+    rt = RaggedTensor.from_nested_row_splits(values, [splits1, splits2])
+    rt_newaxis0 = rt[array_ops.newaxis]
+    rt_newaxis1 = rt[:, array_ops.newaxis]
+    rt_newaxis2 = rt[:, :, array_ops.newaxis]
+    rt_newaxis3 = rt[:, :, :, array_ops.newaxis]
+    rt_newaxis4 = rt[:, :, :, :, array_ops.newaxis]
+
+    self.assertAllEqual(
+        rt, [[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]], []])
+    self.assertAllEqual(
+        rt_newaxis0, [[[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]], []]])
+    self.assertAllEqual(
+        rt_newaxis1,
+        [[[[[b'a', b'b'], [b'c', b'd']], [], [[b'e', b'f']]]], [[]]])
+    self.assertAllEqual(
+        rt_newaxis2,
+        [[[[[b'a', b'b'], [b'c', b'd']]], [[]], [[[b'e', b'f']]]], []])
+    self.assertAllEqual(
+        rt_newaxis3,
+        [[[[[b'a', b'b']], [[b'c', b'd']]], [], [[[b'e', b'f']]]], []])
+    self.assertAllEqual(
+        rt_newaxis4,
+        [[[[[b'a'], [b'b']], [[b'c'], [b'd']]], [], [[[b'e'], [b'f']]]], []])
+
+    self.assertEqual(rt.ragged_rank, 2)
+    self.assertEqual(rt_newaxis0.ragged_rank, 3)
+    self.assertEqual(rt_newaxis1.ragged_rank, 3)
+    self.assertEqual(rt_newaxis2.ragged_rank, 3)
+    self.assertEqual(rt_newaxis3.ragged_rank, 2)
+    self.assertEqual(rt_newaxis4.ragged_rank, 2)
+
+    self.assertEqual(rt_newaxis0.shape.as_list(), [1, 2, None, None, 2])
+    self.assertEqual(rt_newaxis1.shape.as_list(), [2, 1, None, None, 2])
+    self.assertEqual(rt_newaxis2.shape.as_list(), [2, None, 1, None, 2])
+    self.assertEqual(rt_newaxis3.shape.as_list(), [2, None, None, 1, 2])
+    self.assertEqual(rt_newaxis4.shape.as_list(), [2, None, None, 2, 1])
+
+  @parameterized.parameters(
+      # EXAMPLE_RAGGED_TENSOR_3D.shape = [2, 3, None]
+
+      # Indexing into uniform_row_splits dimension:
+      (SLICE_BUILDER[:, 1], [r[1] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, None]),
+      (SLICE_BUILDER[:, 2], [r[2] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, None]),
+      (SLICE_BUILDER[:, -2], [r[-2] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, None]),
+      (SLICE_BUILDER[:, -3], [r[-3] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, None]),
+      (SLICE_BUILDER[1:, 2], [r[2] for r in EXAMPLE_RAGGED_TENSOR_3D[1:]],
+       [1, None]),
+      (SLICE_BUILDER[:, 1, 1:], [r[1][1:] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, None]),
+      (SLICE_BUILDER[1:, 1, 1:],
+       [r[1][1:] for r in EXAMPLE_RAGGED_TENSOR_3D[1:]],
+       [1, None]),
+
+      # Slicing uniform_row_splits dimension:
+      (SLICE_BUILDER[:, 2:], [r[2:] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, 1, None]),
+      (SLICE_BUILDER[:, -2:], [r[-2:] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, 2, None]),
+      (SLICE_BUILDER[:, :, 1:],
+       [[c[1:] for c in r] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, 3, None]),
+      (SLICE_BUILDER[:, 5:], [r[5:] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, 0, None]),
+
+      # Slicing uniform_row_splits dimension with a non-default step size:
+      (SLICE_BUILDER[:, ::2], [r[::2] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, 2, None]),
+      (SLICE_BUILDER[:, ::-1], [r[::-1] for r in EXAMPLE_RAGGED_TENSOR_3D],
+       [2, 3, None]),
+  )  # pyformat: disable
+  def testWithUniformRowLength(self, slice_spec, expected, expected_shape):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    rt = RaggedTensor.from_uniform_row_length(
+        RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_3D_VALUES,
+                                     EXAMPLE_RAGGED_TENSOR_3D_SPLITS),
+        EXAMPLE_RAGGED_TENSOR_3D_ROWLEN)
+    self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_3D)
+    self.assertIsNot(rt.uniform_row_length, None)
+    self._TestGetItem(rt, slice_spec, expected, expected_shape)
+
+    # If the result is 3D, then check that it still has a uniform row length:
+    actual = rt.__getitem__(slice_spec)  # pylint: disable=assignment-from-no-return
+    if actual.shape.rank == 3:
+      self.assertIsNot(actual.uniform_row_length, None)
+      self.assertAllEqual(actual.uniform_row_length, expected_shape[1])
+
+  @parameterized.parameters(
+      (SLICE_BUILDER[:, 3], errors.InvalidArgumentError, 'out of bounds'),
+      (SLICE_BUILDER[:, -4], errors.InvalidArgumentError, 'out of bounds'),
+      (SLICE_BUILDER[:, 10], errors.InvalidArgumentError, 'out of bounds'),
+      (SLICE_BUILDER[:, -10], errors.InvalidArgumentError, 'out of bounds'),
+  )
+  def testErrorsWithUniformRowLength(self, slice_spec, expected, message):
+    """Test that rt.__getitem__(slice_spec) == expected."""
+    rt = RaggedTensor.from_uniform_row_length(
+        RaggedTensor.from_row_splits(EXAMPLE_RAGGED_TENSOR_3D_VALUES,
+                                     EXAMPLE_RAGGED_TENSOR_3D_SPLITS),
+        EXAMPLE_RAGGED_TENSOR_3D_ROWLEN)
+    self.assertAllEqual(rt, EXAMPLE_RAGGED_TENSOR_3D)
+    self._TestGetItemException(rt, slice_spec, expected, message)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_image_ops.py

@@ -0,0 +1,98 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Image operations for RaggedTensors."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import map_fn
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import dispatch
+
+
+@dispatch.dispatch_for_api(image_ops.resize_images_v2)
+def resize_images_v2(images: ragged_tensor.RaggedTensor,
+                     size,
+                     method=image_ops.ResizeMethod.BILINEAR,
+                     preserve_aspect_ratio=False,
+                     antialias=False,
+                     name=None):
+  """RaggedTensor dispatcher for tf.image.resize (tf-v2)."""
+  with ops.name_scope(name, "RaggedResizeImages", [images, size]):
+    return _resize_images(
+        image_ops.resize_images_v2,
+        images,
+        size,
+        method=method,
+        preserve_aspect_ratio=preserve_aspect_ratio,
+        antialias=antialias)
+
+
+@dispatch.dispatch_for_api(image_ops.resize_images)
+def resize_images_v1(images: ragged_tensor.RaggedTensor,
+                     size,
+                     method=image_ops.ResizeMethodV1.BILINEAR,
+                     align_corners=False,
+                     preserve_aspect_ratio=False,
+                     name=None):
+  """RaggedTensor dispatcher for tf.image.resize (tf-v1)."""
+  with ops.name_scope(name, "RaggedResizeImages", [images, size]):
+    return _resize_images(
+        image_ops.resize_images,
+        images,
+        size,
+        method=method,
+        preserve_aspect_ratio=preserve_aspect_ratio,
+        align_corners=align_corners)
+
+
+def _resize_images(resize_op, images, size, **kwargs):
+  """RaggedTensor dispatcher for tf.image.resize."""
+  if images.shape.rank != 4:
+    raise ValueError(
+        "tf.image.resize: images.shape.rank must be 4 if images is ragged.")
+
+  # Determine the output shape (excluding the batch dimension).
+  static_batch_size = tensor_shape.dimension_value(images.shape[0])
+  size = ops.convert_to_tensor(size, dtypes.int32, "size")
+  size_as_shape = tensor_util.constant_value_as_shape(size).with_rank(2)
+  out_shape = size_as_shape + images.shape[-1:]
+  out_spec = tensor_spec.TensorSpec(out_shape, dtypes.float32)
+
+  def resize_one(image):
+    if isinstance(image, ragged_tensor.RaggedTensor):
+      image = image.to_tensor()
+    return resize_op(image, size, **kwargs)
+
+  def resize_with_map():
+    return map_fn.map_fn_v2(resize_one, images, fn_output_signature=out_spec)
+
+  def empty_result():
+    channels = array_ops.shape(images.flat_values)[-1:]
+    return array_ops.zeros(array_ops.concat([[0], size, channels], axis=0))
+
+  if static_batch_size == 0:
+    return empty_result()
+  elif static_batch_size is not None:
+    return resize_with_map()
+  else:
+    empty_batch = math_ops.equal(images.nrows(), 0)
+    return cond.cond(empty_batch, empty_result, resize_with_map)

tensorflow/python/ops/ragged/ragged_map_flat_values_op_test.py

@@ -0,0 +1,236 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_functional_ops.map_flat_values."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedMapInnerValuesOpTest(test_util.TensorFlowTestCase):
+
+  def assertRaggedMapInnerValuesReturns(self,
+                                        op,
+                                        expected,
+                                        args=(),
+                                        kwargs=None):
+    kwargs = kwargs or {}
+    result = ragged_functional_ops.map_flat_values(op, *args, **kwargs)
+    self.assertAllEqual(result, expected)
+
+  def testDocStringExamples(self):
+    """Test the examples in apply_op_to_ragged_values.__doc__."""
+    rt = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5], [6]])
+    v1 = ragged_functional_ops.map_flat_values(array_ops.ones_like, rt)
+    v2 = ragged_functional_ops.map_flat_values(math_ops.multiply, rt, rt)
+    v3 = ragged_functional_ops.map_flat_values(math_ops.add, rt, 5)
+    self.assertAllEqual(v1, [[1, 1, 1], [], [1, 1], [1]])
+    self.assertAllEqual(v2, [[1, 4, 9], [], [16, 25], [36]])
+    self.assertAllEqual(v3, [[6, 7, 8], [], [9, 10], [11]])
+
+  def testOpWithSingleRaggedTensorArg(self):
+    tensor = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=array_ops.zeros_like,
+        args=(tensor,),
+        expected=[[0, 0, 0], [], [0, 0]])
+
+  def testOpWithTwoRaggedTensorArgs(self):
+    x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+    y = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply, args=(x, y), expected=[[3, 2, 12], [], [4, 25]])
+
+  def testOpWithRaggedTensorAndScalarArgs(self):
+    y = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply, args=(5, y), expected=[[5, 10, 15], [], [20, 25]])
+
+  def testOpWithThreeRaggedTensorArgs(self):
+    condition = ragged_factory_ops.constant(
+        [[True, True, False], [], [True, False]])  # pyformat: disable
+    x = ragged_factory_ops.constant([['a', 'b', 'c'], [], ['d', 'e']])
+    y = ragged_factory_ops.constant([['A', 'B', 'C'], [], ['D', 'E']])
+    self.assertRaggedMapInnerValuesReturns(
+        op=array_ops.where_v2,
+        args=(condition, x, y),
+        expected=[[b'a', b'b', b'C'], [], [b'd', b'E']])
+
+  def testOpWithRaggedTensorListArg(self):
+    x = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+    y = ragged_factory_ops.constant([[10, 20, 30], [], [40, 50]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.add_n,
+        args=([x, y, x],),
+        expected=[[12, 24, 36], [], [48, 60]])
+
+  def testOpWithKeywordArgs(self):
+    x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+    y = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply,
+        kwargs=dict(x=x, y=y),
+        expected=[[3, 2, 12], [], [4, 25]])
+
+  def testOpWithMixedPositionalAndKeywordArgs(self):
+    x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+    y = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply,
+        args=(x,),
+        kwargs=dict(y=y),
+        expected=[[3, 2, 12], [], [4, 25]])
+
+  def testNonElementWiseOp(self):
+    x = ragged_factory_ops.constant(
+        [[[3, 1, 4], [1, 5, 9], [2, 6, 5]], [], [[3, 5, 8], [9, 7, 9]]],
+        ragged_rank=1)
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.reduce_sum,
+        kwargs={
+            'input_tensor': x,
+            'axis': 1,
+        },
+        expected=[[8, 15, 13], [], [16, 25]])
+
+  def testOpWithRaggedRankGreaterThanOne(self):
+    # ragged_rank=0
+    x0 = [3, 1, 4, 1, 5, 9, 2, 6, 5]
+    y0 = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+    self.assertAllEqual(
+        math_ops.multiply(x0, y0), [3, 2, 12, 4, 25, 54, 14, 48, 45])
+
+    # ragged_rank=1
+    x1 = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5], [9, 2], [6, 5]])
+    y1 = ragged_factory_ops.constant([[1, 2, 3], [], [4, 5], [6, 7], [8, 9]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply,
+        args=(x1, y1),
+        expected=[[3, 2, 12], [], [4, 25], [54, 14], [48, 45]])
+
+    # ragged_rank=2
+    x2 = ragged_factory_ops.constant([[[3, 1, 4]], [], [[], [1, 5]],
+                                      [[9, 2], [6, 5]]])
+    y2 = ragged_factory_ops.constant([[[1, 2, 3]], [], [[], [4, 5]],
+                                      [[6, 7], [8, 9]]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply,
+        args=(x2, y2),
+        expected=[[[3, 2, 12]],          # row 0
+                  [],                    # row 1
+                  [[], [4, 25]],         # row 2
+                  [[54, 14], [48, 45]]   # row 3
+                 ])  # pyformat: disable
+
+    # ragged_rank=3
+    x3 = ragged_factory_ops.constant([[[[3, 1, 4]], []], [], [[[], [1, 5]]],
+                                      [[[9, 2], [6, 5]]]])
+    y3 = ragged_factory_ops.constant([[[[1, 2, 3]], []], [], [[[], [4, 5]]],
+                                      [[[6, 7], [8, 9]]]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply,
+        args=(x3, y3),
+        expected=[
+            [[[3, 2, 12]], []],       # row 0
+            [],                       # row 1
+            [[[], [4, 25]]],          # row 2
+            [[[54, 14], [48, 45]]]    # row 3
+        ])  # pyformat: disable
+
+  def testOpWithRaggedRankThree(self):
+    x = ragged_factory_ops.constant([[[3, 1, 4]], [], [[], [1, 5]]])
+    y = ragged_factory_ops.constant([[[1, 2, 3]], [], [[], [4, 5]]])
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply,
+        args=(x, y),
+        expected=[[[3, 2, 12]], [], [[], [4, 25]]])
+
+  def testOpWithInnerValuesOnly(self):
+    x = constant_op.constant([[1, 2], [3, 4], [5, 6]])
+    y = constant_op.constant(2)
+    self.assertRaggedMapInnerValuesReturns(
+        op=math_ops.multiply, args=(x, y), expected=[[2, 4], [6, 8], [10, 12]])
+
+  def testRaggedTensorSplitsRaggedRankMismatchError(self):
+    x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+    y = ragged_factory_ops.constant([[[3, 1, 4], []], [], [[1, 5]]])
+    with self.assertRaisesRegex(
+        ValueError, r'All ragged inputs must have the same ragged_rank.'):
+      ragged_functional_ops.map_flat_values(math_ops.add, x, y)
+
+  def testRaggedTensorSplitsValueMismatchError(self):
+    x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
+    y = ragged_factory_ops.constant([[1], [2, 3], [4, 5]])
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                r'partitions have incompatible'):
+      ragged_functional_ops.map_flat_values(math_ops.add, x, y)
+
+    z_splits = array_ops.placeholder_with_default(
+        constant_op.constant([0, 3], dtypes.int64), None)
+    z = ragged_tensor.RaggedTensor.from_row_splits([0, 1, 2], z_splits)
+    with self.assertRaisesRegex(
+        ValueError,
+        r"Input RaggedTensors' flat_values must all have the same "
+        r'outer-dimension size.  Got sizes: \{3, 5\}'):
+      ragged_functional_ops.map_flat_values(math_ops.add, x, z)
+
+  def testRaggedTensorShapeMismatchError(self):
+    x = ragged_factory_ops.constant([[1, 2, 3], [4, 5]])
+    with self.assertRaisesRegex(
+        ValueError, r'tf.ragged.map_flat_values requires that the output of '
+        '`op` have the same outer-dimension size as flat_values of any ragged '
+        r'inputs. \(output shape: \(\); expected outer dimension size: 5\)'):
+      ragged_functional_ops.map_flat_values(math_ops.argmax, x)
+
+  def testRaggedTensorSplitsMismatchErrorAtRuntime(self):
+    splits1 = array_ops.placeholder_with_default(
+        constant_op.constant([0, 3, 3, 5], dtypes.int64), None)
+    splits2 = array_ops.placeholder_with_default(
+        constant_op.constant([0, 1, 3, 5], dtypes.int64), None)
+    x = ragged_tensor.RaggedTensor.from_row_splits([3, 1, 4, 1, 5], splits1)
+    y = ragged_tensor.RaggedTensor.from_row_splits([1, 2, 3, 4, 5], splits2)
+    with self.assertRaisesRegex((ValueError, errors.InvalidArgumentError),
+                                r'partitions have incompatible'):
+      self.evaluate(ragged_functional_ops.map_flat_values(math_ops.add, x, y))
+
+  def testRaggedMapFnPreservesUniformRowLength(self):
+    # x and y are equal, except that x has uniform_row_length and y does not.
+    x = ragged_tensor.RaggedTensor.from_uniform_row_length(
+        ragged_factory_ops.constant([[1, 2], [3]]), uniform_row_length=2)
+    y = ragged_factory_ops.constant([[[1, 2], [3]]])
+
+    a = ragged_functional_ops.map_flat_values(math_ops.add, x, y)
+    self.assertAllEqual(x.uniform_row_length, a.uniform_row_length)
+
+    b = ragged_functional_ops.map_flat_values(math_ops.add, y, x)
+    self.assertAllEqual(x.uniform_row_length, b.uniform_row_length)
+
+    c = ragged_functional_ops.map_flat_values(math_ops.add_n, [x, x])
+    self.assertAllEqual(x.uniform_row_length, c.uniform_row_length)
+
+    d = ragged_functional_ops.map_flat_values(math_ops.add_n, [y, x, y])
+    self.assertAllEqual(x.uniform_row_length, d.uniform_row_length)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_map_fn_op_test.py

@@ -0,0 +1,337 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_map_ops.map_fn."""
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import map_fn as map_fn_lib
+from tensorflow.python.ops import math_ops as mo
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_map_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+def _stack_mean_and_sum(x):
+  return array_ops_stack.stack([mo.reduce_mean(x), mo.reduce_sum(x)])
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedMapOpTest(test_util.TensorFlowTestCase,
+                      parameterized.TestCase):
+
+  @parameterized.parameters([
+      # The following test sets map over a RaggedTensor and apply a
+      # transformation that returns with shape:
+      # [d1, (d2)] -> [d1]
+      dict(
+          fn=mo.reduce_mean,
+          elems=[[1, 2, 3], [4, 5], [6, 7]],
+          elems_dtype=dtypes.int32,
+          expected_output=[2, 4, 6],
+          result_dtype=dtypes.int32,
+      ),
+      dict(
+          fn=string_ops.reduce_join,
+          elems=[['foo', 'bar', 'baz'], ['a'], ['b', 'c']],
+          expected_output=[b'foobarbaz', b'a', b'bc'],
+          elems_dtype=dtypes.string,
+          result_dtype=dtypes.string,
+      ),
+      # [d1, (d2)] -> [d1, 2]
+      dict(
+          fn=_stack_mean_and_sum,
+          elems=[[1, 2, 3], [4, 5], [6, 7]],
+          expected_output=[[2, 6], [4.5, 9], [6.5, 13]],
+          elems_dtype=dtypes.float32,
+          result_dtype=dtypes.float32,
+          expected_ragged_rank=0,
+      ),
+      # [d1, (d2)] -> [d1, (d2)]
+      dict(
+          fn=lambda x: x + np.int64(1),
+          elems=[[1, 2, 3], [4, 5], [6, 7]],
+          expected_output=[[2, 3, 4], [5, 6], [7, 8]],
+          elems_dtype=dtypes.int64,
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=1),
+      ),
+      # [d1, (d2), d3] -> [d1, (d2), d3]
+      dict(
+          fn=lambda x: x + np.int64(1),
+          elems=[[[1, 2], [3, 4]], [], [[5, 6], [7, 8], [9, 0]]],
+          elems_ragged_rank=1,
+          expected_ragged_rank=1,
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=1),
+          expected_output=[[[2, 3], [4, 5]], [], [[6, 7], [8, 9], [10, 1]]],
+      ),
+      # [d1, (d2)] -> [d1, (d2), (d3)]
+      dict(
+          fn=lambda x: ragged_tensor.RaggedTensor.from_row_starts(x, [0]),
+          elems=[[1, 2, 3], [4, 5], [6, 7]],
+          expected_output=[[[1, 2, 3]], [[4, 5]], [[6, 7]]],
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=2),
+      ),
+      # [d1, (d2), (d3)] -> [d1, (d2), (d3)]
+      dict(
+          fn=lambda x: ragged_functional_ops.map_flat_values(mo.add, x, 1),
+          elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
+          expected_output=[[[2, 3, 4]], [[5, 6], [7, 8]]],
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=2),
+      ),
+      # [d1, (d2), (d3)] -> [d1, (d2)]
+      dict(
+          fn=lambda x: ragged_math_ops.reduce_sum(x, axis=1),
+          elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
+          expected_output=[[6], [9, 13]],
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=1),
+      ),
+      # [d1, (d2), (d3)] -> [d1, (d3)]
+      dict(
+          fn=lambda x: ragged_math_ops.reduce_sum(x, axis=0),
+          elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
+          expected_output=[[1, 2, 3], [10, 12]],
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=1),
+      ),
+      # [d1, (d2), (d3)] -> [d1]
+      dict(
+          fn=ragged_math_ops.reduce_sum,
+          elems=[[[1, 2, 3]], [[4, 5], [6, 7]]],
+          expected_output=[6, 22],
+          result_dtype=dtypes.int64,
+      ),
+      # [d1] -> [d1, (d2)]
+      dict(
+          fn=mo.range,
+          elems=[4, 0, 2],
+          expected_output=[[0, 1, 2, 3], [], [0, 1]],
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=1),
+      ),
+      # [d1] -> [d1, (d2), (d3)]
+      dict(
+          fn=lambda x: ragged_math_ops.range(mo.range(x)),
+          elems=[5, 0, 3],
+          expected_output=[[[], [0], [0, 1], [0, 1, 2], [0, 1, 2, 3]], [],
+                           [[], [0], [0, 1]]],
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=2),
+      ),
+      # [d1, (d2), (d3), (d4a), (d5)] ->  [d1, (d2), (d3), (d4b), (d5)]
+      dict(
+          fn=lambda x: x + np.int64(1),
+          elems=[[[[[1, 2, 3]], [[4], [5]]]], [[[[6, 7]]], [[[8], []]]]],
+          expected_output=[[[[[2, 3, 4]], [[5], [6]]]], [[[[7, 8]]], [[[9],
+                                                                       []]]]],
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=4),
+      ),
+      # [d1] -> [d1, (d2), (d3)]
+      dict(
+          fn=ragged_math_ops.range,
+          elems=np.array([1, 2, 3], np.int64),
+          expected_output=[[[0]], [[0, 1]], [[0, 1, 2]]],
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=2)),
+      # [0] -> [0, (d2), (d3)]  (github issue #36232)
+      dict(
+          fn=ragged_math_ops.range,
+          elems=np.zeros([0], np.int64),
+          expected_output=[],
+          expected_ragged_rank=2,
+          result_dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=2)),
+  ])
+
+  def testRaggedMap(
+      self,
+      fn,
+      elems,
+      expected_output,
+      expected_ragged_rank=None,
+      result_ragged_rank=None,
+      elems_ragged_rank=None,
+      elems_dtype=dtypes.int64,
+      result_dtype=None,
+      infer_shape=True,
+  ):
+    elems = ragged_factory_ops.constant(elems, elems_dtype, elems_ragged_rank)
+    output = ragged_map_ops.map_fn(
+        fn=fn, elems=elems, dtype=result_dtype, infer_shape=infer_shape)
+
+    expected_rt = ragged_factory_ops.constant(
+        expected_output, ragged_rank=expected_ragged_rank)
+    self.assertAllEqual(expected_rt, output)
+
+  def testRaggedMapOnStructure(self):
+    batman = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6, 7]])
+    # [[10, 20, 30], [40], [50, 60, 70]]
+    robin = ragged_functional_ops.map_flat_values(mo.multiply, batman, 10)
+
+    features = {'batman': batman, 'robin': robin}
+
+    def _reduce_sum_from_all(f):
+      return mo.reduce_sum(f['batman']) + mo.reduce_sum(f['robin'])
+
+    output = ragged_map_ops.map_fn(
+        fn=_reduce_sum_from_all,
+        elems=features,
+        dtype=dtypes.int32,
+    )
+
+    self.assertAllEqual(output, [66, 44, 198])
+
+  # Test mapping over a dict of RTs can produce a dict of RTs.
+  def testRaggedMapOnStructure_RaggedOutputs(self):
+    batman = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6, 7]])
+    # [[10, 20, 30], [40], [50, 60, 70]]
+    robin = ragged_functional_ops.map_flat_values(mo.multiply, batman, 10)
+
+    features = {'batman': batman, 'robin': robin}
+
+    def _increment(f):
+      return {
+          'batman': f['batman'] + 1,
+          'robin': f['robin'] + 1,
+      }
+
+    output = ragged_map_ops.map_fn(
+        fn=_increment,
+        elems=features,
+        infer_shape=False,
+        dtype={
+            'batman':
+                ragged_tensor.RaggedTensorType(
+                    dtype=dtypes.int32, ragged_rank=1),
+            'robin':
+                ragged_tensor.RaggedTensorType(
+                    dtype=dtypes.int32, ragged_rank=1)
+        },
+    )
+
+    self.assertAllEqual(output['batman'], [[2, 3, 4], [5], [6, 7, 8]])
+    self.assertAllEqual(output['robin'], [[11, 21, 31], [41], [51, 61, 71]])
+
+  def testZip(self):
+    x = ragged_factory_ops.constant(
+        [[10, 20], [30, 40], [50, 60], [70], [80, 90, 100]], dtypes.int64)
+    y = array_ops.expand_dims(mo.range(x.nrows(out_type=dtypes.int64)), axis=1)
+
+    def _zip(foo):
+      y_val, x_val = foo
+      bar = array_ops.tile(y_val, array_ops.shape(x_val))
+      return array_ops_stack.stack([bar, x_val], axis=1)
+
+    output = ragged_map_ops.map_fn(
+        _zip, (y, x),
+        dtype=ragged_tensor.RaggedTensorType(dtype=dtypes.int64, ragged_rank=1),
+        infer_shape=False)
+
+    self.assertAllEqual(
+        output, [[[0, 10], [0, 20]], [[1, 30], [1, 40]], [[2, 50], [2, 60]],
+                 [[3, 70]], [[4, 80], [4, 90], [4, 100]]])
+
+  def testBatchGather(self):
+    tokens = ragged_factory_ops.constant([['hello', '.', 'there'], ['merhaba'],
+                                          ['bonjour', '.', 'ca va', '?']])
+    indices = ragged_factory_ops.constant([[0, 2], [0], [0, 2]])
+
+    def gather(x):
+      tokens_val, indices_val = x
+      return array_ops.gather(tokens_val, indices_val)
+
+    data = tokens, indices
+    out = ragged_map_ops.map_fn(
+        gather,
+        data,
+        dtype=ragged_tensor.RaggedTensorType(
+            dtype=dtypes.string, ragged_rank=1),
+        infer_shape=False)
+
+    self.assertAllEqual(
+        out, [[b'hello', b'there'], [b'merhaba'], [b'bonjour', b'ca va']])
+
+  def testMismatchRaggedRank(self):
+    elems = ragged_factory_ops.constant([[[1, 2, 3]], [[4, 5], [6, 7]]])
+    fn = lambda x: ragged_math_ops.reduce_sum(x, axis=0)
+    with self.assertRaisesRegex(
+        ValueError, r'(?s)Expected `fn` to return.*But it returned.*'):
+      _ = ragged_map_ops.map_fn(
+          fn,
+          elems,
+          dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=23))
+
+  def testMismatchRaggedRank2(self):
+    elems = ragged_factory_ops.constant([[1, 2, 3], [4, 5], [6, 7]])
+    fn = lambda x: ragged_tensor.RaggedTensor.from_row_starts(x, [0])
+    with self.assertRaisesRegex(
+        ValueError, r'(?s)Expected `fn` to return.*But it returned.*'):
+      _ = ragged_map_ops.map_fn(
+          fn,
+          elems,
+          dtype=ragged_tensor.RaggedTensorType(
+              dtype=dtypes.int64, ragged_rank=10))
+
+  def testMapOnSparseTensor(self):
+    s = sparse_tensor.SparseTensor(
+        indices=[[0, 0], [0, 1], [1, 0], [1, 1]],
+        values=[0, 5, 0, 4],
+        dense_shape=[2, 2],
+    )
+    t2 = ragged_tensor.RaggedTensor.from_sparse(s)
+    id_t2 = ragged_map_ops.map_fn(
+        lambda x: x, t2,
+    )
+    self.assertAllEqual(id_t2, [[0, 5], [0, 4]])
+
+  def testRaggedMapWithIncorrectFnOutputSignature(self):
+    x = ragged_factory_ops.constant([[1, 2, 3, 4], [1]])
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                'All flat_values must have compatible shapes'):
+      y = map_fn_lib.map_fn(lambda r: map_fn_lib.map_fn(lambda y: r, r), x)
+      self.evaluate(y)
+
+  def testNestedRaggedMapWithFnOutputSignature(self):
+    ragged1d = ragged_tensor.RaggedTensorSpec([None], dtypes.int32)
+    ragged2d = ragged_tensor.RaggedTensorSpec([None, None], dtypes.int32)
+
+    x = ragged_factory_ops.constant([[1, 2, 3, 4], [1]])
+    # pylint: disable=g-long-lambda
+    y = map_fn_lib.map_fn(
+        lambda r: map_fn_lib.map_fn(
+            lambda y: r, r, fn_output_signature=ragged1d),
+        x,
+        fn_output_signature=ragged2d)
+    expected = [[[1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4], [1, 2, 3, 4]], [[1]]]
+    self.assertAllEqual(y, expected)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_map_ops.py

@@ -0,0 +1,174 @@
+# 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.
+# ==============================================================================
+"""Functional operations for RaggedTensors."""
+
+from tensorflow.python.ops import map_fn as map_fn_lib
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import nest
+
+
+def map_fn(fn,
+           elems,
+           dtype=None,
+           parallel_iterations=None,
+           back_prop=True,
+           swap_memory=False,
+           infer_shape=True,
+           name=None):
+  """map on the list of tensors unpacked from `elems` on dimension 0.
+
+  The simplest version of `map_fn` repeatedly applies the callable `fn` to a
+  sequence of elements from first to last. The elements are made of the
+  tensors unpacked from `elems`. `dtype` is the data type of the return
+  value of `fn`. Users must provide `dtype` if it is different from
+  the data type of `elems`.
+
+  Suppose that `elems` is unpacked into `values`, a list of tensors. The shape
+  of the result tensor is `[values.shape[0]] + fn(values[0]).shape`.
+
+  This method also allows multi-arity `elems` and output of `fn`.  If `elems`
+  is a (possibly nested) list or tuple of tensors, then each of these tensors
+  must have a matching first (unpack) dimension.  The signature of `fn` may
+  match the structure of `elems`.  That is, if `elems` is
+  `(t1, [t2, t3, [t4, t5]])`, then an appropriate signature for `fn` is:
+  `fn = lambda (t1, [t2, t3, [t4, t5]]):`.
+
+  Furthermore, `fn` may emit a different structure than its input.  For example,
+  `fn` may look like: `fn = lambda t1: return (t1 + 1, t1 - 1)`.  In this case,
+  the `dtype` parameter is not optional: `dtype` must be a type or (possibly
+  nested) tuple of types matching the output of `fn`.
+
+  To apply a functional operation to the nonzero elements of a SparseTensor
+  one of the following methods is recommended. First, if the function is
+  expressible as TensorFlow ops, use
+
+  ```python
+    result = SparseTensor(input.indices, fn(input.values), input.dense_shape)
+  ```
+
+  If, however, the function is not expressible as a TensorFlow op, then use
+
+  ```python
+  result = SparseTensor(
+    input.indices, map_fn(fn, input.values), input.dense_shape)
+  ```
+
+  instead.
+
+  When executing eagerly, map_fn does not execute in parallel even if
+  `parallel_iterations` is set to a value > 1. You can still get the
+  performance benefits of running a function in parallel by using the
+  `tf.contrib.eager.defun` decorator,
+
+  ```python
+  # Assume the function being used in map_fn is fn.
+  # To ensure map_fn calls fn in parallel, use the defun decorator.
+  @tf.contrib.eager.defun
+  def func(tensor):
+    return tf.map_fn(fn, tensor)
+  ```
+
+  Note that if you use the defun decorator, any non-TensorFlow Python code
+  that you may have written in your function won't get executed. See
+  `tf.contrib.eager.defun` for more details. The recommendation would be to
+  debug without defun but switch to defun to get performance benefits of
+  running map_fn in parallel.
+
+  Args:
+    fn: The callable to be performed.  It accepts one argument, which will have
+      the same (possibly nested) structure as `elems`.  Its output must have the
+      same structure as `dtype` if one is provided, otherwise it must have the
+      same structure as `elems`.
+    elems: A tensor or (possibly nested) sequence of tensors, each of which will
+      be unpacked along their first dimension.  The nested sequence of the
+      resulting slices will be applied to `fn`.
+    dtype: (optional) The output type(s) of `fn`.  If `fn` returns a structure
+      of Tensors differing from the structure of `elems`, then `dtype` is not
+      optional and must have the same structure as the output of `fn`. Use
+      `RaggedTensorType` to declare an output of type `RaggedTensor`.
+    parallel_iterations: (optional) The number of iterations allowed to run in
+      parallel. When graph building, the default value is 10. While executing
+      eagerly, the default value is set to 1.
+    back_prop: (optional) True enables support for back propagation.
+    swap_memory: (optional) True enables GPU-CPU memory swapping.
+    infer_shape: (optional) False disables tests for consistent output shapes.
+    name: (optional) Name prefix for the returned tensors.
+
+  Returns:
+    A possibly nested sequence of potentially ragged tensors.  Each
+    tensor packs the results of applying `fn` to tensors unpacked from `elems`
+    along the first dimension, from first to last.
+
+  Raises:
+    TypeError: if `fn` is not callable or the structure of the output of
+      `fn` and `dtype` do not match, or if elems is a SparseTensor.
+    ValueError: if the lengths of the output of `fn` and `dtype` do not match.
+
+  #### Examples:
+
+    ```python
+    elems = np.array([1, 2, 3, 4, 5, 6])
+    squares = map_fn(lambda x: x * x, elems)
+    # squares == [1, 4, 9, 16, 25, 36]
+    ```
+
+    ```python
+    elems = (np.array([1, 2, 3]), np.array([-1, 1, -1]))
+    alternate = map_fn(lambda x: x[0] * x[1], elems, dtype=tf.int64)
+    # alternate == [-1, 2, -3]
+    ```
+
+    ```python
+    elems = np.array([1, 2, 3])
+    alternates = map_fn(lambda x: (x, -x), elems, dtype=(tf.int64, tf.int64))
+    # alternates[0] == [1, 2, 3]
+    # alternates[1] == [-1, -2, -3]
+    ```
+
+    ```python
+    elems=ragged.constant([[1, 2, 3], [4, 5], [6, 7]])
+    mean = map_fn(tf.reduce_mean, elems)
+    # mean == [2, 4, 6]
+    ```
+
+    ```python
+    elems=ragged.constant([[1, 2, 3], [4, 5], [6, 7]], dtype=tf.int64)
+    out = map_fn(fn=lambda x: x+1, elems,
+      dtype=ragged.RaggedTensorType(type=tf.int64, ragged_rank=0))
+    # out = tf.ragged.constant([[2, 3, 4], [5, 6], [7, 8]])
+    ```
+  """
+  if dtype is None:
+    dtype = nest.map_structure(lambda e: e.dtype, elems)
+  dtype = nest.map_structure(_ragged_type_to_spec, dtype)
+  return map_fn_lib.map_fn(fn,
+                           elems,
+                           dtype,
+                           parallel_iterations,
+                           back_prop,
+                           swap_memory,
+                           infer_shape,
+                           name)
+
+
+def _ragged_type_to_spec(t):
+  if isinstance(t, ragged_tensor.RaggedTensorType):
+    # Note: need to adjust ragged_rank by 1, since RaggedTensorSpec gives the
+    # type for the mapped `fn` output, but RaggedTensorType gives the type for
+    # the result of stacking the mapped `fn` outputs.
+    return ragged_tensor.RaggedTensorSpec(
+        None, t.dtype, t.ragged_rank - 1, t.row_splits_dtype)
+  else:
+    return t

tensorflow/python/ops/ragged/ragged_math_ops_test.py

@@ -0,0 +1,181 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Tests functionality of math operations on ragged tensors."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_string_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+class RaggedReduceTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      [dict(original=['a b'.split(), 'c d e'.split()], expected='a b c d e')])
+  @test_util.run_in_graph_and_eager_modes
+  def testStringReduceJoin(self, original, expected, separator=' ', axis=None):
+    original_rt = ragged_factory_ops.constant(original)
+    expected_rt = ragged_factory_ops.constant(expected)
+    actual = ragged_string_ops.reduce_join(original_rt, axis=axis,
+                                           separator=separator)
+    self.assertAllEqual(actual, expected_rt)
+
+
+class RaggedSoftmaxTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  def _softmax(self, x):
+    assert len(x.shape) == 2
+    if x.shape[1] == 0:
+      return x
+    m = x.max(1)[:, np.newaxis]
+    u = np.exp(x - m)
+    z = u.sum(1)[:, np.newaxis]
+    return u / z
+
+  @test_util.run_in_graph_and_eager_modes
+  def testOrdinaryValues(self):
+    eps = 1e-5
+    x_list = [np.log([0.5, 0.25, 0.25]), np.log([0.5, 0.5])]
+    x_row_matrices = [[row] for row in x_list]
+    y_row_matrices = [
+        self._softmax(np.array(row_matrix)).tolist()
+        for row_matrix in x_row_matrices
+    ]
+    y_list = [row_matrix[0] for row_matrix in y_row_matrices]
+    y_expected_from_numpy = ragged_factory_ops.constant(
+        y_list, dtype=dtypes.float32)
+    y_expected = ragged_factory_ops.constant([[0.5, 0.25, 0.25], [0.5, 0.5]],
+                                             dtype=dtypes.float32)
+    self.assertAllClose(y_expected_from_numpy, y_expected, eps)
+    x_tf = ragged_factory_ops.constant(x_list, dtype=dtypes.float32)
+    y_tf = nn_ops.softmax_v2(x_tf)
+    self.assertAllClose(y_tf, y_expected_from_numpy, eps)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testLargeValues(self):
+    eps = 1e-5
+    x_list = [[-500, -501, -502], [1729, 1729]]
+    x_row_matrices = [[row] for row in x_list]
+    y_row_matrices = [
+        self._softmax(np.array(row_matrix)).tolist()
+        for row_matrix in x_row_matrices
+    ]
+    y_list = [row_matrix[0] for row_matrix in y_row_matrices]
+    y_expected_from_numpy = ragged_factory_ops.constant(
+        y_list, dtype=dtypes.float32)
+    x_tf = ragged_factory_ops.constant(x_list, dtype=dtypes.float32)
+    y_tf = nn_ops.softmax_v2(x_tf)
+    self.assertAllClose(y_tf, y_expected_from_numpy, eps)
+
+  @test_util.run_in_graph_and_eager_modes
+  def testShortTensors(self):
+    eps = 1e-5
+    x_list = [[], [1]]
+    x_row_matrices = [[row] for row in x_list]
+    y_row_matrices = [
+        self._softmax(np.array(row_matrix)).tolist()
+        for row_matrix in x_row_matrices
+    ]
+    y_list = [row_matrix[0] for row_matrix in y_row_matrices]
+    y_expected_from_numpy = ragged_factory_ops.constant(
+        y_list, dtype=dtypes.float32)
+    x_tf = ragged_factory_ops.constant(x_list, dtype=dtypes.float32)
+    y_tf = nn_ops.softmax_v2(x_tf)
+    self.assertAllClose(y_tf, y_expected_from_numpy, eps)
+
+
+def _cumsum_slow(rt, axis=0, exclusive=False, reverse=False, name=None):
+  dense = rt.to_tensor()
+  result = math_ops.cumsum(dense, axis=axis, exclusive=exclusive,
+                           reverse=reverse, name=name)
+  return ragged_tensor.RaggedTensor.from_tensor(
+      result, lengths=rt.nested_row_lengths())
+
+
+class RaggedCumsumTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[1, 2], [4, 6, 5]]),
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[4, 6], [3, 4, 5]],
+           reverse=True),
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[3, 4], [0, 0, 0]],
+           reverse=True,
+           exclusive=True),
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[0, 0], [1, 2, 0]],
+           exclusive=True),
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[0, 0], [1, 2, 0]],
+           exclusive=True),
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[1, 3], [3, 7, 12]],
+           axis=1),
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[3, 2], [12, 9, 5]],
+           axis=1, reverse=True),
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[2, 0], [9, 5, 0]],
+           axis=1, exclusive=True, reverse=True),
+      dict(original=[[1, 2], [3, 4, 5]],
+           expected=[[0, 1], [0, 3, 7]],
+           axis=1, exclusive=True),
+  ])
+  def test_cumsum(self, original, expected, axis=0, exclusive=False,
+                  reverse=False):
+    original_rt = ragged_factory_ops.constant(original)
+    expected_rt = ragged_factory_ops.constant(expected)
+    actual = ragged_math_ops.ragged_cumsum(
+        original_rt, axis=axis, exclusive=exclusive, reverse=reverse)
+    self.assertAllEqual(actual, expected_rt)
+    baseline = _cumsum_slow(original_rt, axis=axis, exclusive=exclusive,
+                            reverse=reverse)
+    self.assertAllEqual(baseline, expected_rt)
+
+  @parameterized.parameters([
+      dict(expected=[[[0, 1], [2, 5]], [[4, 9], [6, 13], [8, 17]]], axis=2),
+      dict(expected=[[[0, 0], [0, 2]], [[0, 4], [0, 6], [0, 8]]],
+           exclusive=True, axis=2),
+      dict(expected=[[[1, 0], [3, 0]], [[5, 0], [7, 0], [9, 0]]],
+           exclusive=True, reverse=True, axis=2),
+      dict(expected=[[[1, 1], [5, 3]], [[9, 5], [13, 7], [17, 9]]],
+           reverse=True, axis=2),
+  ])
+  def test_cumsum_deep(self, expected, axis=0, exclusive=False, reverse=False):
+    # [[[0, 1], [2, 3]], [[4, 5], [6, 7], [8, 9]]]
+
+    original_rt = ragged_tensor.RaggedTensor.from_row_lengths(
+        array_ops.reshape(math_ops.range(10), (5, 2)), [2, 3])
+
+    actual = ragged_math_ops.ragged_cumsum(
+        original_rt, axis=axis, exclusive=exclusive, reverse=reverse)
+    self.assertAllEqual(actual, expected)
+    baseline = _cumsum_slow(original_rt, axis=axis, exclusive=exclusive,
+                            reverse=reverse)
+    self.assertAllEqual(baseline, expected)
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_matmul_op_test.py

@@ -0,0 +1,289 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Tests for tf.ragged.cross and tf.ragged.matmul."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+def T(shape):
+  """Make a dense test Tensor with the indicated shape."""
+  # Values are arbitrary, but make them nontrivial because we use them to
+  # test matmul against a reference implementation.
+  values = math_ops.range(math_ops.reduce_prod(shape))
+  return array_ops.reshape(values, shape)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedMatmulOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  def eager_ragged_matmul(self, a, b, **kwargs):
+    """Reference implementation for ragged matmul."""
+    if len(a.shape) > 2:
+      return [
+          self.eager_ragged_matmul(a[i], b[i], **kwargs)
+          for i in range(a.shape[0])
+      ]
+    a = self.ensure_non_ragged(a)
+    b = self.ensure_non_ragged(b)
+    return self.evaluate(math_ops.matmul(a, b, **kwargs)).tolist()
+
+  def ensure_non_ragged(self, x):
+    """Returns x as a Tensor.  Fails if x contains ragged rows."""
+    if not isinstance(x, ragged_tensor.RaggedTensor):
+      return x
+    x_uniform = x.to_tensor()
+    self.assertAllEqual(array_ops.size(x), array_ops.size(x_uniform))
+    return x_uniform
+
+  # pylint: disable=g-long-lambda
+  @parameterized.named_parameters([
+      # The test names below have the form '<a.shape>_times_<b.shape>'.
+      # Within <a.shape> and <b.shape>, constants are used for dense dimensions
+      # and letters (e.g. B, I, J, and K) are used for raggged dimensions.
+      dict(  # Uses math_ops.matmul
+          testcase_name='dense',
+          a=lambda: T([3, 4, 5]),
+          b=lambda: T([3, 5, 6]),
+          expected_shape=[3, 4, 6]),
+      dict(  # Uses matmul_2d
+          testcase_name='2x3_times_3x1',
+          a=lambda: ragged_factory_ops.constant([[1, 2, 3], [4, 5, 6]]),
+          b=lambda: ragged_factory_ops.constant([[5], [4], [3]]),
+          expected_shape=[2, None]),
+      dict(  # Uses matmul_3d_with_map_fn
+          testcase_name='2xIxJ_times_2xJxK',
+          a=lambda: ragged_concat_ops.stack([T([15, 32]),
+                                             T([10, 20])]),
+          b=lambda: ragged_concat_ops.stack([T([32, 19]),
+                                             T([20, 13])]),
+          expected_shape=[2, None, None]),
+      dict(  # Uses matmul_3d_with_map_fn
+          testcase_name='2xIxJ_times_2xJx12',
+          a=lambda: ragged_concat_ops.stack([T([15, 4]), T([10, 2])]),
+          b=lambda: ragged_factory_ops.constant([[[1, 2], [3, 4], [5, 6],
+                                                  [7, 8]], [[9, 10], [11, 12]]],
+                                                ragged_rank=1),
+          expected_shape=[2, None, 2]),
+      dict(  # Uses matmul_3d_with_map_fn
+          testcase_name='2xIx8_times_2x8x12',
+          a=lambda: ragged_concat_ops.stack([T([15, 8]), T([10, 8])]),
+          b=lambda: T([2, 8, 12]),
+          expected_shape=[2, None, 12]),
+      dict(  # Uses matmul_3d_with_map_fn
+          testcase_name='2x15x32_times_2x32xK',
+          a=lambda: T([2, 15, 32]),
+          b=lambda: ragged_concat_ops.stack([T([32, 19]),
+                                             T([32, 13])])),
+      dict(  # Uses matmul_3d_with_batch_dim_folding
+          testcase_name='2xIx3_times_2x3x8',
+          a=lambda: ragged_factory_ops.constant(
+              [[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [8, 7, 6], [5, 4, 3]]],
+              ragged_rank=1),
+          b=lambda: T([2, 3, 8]),
+          expected_shape=[2, None, 8]),
+      dict(  # Multiple batch dimensions
+          testcase_name='3xBx5x7_times_3xBx7x9',
+          a=lambda: ragged_tensor.RaggedTensor.from_row_lengths(
+              values=T([10, 5, 7]), row_lengths=[3, 2, 0, 5]),
+          b=lambda: ragged_tensor.RaggedTensor.from_row_lengths(
+              values=T([10, 7, 9]), row_lengths=[3, 2, 0, 5])),
+      dict(  # Multiple batch dimensions
+          testcase_name='3xBx5x7_times_3x2x7x9',
+          a=lambda: ragged_tensor.RaggedTensor.from_tensor(
+              T([3, 2, 5, 7]), ragged_rank=1),
+          b=lambda: T([3, 2, 7, 9])),
+      dict(  # Multiple batch dimensions
+          testcase_name='3xBxIx7_times_3x2x7x9',
+          a=lambda: ragged_tensor.RaggedTensor.from_tensor(
+              T([3, 2, 5, 7]), ragged_rank=2),
+          b=lambda: T([3, 2, 7, 9])),
+      dict(  # Multiple batch dimensions
+          testcase_name='3xBxIxJ_times_3x2x7x9',
+          a=lambda: ragged_tensor.RaggedTensor.from_tensor(
+              T([3, 2, 5, 7]), ragged_rank=3),
+          b=lambda: T([3, 2, 7, 9])),
+      dict(  # Multiple batch dimensions
+          testcase_name='3x2x5x7_times_3xBx7x9',
+          a=lambda: T([3, 2, 5, 7]),
+          b=lambda: ragged_tensor.RaggedTensor.from_tensor(
+              T([3, 2, 7, 9]), ragged_rank=1)),
+      dict(  # Multiple batch dimensions
+          testcase_name='3x2x5x7_times_3xBxJx9',
+          a=lambda: T([3, 2, 5, 7]),
+          b=lambda: ragged_tensor.RaggedTensor.from_tensor(
+              T([3, 2, 7, 9]), ragged_rank=2)),
+      dict(  # Multiple batch dimensions
+          testcase_name='3x2x5x7_times_3xBxJxK',
+          a=lambda: T([3, 2, 5, 7]),
+          b=lambda: ragged_tensor.RaggedTensor.from_tensor(
+              T([3, 2, 7, 9]), ragged_rank=3)),
+      dict(
+          testcase_name='2x3xI_times_2x3x4_transpose_a',
+          a=lambda: ragged_factory_ops.constant(
+              [[[1], [2], [3]], [[4, 5, 6], [7, 8, 9], [10, 11, 12]]]),
+          b=lambda: T([2, 3, 4]),
+          transpose_a=True,
+          expected_shape=[2, None, 4]),
+      dict(
+          testcase_name='2x3xI_times_2x4_3_transpose_a_transpose_b',
+          a=lambda: ragged_factory_ops.constant(
+              [[[1], [2], [3]], [[4, 5, 6], [7, 8, 9], [10, 11, 12]]]),
+          b=lambda: T([2, 4, 3]),
+          transpose_a=True,
+          transpose_b=True,
+          expected_shape=[2, None, 4]),
+      dict(
+          testcase_name='2xIxJ_times_2x5xJ_transpose_b',
+          a=lambda: ragged_factory_ops.constant([[[1, 2], [3, 4], [5, 6]],
+                                                 [[1, 2, 3], [4, 5, 6]]]),
+          b=lambda: ragged_factory_ops.constant([[[3, 1], [4, 1], [5, 9], [
+              1, 2
+          ], [3, 4]], [[2, 4, 6], [1, 3, 5], [7, 8, 9], [1, 2, 3], [3, 2, 1]]]),
+          transpose_b=True,
+          expected_shape=[2, None, 5]),
+      dict(
+          testcase_name='2xIx3_times_2xJx3_transpose_b',
+          a=lambda: ragged_factory_ops.constant([
+              [[1., 2., 3.], [3., 4., 5.]],
+              [[1., 3., 5.], [5., 7., 9.], [9., 11., 13.]]],
+                                                ragged_rank=1),
+          b=lambda: ragged_factory_ops.constant([
+              [[10., 20., 30.], [30., 40., 50.],
+               [50., 60., 70.], [70., 80., 90.]],
+              [[11., 21., 31.]]], ragged_rank=1),
+          transpose_b=True,
+          expected_shape=[2, None, None]),
+      dict(
+          testcase_name='2x2x3_times_2xIx3_transpose_b',
+          a=lambda: constant_op.constant([
+              [[1., 2., 3.], [3., 4., 5.]], [[1., 3., 5.], [5., 7., 9.]]]),
+          b=lambda: ragged_factory_ops.constant([
+              [[10., 20., 30.], [30., 40., 50.],
+               [50., 60., 70.], [70., 80., 90.]],
+              [[11., 21., 31.]]], ragged_rank=1),
+          transpose_b=True,
+          expected_shape=[2, None, None]),
+  ])
+  def testMatmul(self, a, b, expected_shape=None, **kwargs):
+    if callable(a):
+      a = a()
+    if callable(b):
+      b = b()
+    actual = ragged_math_ops.matmul(a, b, **kwargs)
+    expected = self.eager_ragged_matmul(a, b, **kwargs)
+    self.assertAllEqual(actual, expected)
+
+    if expected_shape is not None and not kwargs:
+      if context.executing_eagerly():
+        self.assertTrue(actual.shape.is_compatible_with(expected_shape))
+      else:
+        self.assertEqual(actual.shape.as_list(), expected_shape)
+
+  @parameterized.parameters([
+      dict(
+          a=lambda: ragged_factory_ops.constant([[1, 2, 3], [4, 5]]),
+          b=lambda: ragged_factory_ops.constant([[5], [4], [3]]),
+          exc=errors.InvalidArgumentError,
+          message='The matrices in `a` and `b` may not be ragged in '
+          'their innermost dimension.'),
+      dict(
+          a=lambda: ragged_factory_ops.constant([[1, 2], [4, 5]]),
+          b=lambda: ragged_factory_ops.constant([[5], [4], [3]]),
+          exc=errors.InvalidArgumentError),
+      dict(
+          a=lambda: ragged_concat_ops.stack([T([15, 32]),
+                                             T([10, 20])]),
+          b=lambda: ragged_concat_ops.stack([T([32, 19]),
+                                             T([22, 13])]),
+          exc=errors.InvalidArgumentError),
+      dict(
+          a=[[1]],
+          b=[[1]],
+          transpose_a=True,
+          adjoint_a=True,
+          exc=ValueError,
+          message='Only one of transpose_a and adjoint_a can be True'),
+      dict(
+          a=[[1]],
+          b=[[1]],
+          transpose_b=True,
+          adjoint_b=True,
+          exc=ValueError,
+          message='Only one of transpose_b and adjoint_b can be True'),
+      dict(
+          a=lambda: ragged_factory_ops.constant([[1]]),
+          b=lambda: ragged_factory_ops.constant([[1.0]]),
+          exc=ValueError,
+          message='`a` and `b` must have the same dtype.'),
+      dict(
+          a=lambda: ragged_factory_ops.constant([[1]]),
+          b=lambda: ragged_factory_ops.constant([[[1]]]),
+          exc=ValueError,
+          message='`a` and `b` must have the same rank.'),
+      dict(  # Multiple batch dimensions
+          a=lambda: ragged_tensor.RaggedTensor.from_tensor(T([3, 2, 5, 7])),
+          b=lambda: ragged_tensor.RaggedTensor.from_tensor(T([3, 3, 7, 9])),
+          exc=errors.InvalidArgumentError,
+          message='Batch dimensions of `a` and `b` do not have the same size'),
+  ])
+  def testMatmulError(self, a, b, exc, message=None, **kwargs):
+    if callable(a):
+      a = a()
+    if callable(b):
+      b = b()
+    with self.assertRaisesRegex(exc, message):
+      self.evaluate(ragged_math_ops.matmul(a, b, **kwargs))
+
+  def testUnknownRank(self):
+    no_rank_spec = ragged_tensor.RaggedTensorSpec(None, dtypes.int32, 1)
+    rank_only_spec = ragged_tensor.RaggedTensorSpec([None, None], dtypes.int32,
+                                                    1)
+
+    matmul_no_rank_for_a = def_function.function(
+        input_signature=[rank_only_spec, no_rank_spec])(
+            ragged_math_ops.matmul)
+    matmul_no_rank_for_b = def_function.function(
+        input_signature=[no_rank_spec, rank_only_spec])(
+            ragged_math_ops.matmul)
+    matmul_no_rank_for_a_or_b = def_function.function(
+        input_signature=[no_rank_spec, no_rank_spec])(
+            ragged_math_ops.matmul)
+
+    a = ragged_factory_ops.constant([[1, 2]])
+    b = ragged_factory_ops.constant([[3], [4]])
+    self.assertAllEqual(matmul_no_rank_for_a(a, b), [[11]])
+    self.assertAllEqual(matmul_no_rank_for_b(a, b), [[11]])
+    with self.assertRaisesRegex(
+        ValueError, 'matmul requires at least one input to have known '
+        'rank if either input is ragged.'):
+      matmul_no_rank_for_a_or_b(a, b)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_merge_dims_op_test.py

@@ -0,0 +1,268 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for RaggedTensor.merge_dims."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+from tensorflow.python.util import nest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedMergeDimsOpTest(test_util.TensorFlowTestCase,
+                            parameterized.TestCase):
+
+  @parameterized.named_parameters([
+      {
+          'testcase_name': '2DAxis0To1',
+          'rt': [[1, 2], [], [3, 4, 5]],
+          'outer_axis': 0,
+          'inner_axis': 1,
+          'expected': [1, 2, 3, 4, 5],
+      },
+      {
+          'testcase_name': '3DAxis0To1',
+          'rt': [[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []]],
+          'outer_axis': 0,
+          'inner_axis': 1,
+          'expected': [[1, 2], [], [3, 4, 5], [6], [7, 8], []],
+      },
+      {
+          'testcase_name': '3DAxis1To2',
+          'rt': [[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []]],
+          'outer_axis': 1,
+          'inner_axis': 2,
+          'expected': [[1, 2, 3, 4, 5], [6, 7, 8]],
+      },
+      {
+          'testcase_name': '3DAxis0To2',
+          'rt': [[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []]],
+          'outer_axis': 0,
+          'inner_axis': 2,
+          'expected': [1, 2, 3, 4, 5, 6, 7, 8],
+      },
+      {
+          'testcase_name': '3DAxis0To1WithDenseValues',
+          'rt': [[[1, 2], [3, 4], [5, 6]], [[7, 8]]],
+          'ragged_ranks': (1, 2),
+          'outer_axis': 0,
+          'inner_axis': 1,
+          'expected': [[1, 2], [3, 4], [5, 6], [7, 8]],
+      },
+      {
+          'testcase_name': '3DAxis1To2WithDenseValues',
+          'rt': [[[1, 2], [3, 4], [5, 6]], [[7, 8]]],
+          'ragged_ranks': (1, 2),
+          'outer_axis': 1,
+          'inner_axis': 2,
+          'expected': [[1, 2, 3, 4, 5, 6], [7, 8]],
+      },
+      {
+          'testcase_name': '4DAxis0To1',
+          'rt': [[[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []]], [[[9], [0]]]],
+          'outer_axis': 0,
+          'inner_axis': 1,
+          'expected': [[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []], [[9], [0]]],
+      },
+      {
+          'testcase_name': '4DAxis1To2',
+          'rt': [[[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []]], [[[9], [0]]]],
+          'outer_axis': 1,
+          'inner_axis': 2,
+          'expected': [[[1, 2], [], [3, 4, 5], [6], [7, 8], []], [[9], [0]]],
+      },
+      {
+          'testcase_name': '4DAxis2To3',
+          'rt': [[[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []]], [[[9], [0]]]],
+          'outer_axis': 2,
+          'inner_axis': 3,
+          'expected': [[[1, 2, 3, 4, 5], [6, 7, 8]], [[9, 0]]],
+      },
+      {
+          'testcase_name': '4DAxis1To3',
+          'rt': [[[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []]], [[[9], [0]]]],
+          'outer_axis': 1,
+          'inner_axis': 3,
+          'expected': [[1, 2, 3, 4, 5, 6, 7, 8], [9, 0]],
+      },
+      {
+          'testcase_name': '4DAxis1ToNeg1',
+          'rt': [[[[1, 2], [], [3, 4, 5]], [[6], [7, 8], []]], [[[9], [0]]]],
+          'outer_axis': 1,
+          'inner_axis': -1,
+          'expected': [[1, 2, 3, 4, 5, 6, 7, 8], [9, 0]],
+      },
+      {
+          'testcase_name': '4DAxis1To2WithDenseValues',
+          'rt': [[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], [[[9, 10], [11, 12]]]],
+          'ragged_ranks': (1, 2, 3),
+          'outer_axis': 1,
+          'inner_axis': 2,
+          'expected': [[[1, 2], [3, 4], [5, 6], [7, 8]], [[9, 10], [11, 12]]],
+      },
+      {
+          'testcase_name': '4DAxis2To3WithDenseValues',
+          'rt': [[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], [[[9, 10], [11, 12]]]],
+          'ragged_ranks': (1, 2, 3),
+          'outer_axis': 2,
+          'inner_axis': 3,
+          'expected': [[[1, 2, 3, 4], [5, 6, 7, 8]], [[9, 10, 11, 12]]],
+      },
+      {
+          'testcase_name': '4DAxis1To3WithDenseValues',
+          'rt': [[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], [[[9, 10], [11, 12]]]],
+          'ragged_ranks': (1, 2, 3),
+          'outer_axis': 1,
+          'inner_axis': 3,
+          'expected': [[1, 2, 3, 4, 5, 6, 7, 8], [9, 10, 11, 12]],
+      },
+      {
+          'testcase_name': '5DAxis2To3WithDenseValues',
+          'rt': [[[[[1, 2], [3, 4]]], [[[5, 6], [7, 8]]]],
+                 [[[[9, 10], [11, 12]]]]],
+          'ragged_ranks': (1, 2, 3, 4),
+          'outer_axis': 2,
+          'inner_axis': 3,
+          'expected': [[[[1, 2], [3, 4]], [[5, 6], [7, 8]]],
+                       [[[9, 10], [11, 12]]]],
+      },
+      {
+          'testcase_name': '5DAxis3To4WithDenseValues',
+          'rt': [[[[[1, 2], [3, 4]]], [[[5, 6], [7, 8]]]],
+                 [[[[9, 10], [11, 12]]]]],
+          'ragged_ranks': (1, 2, 3, 4),
+          'outer_axis': 3,
+          'inner_axis': 4,
+          'expected': [[[[1, 2, 3, 4]], [[5, 6, 7, 8]]], [[[9, 10, 11, 12]]]],
+      },
+      {
+          'testcase_name': '5DAxis1To3WithDenseValues',
+          'rt': [[[[[1, 2], [3, 4]]], [[[5, 6], [7, 8]]]],
+                 [[[[9, 10], [11, 12]]]]],
+          'ragged_ranks': (1, 2, 3, 4),
+          'outer_axis': 1,
+          'inner_axis': 3,
+          'expected': [[[1, 2], [3, 4], [5, 6], [7, 8]], [[9, 10], [11, 12]]],
+      },
+      {
+          'testcase_name': 'OuterEqualsInner',
+          'rt': [[1], [2], [3, 4]],
+          'outer_axis': 0,
+          'inner_axis': 0,
+          'expected': [[1], [2], [3, 4]],
+      },
+      {
+          'testcase_name': 'OuterEqualsInnerWithNegativeAxis',
+          'rt': [[1], [2], [3, 4]],
+          'outer_axis': 1,
+          'inner_axis': -1,
+          'expected': [[1], [2], [3, 4]],
+      },
+  ])  # pyformat: disable
+  def testRaggedMergeDims(self,
+                          rt,
+                          outer_axis,
+                          inner_axis,
+                          expected,
+                          ragged_ranks=(None,)):
+    for ragged_rank in ragged_ranks:
+      x = ragged_factory_ops.constant(rt, ragged_rank=ragged_rank)
+
+      # Check basic behavior.
+      actual = x.merge_dims(outer_axis, inner_axis)
+      self.assertAllEqual(expected, actual)
+      if outer_axis >= 0 and inner_axis >= 0:
+        self.assertEqual(actual.shape.rank,
+                         x.shape.rank - (inner_axis - outer_axis))
+
+      # Check behavior with negative axis.
+      if outer_axis >= 0 and inner_axis >= 0:
+        actual_with_neg_axis = x.merge_dims(outer_axis - x.shape.rank,
+                                            inner_axis - x.shape.rank)
+        self.assertAllEqual(expected, actual_with_neg_axis)
+
+      # Check behavior with placeholder input (no shape info).
+      if (not context.executing_eagerly() and outer_axis >= 0 and
+          inner_axis >= 0):
+        x_with_placeholders = nest.map_structure(
+            lambda t: array_ops.placeholder_with_default(t, None),
+            x,
+            expand_composites=True)
+        actual_with_placeholders = x_with_placeholders.merge_dims(
+            outer_axis, inner_axis)
+        self.assertAllEqual(expected, actual_with_placeholders)
+
+  @parameterized.parameters([
+      {
+          'rt': [[1]],
+          'outer_axis': {},
+          'inner_axis': 1,
+          'exception': TypeError,
+          'message': 'outer_axis must be an int',
+      },
+      {
+          'rt': [[1]],
+          'outer_axis': 1,
+          'inner_axis': {},
+          'exception': TypeError,
+          'message': 'inner_axis must be an int',
+      },
+      {
+          'rt': [[1]],
+          'outer_axis': 1,
+          'inner_axis': 3,
+          'exception': ValueError,
+          'message': 'inner_axis=3 out of bounds: expected -2<=inner_axis<2',
+      },
+      {
+          'rt': [[1]],
+          'outer_axis': 1,
+          'inner_axis': -3,
+          'exception': ValueError,
+          'message': 'inner_axis=-3 out of bounds: expected -2<=inner_axis<2',
+      },
+      {
+          'rt': [[1]],
+          'outer_axis': 1,
+          'inner_axis': 0,
+          'exception': ValueError,
+          'message': 'Expected outer_axis .* to be less than or equal to .*',
+      },
+      {
+          'rt': [[1]],
+          'outer_axis': -1,
+          'inner_axis': -2,
+          'exception': ValueError,
+          'message': 'Expected outer_axis .* to be less than or equal to .*',
+      },
+  ])  # pyformat: disable
+  def testRaggedMergeDimsError(self,
+                               rt,
+                               outer_axis,
+                               inner_axis,
+                               exception,
+                               message=None,
+                               ragged_rank=None):
+    x = ragged_factory_ops.constant(rt, ragged_rank=ragged_rank)
+    with self.assertRaisesRegex(exception, message):
+      self.evaluate(x.merge_dims(outer_axis, inner_axis))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_one_hot_op_test.py

@@ -0,0 +1,196 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Tests for ragged_one_hot."""
+
+from absl.testing import parameterized
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.with_eager_op_as_function
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedOneHotTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      # 2D Indices (ragged_rank=1)
+      dict(indices=[[0, 2, -1], [3]],
+           depth=4,
+           expected=[[[1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]],
+                     [[0, 0, 0, 1]]]),
+      dict(indices=[[0, 2, -1], [3]],
+           depth=4,
+           axis=-1,
+           expected=[[[1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]],
+                     [[0, 0, 0, 1]]]),
+      dict(indices=[[0, 2, -1], [3]],
+           depth=4,
+           axis=2,
+           expected=[[[1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]],
+                     [[0, 0, 0, 1]]]),
+      dict(indices=[[0, 2, -1], [3]],
+           depth=4,
+           on_value=8,
+           off_value=4,
+           expected=[[[8, 4, 4, 4], [4, 4, 8, 4], [4, 4, 4, 4]],
+                     [[4, 4, 4, 8]]]),
+      dict(indices=[[0, 2, -1], [3]],
+           depth=0,
+           expected=[[[], [], []], [[]]]),
+      # 3D Indices (ragged_rank=2)
+      dict(indices=[[[0, 2, -1], [3]], [[2, 8]]],
+           depth=4,
+           expected=[[[[1, 0, 0, 0], [0, 0, 1, 0], [0, 0, 0, 0]],
+                      [[0, 0, 0, 1]]],
+                     [[[0, 0, 1, 0], [0, 0, 0, 0]]]]),
+      # 3D Indices (ragged_rank=1)
+      dict(indices=[[[0, 2], [-1, 3]], [[2, 8]]],
+           ragged_rank=1,
+           depth=4,
+           expected=[[[[1, 0, 0, 0], [0, 0, 1, 0]],
+                      [[0, 0, 0, 0], [0, 0, 0, 1]]],
+                     [[[0, 0, 1, 0], [0, 0, 0, 0]]]]),
+      dict(indices=[[[0, 2], [-1, 3]], [[2, 8]]],
+           ragged_rank=1,
+           axis=2,
+           depth=4,
+           expected=[[[[1, 0], [0, 0], [0, 1], [0, 0]],
+                      [[0, 0], [0, 0], [0, 0], [0, 1]]],
+                     [[[0, 0], [0, 0], [1, 0], [0, 0]]]]),
+  ])  # pyformat: disable
+  def testRaggedOneHot(self,
+                       indices,
+                       depth,
+                       on_value=None,
+                       off_value=None,
+                       axis=None,
+                       dtype=None,
+                       expected=None,
+                       ragged_rank=None):
+    ragged_indices = ragged_factory_ops.constant(
+        indices, ragged_rank=ragged_rank)
+    result = ragged_array_ops.ragged_one_hot(
+        ragged_indices,
+        depth,
+        on_value=on_value,
+        off_value=off_value,
+        axis=axis,
+        dtype=dtype)
+    self.assertAllEqual(result, expected)
+    self.assertEqual(result.ragged_rank, ragged_indices.ragged_rank)
+
+  @parameterized.parameters([
+      dict(indices=[[1]], depth=4, axis=0,  # axis < ragged_rank
+           message=r'axis \(0\) must be greater than indices.ragged_rank'),
+      dict(indices=[[1]], depth=4, axis=1,  # axis == ragged_rank
+           message=r'axis \(1\) must be greater than indices.ragged_rank'),
+      dict(indices=[[1]], depth=4, axis=-2,
+           # Note: the only negative `axis` value supported by
+           # array_ops.one_hot is -1.
+           message=(r'(?i)axis must be >= -1|'  # graph mode
+                    r'Expected axis.* to be -1 or between.*'),  # eager mode
+           exception=(ValueError, errors.InvalidArgumentError,
+                      errors.UnknownError)),
+  ])  # pyformat: disable
+  def testErrors(self,
+                 indices,
+                 depth,
+                 on_value=None,
+                 off_value=None,
+                 axis=None,
+                 dtype=None,
+                 exception=ValueError,
+                 message=None,
+                 ragged_rank=None):
+    ragged_indices = ragged_factory_ops.constant(
+        indices, ragged_rank=ragged_rank)
+    with self.assertRaisesRegex(exception, message):
+      array_ops.one_hot(
+          ragged_indices,
+          depth,
+          on_value=on_value,
+          off_value=off_value,
+          axis=axis,
+          dtype=dtype)
+
+  @parameterized.parameters([
+      dict(indices_shape=[5, 7], depth=6, axis=-1),
+      dict(indices_shape=[5, 7], depth=6, axis=2),
+      dict(indices_shape=[5, 2, 7], depth=3, axis=-1),
+      dict(indices_shape=[5, 2, 7], depth=3, axis=3),
+      dict(indices_shape=[5, 2, 7], depth=3, axis=2),
+      dict(indices_shape=[5, 2, 7, 4], depth=3, axis=-1),
+      dict(indices_shape=[5, 2, 7, 4], depth=3, axis=4),
+      dict(indices_shape=[5, 2, 7, 4], depth=3, axis=3),
+      dict(indices_shape=[5, 2, 7, 4], depth=3, axis=2),
+      dict(indices_shape=[5, 2, 7], depth=3, on_value=True, off_value=False),
+      dict(indices_shape=[5, 2, 7], depth=3, dtype=dtypes.float32),
+  ])  # pyformat: disable
+  def testRaggedOneHotMatchesArrayOpsOneHot(self,
+                                            indices_shape,
+                                            depth,
+                                            on_value=None,
+                                            off_value=None,
+                                            axis=None,
+                                            dtype=None):
+    """Tests that tf.one_hot gives the same result for ragged & uniform tensors.
+
+    Runs tf.one_hot with a uniform tensor, and compares the output with the
+    results of calling tf.one_hot with ragged version of that tensor with
+    varying ragged ranks.
+
+    Args:
+      indices_shape: Shape for `indices` arg to `tf.one_hot`
+      depth: `depth` arg to `tf.one_hot`
+      on_value: `on_value` arg to `tf.one_hot`
+      off_value: `off_value` arg to `tf.one_hot`
+      axis: `axis` arg to `tf.one_hot`
+      dtype: `dtype` arg to `tf.one_hot`
+    """
+    indices_shape = tensor_shape.as_shape(indices_shape)
+    indices = np.random.randint(depth + 1, size=indices_shape)
+    expected = array_ops.one_hot(
+        indices,
+        depth,
+        on_value=on_value,
+        off_value=off_value,
+        axis=axis,
+        dtype=dtype)
+    for ragged_rank in range(1, len(indices_shape)):
+      if axis is not None and 0 <= axis <= ragged_rank:
+        continue  # axis <= ragged_rank is not supported.
+      ragged_indices = ragged_tensor.RaggedTensor.from_tensor(
+          indices, ragged_rank=ragged_rank)
+      result = ragged_array_ops.ragged_one_hot(
+          ragged_indices,
+          depth,
+          on_value=on_value,
+          off_value=off_value,
+          axis=axis,
+          dtype=dtype)
+      self.assertAllEqual(result.to_tensor(), expected)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_operators.py

@@ -0,0 +1,342 @@
+# 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.
+# ==============================================================================
+"""Operator overloads for `RaggedTensor`."""
+
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_getitem
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import tf_decorator
+
+
+# =============================================================================
+# Equality Docstring
+# =============================================================================
+def ragged_eq(self, other):  # pylint: disable=g-doc-args
+  """Returns result of elementwise `==` or False if not broadcast-compatible.
+
+  Compares two ragged tensors elemewise for equality if they are
+  broadcast-compatible; or returns False if they are not
+  [broadcast-compatible](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html).
+
+  Note that this behavior differs from `tf.math.equal`, which raises an
+  exception if the two ragged tensors are not broadcast-compatible.
+
+  For example:
+
+  >>> rt1 = tf.ragged.constant([[1, 2], [3]])
+  >>> rt1 == rt1
+  <tf.RaggedTensor [[True, True], [True]]>
+
+  >>> rt2 = tf.ragged.constant([[1, 2], [4]])
+  >>> rt1 == rt2
+  <tf.RaggedTensor [[True, True], [False]]>
+
+  >>> rt3 = tf.ragged.constant([[1, 2], [3, 4]])
+  >>> # rt1 and rt3 are not broadcast-compatible.
+  >>> rt1 == rt3
+  False
+
+  >>> # You can also compare a `tf.RaggedTensor` to a `tf.Tensor`.
+  >>> t = tf.constant([[1, 2], [3, 4]])
+  >>> rt1 == t
+  False
+  >>> t == rt1
+  False
+  >>> rt4 = tf.ragged.constant([[1, 2], [3, 4]])
+  >>> rt4 == t
+  <tf.RaggedTensor [[True, True], [True, True]]>
+  >>> t == rt4
+  <tf.RaggedTensor [[True, True], [True, True]]>
+
+  Args:
+    other: The right-hand side of the `==` operator.
+
+  Returns:
+    The ragged tensor result of the elementwise `==` operation, or `False` if
+    the arguments are not broadcast-compatible.
+  """
+  return math_ops.tensor_equals(self, other)
+
+
+# =============================================================================
+# Ordering Docstring
+# =============================================================================
+def ragged_ge(self, other):  # pylint: disable=g-doc-args
+  """Elementwise `>=` comparison of two convertible-to-ragged-tensor values.
+
+  Computes the elemewise `>=` comparison of two values that are convertible to
+  ragged tenors, with [broadcasting]
+  (http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html) support.
+  Raises an exception if two values are not broadcast-compatible.
+
+  For example:
+
+  >>> rt1 = tf.ragged.constant([[1, 2], [3]])
+  >>> rt1 >= rt1
+  <tf.RaggedTensor [[True, True], [True]]>
+
+  >>> rt2 = tf.ragged.constant([[2, 1], [3]])
+  >>> rt1 >= rt2
+  <tf.RaggedTensor [[False, True], [True]]>
+
+  >>> rt3 = tf.ragged.constant([[1, 2], [3, 4]])
+  >>> # rt1 and rt3 are not broadcast-compatible.
+  >>> rt1 >= rt3
+  Traceback (most recent call last):
+  ...
+  InvalidArgumentError: ...
+
+  >>> # You can also compare a `tf.RaggedTensor` to a `tf.Tensor`.
+  >>> rt4 = tf.ragged.constant([[1, 2],[3, 4]])
+  >>> t1 = tf.constant([[2, 1], [4, 3]])
+  >>> rt4 >= t1
+  <tf.RaggedTensor [[False, True],
+   [False, True]]>
+  >>> t1 >= rt4
+  <tf.RaggedTensor [[True, False],
+   [True, False]]>
+
+  >>> # Compares a `tf.RaggedTensor` to a `tf.Tensor` with broadcasting.
+  >>> t2 = tf.constant([[2]])
+  >>> rt4 >= t2
+  <tf.RaggedTensor [[False, True],
+   [True, True]]>
+  >>> t2 >= rt4
+  <tf.RaggedTensor [[True, True],
+   [False, False]]>
+
+  Args:
+    other: The right-hand side of the `>=` operator.
+
+  Returns:
+    A `tf.RaggedTensor` of dtype `tf.bool` with the shape that `self` and
+    `other` broadcast to.
+
+  Raises:
+    InvalidArgumentError: If `self` and `other` are not broadcast-compatible.
+  """
+  return math_ops.greater_equal(self, other)
+
+
+# =============================================================================
+# Logical Docstring
+# =============================================================================
+
+
+# =============================================================================
+# Arithmetic Docstring
+# =============================================================================
+def ragged_abs(self, name=None):  # pylint: disable=g-doc-args
+  r"""Computes the absolute value of a ragged tensor.
+
+  Given a ragged tensor of integer or floating-point values, this operation
+  returns a ragged tensor of the same type, where each element contains the
+  absolute value of the corresponding element in the input.
+
+  Given a ragged tensor `x` of complex numbers, this operation returns a tensor
+  of type `float32` or `float64` that is the absolute value of each element in
+  `x`. For a complex number \\(a + bj\\), its absolute value is computed as
+  \\(\sqrt{a^2 + b^2}\\).
+
+  For example:
+
+  >>> # real number
+  >>> x = tf.ragged.constant([[-2.2, 3.2], [-4.2]])
+  >>> tf.abs(x)
+  <tf.RaggedTensor [[2.2, 3.2], [4.2]]>
+
+  >>> # complex number
+  >>> x = tf.ragged.constant([[-2.2 + 4.7j], [-3.2 + 5.7j], [-4.2 + 6.7j]])
+  >>> tf.abs(x)
+  <tf.RaggedTensor [[5.189412298131649],
+   [6.536818798161687],
+   [7.907591289387685]]>
+
+  Args:
+    name: A name for the operation (optional).
+
+  Returns:
+    A `RaggedTensor` of the same size and type as `x`, with absolute values.
+    Note, for `complex64` or `complex128` input, the returned `RaggedTensor`
+    will be of type `float32` or `float64`, respectively.
+  """
+  return math_ops.abs(self, name=name)
+
+
+# ===========================================================================
+def ragged_and(self, y, name=None):  # pylint: disable=g-doc-args
+  r"""Returns the truth value of elementwise `x & y`.
+
+  Logical AND function.
+
+  Requires that `x` and `y` have the same shape or have
+  [broadcast-compatible](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+  shapes. For example, `y` can be:
+
+    - A single Python boolean, where the result will be calculated by applying
+      logical AND with the single element to each element in `x`.
+    - A `tf.Tensor` object of dtype `tf.bool` of the same shape or
+      [broadcast-compatible](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+      shape. In this case, the result will be the element-wise logical AND of
+      `x` and `y`.
+    - A `tf.RaggedTensor` object of dtype `tf.bool` of the same shape or
+      [broadcast-compatible](http://docs.scipy.org/doc/numpy/user/basics.broadcasting.html)
+      shape. In this case, the result will be the element-wise logical AND of
+      `x` and `y`.
+
+  For example:
+
+  >>> # `y` is a Python boolean
+  >>> x = tf.ragged.constant([[True, False], [True]])
+  >>> y = True
+  >>> x & y
+  <tf.RaggedTensor [[True, False], [True]]>
+  >>> tf.math.logical_and(x, y)  # Equivalent of x & y
+  <tf.RaggedTensor [[True, False], [True]]>
+  >>> y & x
+  <tf.RaggedTensor [[True, False], [True]]>
+  >>> tf.math.reduce_all(x & y)  # Reduce to a scalar bool Tensor.
+  <tf.Tensor: shape=(), dtype=bool, numpy=False>
+
+  >>> # `y` is a tf.Tensor of the same shape.
+  >>> x = tf.ragged.constant([[True, False], [True, False]])
+  >>> y = tf.constant([[True, False], [False, True]])
+  >>> x & y
+  <tf.RaggedTensor [[True, False], [False, False]]>
+
+  >>> # `y` is a tf.Tensor of a broadcast-compatible shape.
+  >>> x = tf.ragged.constant([[True, False], [True]])
+  >>> y = tf.constant([[True], [False]])
+  >>> x & y
+  <tf.RaggedTensor [[True, False], [False]]>
+
+  >>> # `y` is a `tf.RaggedTensor` of the same shape.
+  >>> x = tf.ragged.constant([[True, False], [True]])
+  >>> y = tf.ragged.constant([[False, True], [True]])
+  >>> x & y
+  <tf.RaggedTensor [[False, False], [True]]>
+
+  >>> # `y` is a `tf.RaggedTensor` of a broadcast-compatible shape.
+  >>> x = tf.ragged.constant([[[True, True, False]], [[]], [[True, False]]])
+  >>> y = tf.ragged.constant([[[True]], [[True]], [[False]]], ragged_rank=1)
+  >>> x & y
+  <tf.RaggedTensor [[[True, True, False]], [[]], [[False, False]]]>
+
+  Args:
+    y: A Python boolean or a `tf.Tensor` or `tf.RaggedTensor` of dtype
+      `tf.bool`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `tf.RaggedTensor` of dtype `tf.bool` with the shape that `x` and `y`
+    broadcast to.
+  """
+  return math_ops.logical_and(self, y, name)
+
+
+# Helper Methods.
+def _right(operator):
+  """Right-handed version of an operator: swap args x and y."""
+  return tf_decorator.make_decorator(operator, lambda y, x: operator(x, y))
+
+
+def ragged_hash(self):
+  """The operation invoked by the `RaggedTensor.__hash__` operator."""
+  g = getattr(self.row_splits, "graph", None)
+  # pylint: disable=protected-access
+  if (
+      tensor.Tensor._USE_EQUALITY
+      and ops.executing_eagerly_outside_functions()
+      and (g is None or g.building_function)
+  ):
+    raise TypeError("RaggedTensor is unhashable.")
+  else:
+    return id(self)
+
+
+# Indexing
+ragged_tensor.RaggedTensor.__getitem__ = ragged_getitem.ragged_tensor_getitem
+
+# Equality
+ragged_tensor.RaggedTensor.__eq__ = ragged_eq
+ragged_tensor.RaggedTensor.__ne__ = math_ops.tensor_not_equals
+ragged_tensor.RaggedTensor.__hash__ = ragged_hash
+
+# Ordering operators
+ragged_tensor.RaggedTensor.__ge__ = ragged_ge
+ragged_tensor.RaggedTensor.__gt__ = math_ops.greater
+ragged_tensor.RaggedTensor.__le__ = math_ops.less_equal
+ragged_tensor.RaggedTensor.__lt__ = math_ops.less
+
+# Logical operators
+ragged_tensor.RaggedTensor.__and__ = ragged_and
+ragged_tensor.RaggedTensor.__rand__ = _right(ragged_and)
+
+ragged_tensor.RaggedTensor.__invert__ = math_ops.logical_not
+ragged_tensor.RaggedTensor.__ror__ = _right(math_ops.logical_or)
+ragged_tensor.RaggedTensor.__or__ = math_ops.logical_or
+ragged_tensor.RaggedTensor.__xor__ = math_ops.logical_xor
+ragged_tensor.RaggedTensor.__rxor__ = _right(math_ops.logical_xor)
+
+# Arithmetic operators
+ragged_tensor.RaggedTensor.__abs__ = ragged_abs
+ragged_tensor.RaggedTensor.__add__ = math_ops.add
+ragged_tensor.RaggedTensor.__radd__ = _right(math_ops.add)
+ragged_tensor.RaggedTensor.__div__ = math_ops.div
+ragged_tensor.RaggedTensor.__rdiv__ = _right(math_ops.div)
+ragged_tensor.RaggedTensor.__floordiv__ = math_ops.floordiv
+ragged_tensor.RaggedTensor.__rfloordiv__ = _right(math_ops.floordiv)
+ragged_tensor.RaggedTensor.__mod__ = math_ops.floormod
+ragged_tensor.RaggedTensor.__rmod__ = _right(math_ops.floormod)
+ragged_tensor.RaggedTensor.__mul__ = math_ops.multiply
+ragged_tensor.RaggedTensor.__rmul__ = _right(math_ops.multiply)
+ragged_tensor.RaggedTensor.__neg__ = math_ops.negative
+ragged_tensor.RaggedTensor.__pow__ = math_ops.pow
+ragged_tensor.RaggedTensor.__rpow__ = _right(math_ops.pow)
+ragged_tensor.RaggedTensor.__sub__ = math_ops.subtract
+ragged_tensor.RaggedTensor.__rsub__ = _right(math_ops.subtract)
+ragged_tensor.RaggedTensor.__truediv__ = math_ops.truediv
+ragged_tensor.RaggedTensor.__rtruediv__ = _right(math_ops.truediv)
+
+
+def ragged_bool(self):  # pylint: disable=g-doc-args
+  """Raises TypeError when a RaggedTensor is used as a Python bool.
+
+  To prevent RaggedTensor from being used as a bool, this function always raise
+  TypeError when being called.
+
+  For example:
+
+  >>> x = tf.ragged.constant([[1, 2], [3]])
+  >>> result = True if x else False  # Evaluate x as a bool value.
+  Traceback (most recent call last):
+  ...
+  TypeError: RaggedTensor may not be used as a boolean.
+
+  >>> x = tf.ragged.constant([[1]])
+  >>> r = (x == 1)  # tf.RaggedTensor [[True]]
+  >>> if r:  # Evaluate r as a bool value.
+  ...   pass
+  Traceback (most recent call last):
+  ...
+  TypeError: RaggedTensor may not be used as a boolean.
+  """
+  raise TypeError("RaggedTensor may not be used as a boolean.")
+
+
+ragged_tensor.RaggedTensor.__bool__ = ragged_bool  # Python3 bool conversion.
+ragged_tensor.RaggedTensor.__nonzero__ = ragged_bool  # Python2 bool conversion.

tensorflow/python/ops/ragged/ragged_operators_test.py

@@ -0,0 +1,118 @@
+# 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.
+# ==============================================================================
+"""Tests for overloaded RaggedTensor operators."""
+
+from tensorflow.python import tf2
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedElementwiseOpsTest(test_util.TensorFlowTestCase):
+
+  def testEqualityOperators(self):
+    a = ragged_factory_ops.constant([[1, 2], [3]])
+    b = ragged_factory_ops.constant([[4, 5], [3]])
+    c = 2
+
+    if tf2.enabled() and ops.executing_eagerly_outside_functions():
+      # Value-based equality:
+      self.assertAllEqual(a == b, [[False, False], [True]])
+      self.assertAllEqual(a != b, [[True, True], [False]])
+
+      # Value-based equality (w/ broadcasting):
+      self.assertAllEqual(a == c, [[False, True], [False]])
+      self.assertAllEqual(a != c, [[True, False], [True]])
+    else:
+      # Identity-based equality:
+      self.assertAllEqual(a == b, False)
+      self.assertAllEqual(a != b, True)
+
+  def testOrderingOperators(self):
+    x = ragged_factory_ops.constant([[1, 5], [3]])
+    y = ragged_factory_ops.constant([[4, 5], [1]])
+    self.assertAllEqual((x > y), [[False, False], [True]])
+    self.assertAllEqual((x >= y), [[False, True], [True]])
+    self.assertAllEqual((x < y), [[True, False], [False]])
+    self.assertAllEqual((x <= y), [[True, True], [False]])
+
+  def testArithmeticOperators(self):
+    x = ragged_factory_ops.constant([[1.0, -2.0], [8.0]])
+    y = ragged_factory_ops.constant([[4.0, 4.0], [2.0]])
+    self.assertAllEqual(abs(x), [[1.0, 2.0], [8.0]])
+
+    # pylint: disable=invalid-unary-operand-type
+    self.assertAllEqual((-x), [[-1.0, 2.0], [-8.0]])
+
+    self.assertAllEqual((x + y), [[5.0, 2.0], [10.0]])
+    self.assertAllEqual((3.0 + y), [[7.0, 7.0], [5.0]])
+    self.assertAllEqual((x + 3.0), [[4.0, 1.0], [11.0]])
+
+    self.assertAllEqual((x - y), [[-3.0, -6.0], [6.0]])
+    self.assertAllEqual((3.0 - y), [[-1.0, -1.0], [1.0]])
+    self.assertAllEqual((x + 3.0), [[4.0, 1.0], [11.0]])
+
+    self.assertAllEqual((x * y), [[4.0, -8.0], [16.0]])
+    self.assertAllEqual((3.0 * y), [[12.0, 12.0], [6.0]])
+    self.assertAllEqual((x * 3.0), [[3.0, -6.0], [24.0]])
+
+    self.assertAllEqual((x / y), [[0.25, -0.5], [4.0]])
+    self.assertAllEqual((y / x), [[4.0, -2.0], [0.25]])
+    self.assertAllEqual((2.0 / y), [[0.5, 0.5], [1.0]])
+    self.assertAllEqual((x / 2.0), [[0.5, -1.0], [4.0]])
+
+    self.assertAllEqual((x // y), [[0.0, -1.0], [4.0]])
+    self.assertAllEqual((y // x), [[4.0, -2.0], [0.0]])
+    self.assertAllEqual((2.0 // y), [[0.0, 0.0], [1.0]])
+    self.assertAllEqual((x // 2.0), [[0.0, -1.0], [4.0]])
+
+    self.assertAllEqual((x % y), [[1.0, 2.0], [0.0]])
+    self.assertAllEqual((y % x), [[0.0, -0.0], [2.0]])
+    self.assertAllEqual((2.0 % y), [[2.0, 2.0], [0.0]])
+    self.assertAllEqual((x % 2.0), [[1.0, 0.0], [0.0]])
+
+  def testLogicalOperators(self):
+    # pylint: disable=invalid-unary-operand-type
+    a = ragged_factory_ops.constant([[True, True], [False]])
+    b = ragged_factory_ops.constant([[True, False], [False]])
+    self.assertAllEqual((~a), [[False, False], [True]])
+
+    self.assertAllEqual((a & b), [[True, False], [False]])
+    self.assertAllEqual((a & True), [[True, True], [False]])
+    self.assertAllEqual((True & b), [[True, False], [False]])
+
+    self.assertAllEqual((a | b), [[True, True], [False]])
+    self.assertAllEqual((a | False), [[True, True], [False]])
+    self.assertAllEqual((False | b), [[True, False], [False]])
+
+    self.assertAllEqual((a ^ b), [[False, True], [False]])
+    self.assertAllEqual((a ^ True), [[False, False], [True]])
+    self.assertAllEqual((True ^ b), [[False, True], [True]])
+
+  def testDummyOperators(self):
+    a = ragged_factory_ops.constant([[True, True], [False]])
+    with self.assertRaisesRegex(TypeError,
+                                'RaggedTensor may not be used as a boolean.'):
+      bool(a)
+    with self.assertRaisesRegex(TypeError,
+                                'RaggedTensor may not be used as a boolean.'):
+      if a:
+        pass
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_ops.py

@@ -0,0 +1,51 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Import all modules in the `ragged` package that define exported symbols.
+
+Additional, import ragged_dispatch (which has the side-effect of registering
+dispatch handlers for many standard TF ops) and ragged_operators (which has the
+side-effect of overriding RaggedTensor operators, such as RaggedTensor.__add__).
+
+We don't import these modules from ragged/__init__.py, since we want to avoid
+circular dependencies.
+"""
+
+
+# pylint: disable=unused-import
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_autograph
+from tensorflow.python.ops.ragged import ragged_batch_gather_ops
+from tensorflow.python.ops.ragged import ragged_batch_gather_with_default_op
+from tensorflow.python.ops.ragged import ragged_bincount_ops
+from tensorflow.python.ops.ragged import ragged_check_ops
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_conversion_ops
+from tensorflow.python.ops.ragged import ragged_dispatch
+from tensorflow.python.ops.ragged import ragged_embedding_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_gather_ops
+from tensorflow.python.ops.ragged import ragged_getitem
+from tensorflow.python.ops.ragged import ragged_image_ops
+from tensorflow.python.ops.ragged import ragged_map_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_operators
+from tensorflow.python.ops.ragged import ragged_squeeze_op
+from tensorflow.python.ops.ragged import ragged_string_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_shape
+from tensorflow.python.ops.ragged import ragged_tensor_value
+from tensorflow.python.ops.ragged import ragged_where_op
+from tensorflow.python.ops.ragged import segment_id_ops

tensorflow/python/ops/ragged/ragged_placeholder_op_test.py

@@ -0,0 +1,74 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_placeholder op."""
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedPlaceholderOpTest(test_util.TensorFlowTestCase,
+                              parameterized.TestCase):
+
+  @parameterized.parameters([
+      # dtype, ragged_rank, value_shape, name -> expected
+      (dtypes.int32, 0, [5], None,
+       'Tensor("Placeholder:0", shape=(5,), dtype=int32)'),
+      (dtypes.int32, 1, [], 'ph', 'tf.RaggedTensor('
+       'values=Tensor("ph/flat_values:0", shape=(None,), dtype=int32), '
+       'row_splits=Tensor("ph/row_splits_0:0", shape=(None,), dtype=int64))'),
+      (dtypes.string, 1, [5], 'ph', 'tf.RaggedTensor('
+       'values=Tensor("ph/flat_values:0", shape=(None, 5), dtype=string), '
+       'row_splits=Tensor("ph/row_splits_0:0", shape=(None,), dtype=int64))'),
+      (dtypes.float32, 2, [], 'ph', 'tf.RaggedTensor(values=tf.RaggedTensor('
+       'values=Tensor("ph/flat_values:0", shape=(None,), dtype=float32), '
+       'row_splits=Tensor("ph/row_splits_1:0", shape=(None,), dtype=int64)), '
+       'row_splits=Tensor("ph/row_splits_0:0", shape=(None,), dtype=int64))'),
+      (dtypes.int32, 2, [3, 5], 'ph', 'tf.RaggedTensor(values=tf.RaggedTensor('
+       'values=Tensor("ph/flat_values:0", shape=(None, 3, 5), dtype=int32), '
+       'row_splits=Tensor("ph/row_splits_1:0", shape=(None,), dtype=int64)), '
+       'row_splits=Tensor("ph/row_splits_0:0", shape=(None,), dtype=int64))'),
+  ])
+  def testRaggedPlaceholder(self, dtype, ragged_rank, value_shape, name,
+                            expected):
+    if not context.executing_eagerly():
+      placeholder = ragged_factory_ops.placeholder(
+          dtype, ragged_rank, value_shape, name)
+      result = str(placeholder).replace('?', 'None')
+      self.assertEqual(result, expected)
+
+  def testRaggedPlaceholderRaisesExceptionInEagerMode(self):
+    if context.executing_eagerly():
+      with self.assertRaises(RuntimeError):
+        ragged_factory_ops.placeholder(dtypes.int32, 1, [])
+
+  def testRaggedPlaceholderDoesNotIncludeValidationOps(self):
+    if context.executing_eagerly():
+      return
+    graph = ops.Graph()
+    with graph.as_default():
+      ragged_factory_ops.placeholder(
+          dtypes.float32, ragged_rank=1, value_shape=[])
+      self.assertEqual([op.type for op in graph.get_operations()],
+                       ['Placeholder', 'Placeholder'])
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_print_op_test.py

@@ -0,0 +1,195 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Tests for tf.print with ragged tensors.
+
+Note: ragged support for tf.print is implemented by RaggedPrintV2Dispatcher in
+ragged_dispatch.py.
+"""
+
+import os.path
+import tempfile
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import logging_ops
+from tensorflow.python.ops import sparse_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_string_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedPrintV2Test(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  # pylint: disable=g-long-lambda
+  @parameterized.named_parameters([
+      dict(
+          testcase_name='2d_int_values',
+          inputs=lambda: [ragged_factory_ops.constant([[1, 2], [3]])],
+          expected='[[1, 2], [3]]\n'),
+      dict(
+          testcase_name='3d_int_values',
+          inputs=lambda: [ragged_factory_ops.constant([[[1, 2], [3]], [[4]]])],
+          expected='[[[1, 2], [3]], [[4]]]\n'),
+      dict(
+          testcase_name='2d_str_values',
+          inputs=lambda: [ragged_factory_ops.constant([['a', 'b'], ['c']])],
+          expected="[['a', 'b'], ['c']]\n"),
+      dict(
+          testcase_name='2d_str_values_with_escaping',
+          inputs=lambda: [ragged_factory_ops.constant([["a'b"], ['c"d']])],
+          expected="[['a\\'b'], ['c\"d']]\n"),
+      dict(
+          testcase_name='two_ragged_values',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2], [3]]),
+              ragged_factory_ops.constant([[5], [], [6, 7, 8]])
+          ],
+          expected='[[1, 2], [3]] [[5], [], [6, 7, 8]]\n'),
+      dict(
+          testcase_name='ragged_value_and_non_tensor_values',
+          inputs=lambda:
+          ['a', 5, True,
+           ragged_factory_ops.constant([[1, 2], [3]]), 'c'],
+          expected='a 5 True [[1, 2], [3]] c\n'),
+      dict(
+          testcase_name='ragged_value_and_dense_value',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2], [3]]),
+              constant_op.constant([[1, 2], [3, 4]])
+          ],
+          expected='[[1, 2], [3]] [[1 2]\n [3 4]]\n'),
+      dict(
+          testcase_name='ragged_value_and_sparse_value',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2], [3]]),
+              sparse_ops.from_dense([[1]])
+          ],
+          expected=(
+              '[[1, 2], [3]] '
+              "'SparseTensor(indices=[[0 0]], values=[1], shape=[1 1])'\n")),
+      dict(
+          testcase_name='summarize_default',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2, 3, 4, 5, 6, 7, 8, 9], [10], [
+              ], [], [], [], [11, 12]])
+          ],
+          expected=('[[1, 2, 3, ..., 7, 8, 9], [10], [], '
+                    '..., '
+                    '[], [], [11, 12]]\n')),
+      dict(
+          testcase_name='summarize_2',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2, 3, 4, 5, 6, 7, 8, 9], [10], [
+              ], [], [], [], [11, 12]])
+          ],
+          summarize=2,
+          expected='[[1, 2, ..., 8, 9], [10], ..., [], [11, 12]]\n'),
+      dict(
+          testcase_name='summarize_neg1',
+          inputs=lambda: [
+              ragged_factory_ops.constant([[1, 2, 3, 4, 5, 6, 7, 8, 9], [10], [
+              ], [], [], [], [11, 12]])
+          ],
+          summarize=-1,
+          expected=('[[1, 2, 3, 4, 5, 6, 7, 8, 9], [10], '
+                    '[], [], [], [], [11, 12]]\n')),
+  ])
+  def testRaggedPrint(self, inputs, expected, summarize=None):
+    if callable(inputs):
+      inputs = inputs()
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      path = os.path.join(tmpdirname, 'print_output')
+      kwargs = {'output_stream': 'file://{}'.format(path)}
+      if summarize is not None:
+        kwargs.update(summarize=summarize)
+      self.evaluate(logging_ops.print_v2(*inputs, **kwargs))
+      actual = open(path, 'r').read()
+      self.assertEqual(repr(actual), repr(expected))
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedToStringTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.named_parameters([
+      ('2d_int', [[1, 2], [], [3, 4, 5]], '[[1, 2], [], [3, 4, 5]]'),
+      ('2d_str', [['a'], ['b'], ['c', 'd']], "[['a'], ['b'], ['c', 'd']]"),
+      ('3d_int', [[[1, 2], []], [[3, 4, 5]]], '[[[1, 2], []], [[3, 4, 5]]]'),
+      ('escape', [["a'b"], [r'c\d']], r"[['a\'b'], ['c\\d']]"),
+      dict(testcase_name='2d_empty', rt=[], ragged_rank=1, expected='[]'),
+      dict(testcase_name='3d_empty', rt=[], ragged_rank=2, expected='[]'),
+      dict(
+          testcase_name='3d_rrank1',
+          rt=[[[1, 2], [3, 4]], [], [[5, 6]]],
+          ragged_rank=1,
+          expected='[[[1, 2], [3, 4]], [], [[5, 6]]]'),
+      dict(
+          testcase_name='2d_empty_row', rt=[[]], ragged_rank=1,
+          expected='[[]]'),
+      dict(
+          testcase_name='3d_empty_row', rt=[[]], ragged_rank=2,
+          expected='[[]]'),
+      dict(
+          testcase_name='summarize_1',
+          rt=[[1, 2, 3, 4, 5], [], [6], [7], [8, 9]],
+          summarize=1,
+          expected='[[1, ..., 5], ..., [8, 9]]'),
+      dict(
+          testcase_name='summarize_2',
+          rt=[[1, 2, 3, 4, 5], [], [6], [7], [8, 9]],
+          summarize=2,
+          expected='[[1, 2, ..., 4, 5], [], ..., [7], [8, 9]]'),
+  ])
+  def testRaggedToString(self, rt, expected, summarize=None, ragged_rank=None):
+    rt = ragged_factory_ops.constant(rt, ragged_rank=ragged_rank)
+    actual = ragged_string_ops.ragged_tensor_to_string(rt, summarize=summarize)
+    self.assertAllEqual(actual, expected)
+
+  @parameterized.named_parameters([
+      ('maxelts_BadType', [[1]], "Expected summarize .*, got 'foo'", 'foo'),
+      ('maxelts_0', [[1]], 'Expected summarize to be .*, got 0', 0),
+      ('maxelts_Neg2', [[1]], 'Expected summarize to be .*, got -2', -2),
+  ])
+  def testRaggedToStringErrors(self,
+                               rt,
+                               error,
+                               summarize=None,
+                               exception=ValueError):
+    rt = ragged_factory_ops.constant(rt)
+    with self.assertRaisesRegex(exception, error):
+      self.evaluate(
+          ragged_string_ops.ragged_tensor_to_string(rt, summarize=summarize))
+
+  def testRaggedToStringUnknownRank(self):
+
+    @def_function.function(input_signature=[
+        ragged_tensor.RaggedTensorSpec(ragged_rank=1, dtype=dtypes.int32)
+    ])
+    def f(rt):
+      return ragged_string_ops.ragged_tensor_to_string(rt)
+
+    with self.assertRaisesRegex(
+        ValueError, 'RaggedTensor to_string requires '
+        'that rt.shape.rank is not None'):
+      f(ragged_factory_ops.constant([[1, 2], [3]]))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_range_op_test.py

@@ -0,0 +1,144 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_range op."""
+
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import gen_ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedRangeOpTest(test_util.TensorFlowTestCase):
+
+  def testDocStringExamples(self):
+    """Examples from ragged_range.__doc__."""
+    rt1 = ragged_math_ops.range([3, 5, 2])
+    self.assertAllEqual(rt1, [[0, 1, 2], [0, 1, 2, 3, 4], [0, 1]])
+
+    rt2 = ragged_math_ops.range([0, 5, 8], [3, 3, 12])
+    self.assertAllEqual(rt2, [[0, 1, 2], [], [8, 9, 10, 11]])
+
+    rt3 = ragged_math_ops.range([0, 5, 8], [3, 3, 12], 2)
+    self.assertAllEqual(rt3, [[0, 2], [], [8, 10]])
+
+  def testBasicRanges(self):
+    # Specify limits only.
+    self.assertAllEqual(
+        ragged_math_ops.range([0, 3, 5]),
+        [list(range(0)), list(range(3)),
+         list(range(5))])
+
+    # Specify starts and limits.
+    self.assertAllEqual(
+        ragged_math_ops.range([0, 3, 5], [2, 3, 10]),
+        [list(range(0, 2)),
+         list(range(3, 3)),
+         list(range(5, 10))])
+
+    # Specify starts, limits, and deltas.
+    self.assertAllEqual(
+        ragged_math_ops.range([0, 3, 5], [4, 4, 15], [2, 3, 4]),
+        [list(range(0, 4, 2)),
+         list(range(3, 4, 3)),
+         list(range(5, 15, 4))])
+
+  def testFloatRanges(self):
+    expected = [[0.0, 0.4, 0.8, 1.2, 1.6, 2.0, 2.4, 2.8, 3.2, 3.6], [3.0],
+                [5.0, 7.2, 9.4, 11.6, 13.8]]
+    actual = ragged_math_ops.range([0.0, 3.0, 5.0], [3.9, 4.0, 15.0],
+                                   [0.4, 1.5, 2.2])
+    self.assertAllClose(actual, expected)
+
+  def testNegativeDeltas(self):
+    self.assertAllEqual(
+        ragged_math_ops.range([0, 3, 5], limits=0, deltas=-1),
+        [list(range(0, 0, -1)),
+         list(range(3, 0, -1)),
+         list(range(5, 0, -1))])
+
+    self.assertAllEqual(
+        ragged_math_ops.range([0, -3, 5], limits=0, deltas=[-1, 1, -2]),
+        [list(range(0, 0, -1)),
+         list(range(-3, 0, 1)),
+         list(range(5, 0, -2))])
+
+  def testBroadcast(self):
+    # Specify starts and limits, broadcast deltas.
+    self.assertAllEqual(
+        ragged_math_ops.range([0, 3, 5], [4, 4, 15], 3),
+        [list(range(0, 4, 3)),
+         list(range(3, 4, 3)),
+         list(range(5, 15, 3))])
+
+    # Broadcast all arguments.
+    self.assertAllEqual(ragged_math_ops.range(0, 5, 1), [list(range(0, 5, 1))])
+
+  def testEmptyRanges(self):
+    rt1 = ragged_math_ops.range([0, 5, 3], [0, 3, 5])
+    rt2 = ragged_math_ops.range([0, 5, 5], [0, 3, 5], -1)
+    self.assertAllEqual(rt1, [[], [], [3, 4]])
+    self.assertAllEqual(rt2, [[], [5, 4], []])
+
+  def testShapeFnErrors(self):
+    self.assertRaises((ValueError, errors.InvalidArgumentError),
+                      ragged_math_ops.range, [[0]], 5)
+    self.assertRaises((ValueError, errors.InvalidArgumentError),
+                      ragged_math_ops.range, 0, [[5]])
+    self.assertRaises((ValueError, errors.InvalidArgumentError),
+                      ragged_math_ops.range, 0, 5, [[0]])
+    self.assertRaises((ValueError, errors.InvalidArgumentError),
+                      ragged_math_ops.range, [0], [1, 2])
+
+  def testKernelErrors(self):
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r'Requires delta != 0'):
+      self.evaluate(ragged_math_ops.range(0, 0, 0))
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                r'Requires \(\(limit - start\) / delta\) <='):
+      self.evaluate(ragged_math_ops.range(0.1, 1e10, 1e-10))
+
+    with self.assertRaisesRegex(errors.InvalidArgumentError, 'overflowed'):
+      self.evaluate(
+          gen_ragged_math_ops.ragged_range(
+              starts=[0, 0],
+              limits=[2**31 - 1, 1],
+              deltas=[1, 1],
+              Tsplits=dtypes.int32))
+
+  def testShape(self):
+    self.assertAllEqual(
+        ragged_math_ops.range(0, 0, 1).shape.as_list(), [1, None])
+    self.assertAllEqual(
+        ragged_math_ops.range([1, 2, 3]).shape.as_list(), [3, None])
+    self.assertAllEqual(
+        ragged_math_ops.range([1, 2, 3], [4, 5, 6]).shape.as_list(), [3, None])
+
+  def testInt32Overflow(self):
+    start = 1136033460
+    end = -2110457150
+    step = -1849827689
+    expected = [np.arange(start, end, step)]
+    actual = ragged_math_ops.range(start, end, step)
+    self.assertAllEqual(expected, self.evaluate(actual))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_rank_op_test.py

@@ -0,0 +1,85 @@
+# Copyright 2019 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.
+# ==============================================================================
+"""Tests for ragged.rank op."""
+
+from absl.testing import parameterized
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedRankOpTest(test_util.TensorFlowTestCase,
+                       parameterized.TestCase):
+
+  @parameterized.parameters([
+      # Rank 0
+      dict(
+          test_input=1,
+          expected_rank=0,
+      ),
+      # Rank 1
+      dict(
+          test_input=[1],
+          expected_rank=1,
+      ),
+      dict(
+          test_input=[1, 2, 3, 4],
+          expected_rank=1,
+      ),
+      # Rank 2
+      dict(
+          test_input=[[1], [2], [3]],
+          expected_rank=2,
+      ),
+      # Rank 3
+      dict(
+          test_input=[[[1], [2, 3]], [[4], [5, 6, 7]]],
+          expected_rank=3,
+      ),
+      # Rank 3, ragged_rank=2
+      dict(
+          test_input=[[[1], [2, 3], [10, 20]],
+                      [[4], [5, 6, 7]]],
+          expected_rank=3,
+          ragged_rank=2,
+      ),
+      # Rank 4, ragged_rank=3 with dimensions: {2, (1, 2), (2), (1, 2)}
+      dict(
+          test_input=[[[[1], [2]]],
+                      [[[3, 4], [5, 6]], [[7, 8], [9, 10]]]],
+          expected_rank=4,
+      ),
+      # Rank 4, ragged_rank=2 with dimensions: {2, (1, 2), (1, 2), 2}
+      dict(
+          test_input=[
+              [[[1, 2]]],
+              [[[5, 6], [7, 8]],
+               [[9, 10], [11, 12]]]],
+          expected_rank=4,
+          ragged_rank=2,
+      ),
+
+  ])
+  def testRaggedRank(self, test_input, expected_rank, ragged_rank=None):
+    test_input = ragged_factory_ops.constant(
+        test_input, ragged_rank=ragged_rank)
+    self.assertAllEqual(ragged_array_ops.rank(
+        test_input), expected_rank)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_reduce_op_test.py

@@ -0,0 +1,710 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_math_ops.reduce_<AGGREGATE> ops."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.platform import googletest
+
+_MAX_INT32 = dtypes.int32.max
+_MIN_INT32 = dtypes.int32.min
+_NAN = np.nan
+
+
+def mean(*values):
+  return 1.0 * sum(values) / len(values)
+
+
+def variance(*values):
+  return np.var(values, dtype=np.float64)
+
+
+def std(*values):
+  return np.std(values, dtype=np.float64)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedReduceOpsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters(
+      #=========================================================================
+      # Docstring examples.  RaggedTensor for testing is:
+      #   [[3, 1, 4],
+      #    [1, 5,  ],
+      #    [9,     ],
+      #    [2, 6   ]]
+      #=========================================================================
+      # keepdims=True
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=False,
+          expected=[15, 12, 4]  # = [3+1+9+2, 1+5+6, 4]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=-2,
+          keepdims=False,
+          expected=[15, 12, 4]  # = [3+1+9+2, 1+5+6, 4]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=1,
+          keepdims=False,
+          expected=[8, 6, 9, 8]  # = [3+1+4, 1+5, 9, 2+6]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=-1,
+          keepdims=False,
+          expected=[8, 6, 9, 8]  # = [3+1+4, 1+5, 9, 2+6]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_prod,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=False,
+          expected=[54, 30, 4]  # = [3*1*9*2, 1*5*6, 4]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_prod,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=1,
+          keepdims=False,
+          expected=[12, 5, 9, 12]  # = [3*1*4, 1*5, 9, 2*6]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_min,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=False,
+          expected=[1, 1, 4]  # = [min(3, 1, 9, 2), min(1, 5, 6), 4]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_min,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=1,
+          keepdims=False,
+          expected=[1, 1, 9, 2]  # = [min(3, 1, 4), min(1, 5), 9, min(2, 6)]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_max,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=False,
+          expected=[9, 6, 4]  # = [max(3, 1, 9, 2), max(1, 5, 6), 4]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_max,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=1,
+          keepdims=False,
+          expected=[4, 5, 9, 6]  # = [max(3, 1, 4), max(1, 5), 9, max(2, 6)]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_mean,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=False,
+          expected=[3.75, 4, 4]  # = [mean(3, 1, 9, 2), mean(1, 5, 6), 4]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_variance,
+          rt_input=[[3, 1, 4], [1, 1], [9], [2, 1]],
+          axis=0,
+          keepdims=False,
+          expected=[9.6875, 0.0, 0.0]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_std,
+          rt_input=[[3, 1, 4], [3, 1], [2], [2, 1]],
+          axis=0,
+          keepdims=False,
+          expected=[0.5, 0., 0.]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_any,
+          rt_input=[[True, True], [True, True, False, True], [False, True]],
+          axis=0,
+          keepdims=False,
+          expected=[True, True, False, True]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_any,
+          rt_input=[[True, True], [True, True, False, True], [False, True]],
+          axis=1,
+          keepdims=False,
+          expected=[True, True, True]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_all,
+          rt_input=[[True, True], [True, True, False, True], [False, True]],
+          axis=0,
+          keepdims=False,
+          expected=[False, True, False, True]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_all,
+          rt_input=[[True, True], [True, True, False, True], [False, True]],
+          axis=1,
+          keepdims=False,
+          expected=[True, False, False]),
+      # keepdims=True
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=True,
+          expected=[[15, 12, 4]]  # = [[3+1+9+2, 1+5+6, 4]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=-2,
+          keepdims=True,
+          expected=[[15, 12, 4]]  # = [[3+1+9+2, 1+5+6, 4]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=1,
+          keepdims=True,
+          expected=[[8], [6], [9], [8]]  # = [[3+1+4], [1+5], [9], [2+6]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=-1,
+          keepdims=True,
+          expected=[[8], [6], [9], [8]]  # = [[3+1+4], [1+5], [9], [2+6]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_prod,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=True,
+          expected=[[54, 30, 4]]  # = [[3*1*9*2, 1*5*6, 4]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_prod,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=1,
+          keepdims=True,
+          expected=[[12], [5], [9], [12]]  # = [[3*1*4], [1*5], [9], [2*6]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_min,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=True,
+          expected=[[1, 1, 4]]  # = [[min(3, 1, 9, 2), min(1, 5, 6), 4]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_min,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=1,
+          keepdims=True,
+          expected=[[1], [1], [9],
+                    [2]]  # = [[min(3, 1, 4)], [min(1, 5)], [9], [min(2, 6)]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_max,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=True,
+          expected=[[9, 6, 4]]  # = [[max(3, 1, 9, 2), max(1, 5, 6), 4]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_max,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=1,
+          keepdims=True,
+          expected=[[4], [5], [9],
+                    [6]]  # = [[max(3, 1, 4)], [max(1, 5)], [9], [max(2, 6)]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_mean,
+          rt_input=[[3, 1, 4], [1, 5], [9], [2, 6]],
+          axis=0,
+          keepdims=True,
+          expected=[[3.75, 4, 4]]  # = [[mean(3, 1, 9, 2), mean(1, 5, 6), 4]]
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_variance,
+          rt_input=[[3, 1, 4], [1, 1], [9], [2, 1]],
+          axis=0,
+          keepdims=True,
+          expected=[[9.6875, 0., 0.]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_std,
+          rt_input=[[3, 1, 4], [3, 1], [2], [2, 1]],
+          axis=0,
+          keepdims=True,
+          expected=[[0.5, 0., 0.]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_any,
+          rt_input=[[True, True], [True, True, False, True], [False, True]],
+          axis=0,
+          keepdims=True,
+          expected=[[True, True, False, True]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_any,
+          rt_input=[[True, True], [True, True, False, True], [False, True]],
+          axis=1,
+          keepdims=True,
+          expected=[[True], [True], [True]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_all,
+          rt_input=[[True, True], [True, True, False, True], [False, True]],
+          axis=0,
+          keepdims=True,
+          expected=[[False, True, False, True]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_all,
+          rt_input=[[True, True], [True, True, False, True], [False, True]],
+          axis=1,
+          keepdims=True,
+          expected=[[True], [False], [False]]),
+
+      #=========================================================================
+      # Examples with the following RaggedTensor (ragged_rank=1):
+      #   [[0, 1, 2, 3],
+      #    [4         ],
+      #    [          ],
+      #    [5, 6      ],
+      #    [7         ],
+      #    [8, 9      ]]
+      #=========================================================================
+
+      # axis=None
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=None,
+          keepdims=False,
+          expected=0 + 1 + 2 + 3 + 4 + 5 + 6 + 7 + 8 + 9),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_prod,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=None,
+          keepdims=False,
+          expected=0 * 1 * 2 * 3 * 4 * 5 * 6 * 7 * 8 * 9),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_min,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=None,
+          keepdims=False,
+          expected=min(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_max,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=None,
+          keepdims=False,
+          expected=max(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_mean,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=None,
+          keepdims=False,
+          expected=mean(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_variance,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=None,
+          keepdims=False,
+          expected=variance(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_std,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=None,
+          keepdims=False,
+          expected=std(0, 1, 2, 3, 4, 5, 6, 7, 8, 9)),
+      # axis=0
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=0,
+          keepdims=False,
+          expected=[0 + 4 + 5 + 7 + 8, 1 + 6 + 9, 2, 3]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_prod,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=0,
+          keepdims=False,
+          expected=[0 * 4 * 5 * 7 * 8, 1 * 6 * 9, 2, 3]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_min,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=0,
+          keepdims=False,
+          expected=[min(0, 4, 5, 7, 8), min(1, 6, 9), 2, 3]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_max,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=0,
+          keepdims=False,
+          expected=[max(0, 4, 5, 7, 8), max(1, 6, 9), 2, 3]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_mean,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=0,
+          keepdims=False,
+          expected=[mean(0, 4, 5, 7, 8),
+                    mean(1, 6, 9), 2, 3]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_variance,
+          rt_input=[[0, 1, 2, 3], [1], [], [2, 1], [3], [4, 1]],
+          axis=0,
+          keepdims=False,
+          expected=[variance(0, 1, 2, 3, 4),
+                    variance(1, 1, 1), 0, 0]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_std,
+          rt_input=[[1, 1, 2, 3], [1], [], [1, 1], [1], [1, 1]],
+          axis=0,
+          keepdims=False,
+          expected=[std(1, 1, 1, 1, 1), std(1, 1, 1), 0, 0]),
+      # axis=1
+      # Note: we don't test mean here because it gives a NaN, and this will
+      # cause assertEqual to fail (since NaN != NaN).  See testMeanNan().
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=1,
+          keepdims=False,
+          expected=[0 + 1 + 2 + 3, 4, 0, 5 + 6, 7, 8 + 9]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_prod,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=1,
+          keepdims=False,
+          expected=[0 * 1 * 2 * 3, 4, 1, 5 * 6, 7, 8 * 9]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_min,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=1,
+          keepdims=False,
+          expected=[min(0, 1, 2, 3), 4, _MAX_INT32,
+                    min(5, 6), 7,
+                    min(8, 9)]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_max,
+          rt_input=[[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]],
+          axis=1,
+          keepdims=False,
+          expected=[max(0, 1, 2, 3), 4, _MIN_INT32,
+                    max(5, 6), 7,
+                    max(8, 9)]),
+
+      #=========================================================================
+      # Examples with ragged_rank=2:
+      # [[[1, 2], [ ], [3, 4, 5]],
+      #  [[6, 7], [ ], [8      ]],
+      #  [                      ],
+      #  [[9   ]                ]]
+      #=========================================================================
+      # keepdims=False
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[],
+          keepdims=False,
+          expected=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=None,
+          keepdims=False,
+          expected=sum([1, 2, 3, 4, 5, 6, 7, 8, 9])),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=0,
+          keepdims=False,
+          expected=[[1 + 6 + 9, 2 + 7], [], [3 + 8, 4, 5]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=1,
+          keepdims=False,
+          expected=[[1 + 3, 2 + 4, 5], [6 + 8, 7], [], [9]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=2,
+          keepdims=False,
+          expected=[[1 + 2, 0, 3 + 4 + 5], [6 + 7, 0, 8], [], [9]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[0, 1],
+          keepdims=False,
+          expected=[1 + 3 + 6 + 8 + 9, 2 + 4 + 7, 5]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[0, 2],
+          keepdims=False,
+          expected=[1 + 6 + 9 + 2 + 7, 0, 3 + 8 + 4 + 5]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[1, 2],
+          keepdims=False,
+          expected=[1 + 2 + 3 + 4 + 5, 6 + 7 + 8, 0, 9]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[0, 1, 2],
+          keepdims=False,
+          expected=sum([1, 2, 3, 4, 5, 6, 7, 8, 9])),
+      # keepdims=True
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[],
+          keepdims=True,
+          expected=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=None,
+          keepdims=True,
+          expected=[[[sum([1, 2, 3, 4, 5, 6, 7, 8, 9])]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=0,
+          keepdims=True,
+          expected=[[[1 + 6 + 9, 2 + 7], [], [3 + 8, 4, 5]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=1,
+          keepdims=True,
+          expected=[[[1 + 3, 2 + 4, 5]], [[6 + 8, 7]], [[]], [[9]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=2,
+          keepdims=True,
+          expected=[[[1 + 2], [0], [3 + 4 + 5]], [[6 + 7], [0], [8]], [],
+                    [[9]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[0, 1],
+          keepdims=True,
+          expected=[[[1 + 3 + 6 + 8 + 9, 2 + 4 + 7, 5]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[0, 2],
+          keepdims=True,
+          expected=[[[1 + 6 + 9 + 2 + 7], [0], [3 + 8 + 4 + 5]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[1, 2],
+          keepdims=True,
+          expected=[[[1 + 2 + 3 + 4 + 5]], [[6 + 7 + 8]], [[0]], [[9]]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[0, 1, 2],
+          keepdims=True,
+          expected=[[[sum([1, 2, 3, 4, 5, 6, 7, 8, 9])]]]),
+
+      #=========================================================================
+      # Examples for ragged_reduce_mean ragged_rank=2:
+      # [[[1, 2], [3, 4, 5]],
+      #  [[6, 7], [8      ]],
+      #  [[9   ]          ]]
+      #=========================================================================
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_mean,
+          rt_input=[[[1, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
+          axis=0,
+          keepdims=False,
+          expected=[[mean(1, 6, 9), mean(2, 7)], [mean(3, 8), 4, 5]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_mean,
+          rt_input=[[[1, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
+          axis=1,
+          keepdims=False,
+          expected=[[mean(1, 3), mean(2, 4), 5], [mean(6, 8), 7], [9]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_mean,
+          rt_input=[[[1, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
+          axis=2,
+          keepdims=False,
+          expected=[[mean(1, 2), mean(3, 4, 5)], [mean(6, 7), 8], [9]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_variance,
+          rt_input=[[[6, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
+          axis=0,
+          keepdims=False,
+          expected=[[variance(6, 6, 9), variance(2, 7)],
+                    [variance(3, 8), 0., 0.]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_variance,
+          rt_input=[[[6, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
+          axis=1,
+          keepdims=False,
+          expected=[[variance(6, 3), variance(2, 4), 0.], [variance(6, 8), 0.],
+                    [0.]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_variance,
+          rt_input=[[[6, 2], [6, 9, 9]], [[6, 7], [8]], [[9]]],
+          axis=2,
+          keepdims=False,
+          expected=[[variance(6, 2), variance(6, 9, 9)], [variance(6, 7), 0.],
+                    [0.]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_std,
+          rt_input=[[[6, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
+          axis=0,
+          keepdims=False,
+          expected=[[std(6, 6, 9), std(2, 7)], [std(3, 8), 0., 0.]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_std,
+          rt_input=[[[6, 2], [3, 4, 5]], [[6, 7], [8]], [[9]]],
+          axis=1,
+          keepdims=False,
+          expected=[[std(6, 3), std(2, 4), 0.], [std(6, 8), 0.], [0.]]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_std,
+          rt_input=[[[6, 2], [6, 9, 9]], [[6, 7], [8]], [[9]]],
+          axis=2,
+          keepdims=False,
+          expected=[[std(6, 2), std(6, 9, 9)], [std(6, 7), 0.], [0.]]),
+
+      # Test case for GitHub issue 27497, multiple negative axes.
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[-2, -1],
+          keepdims=False,
+          expected=[1 + 2 + 3 + 4 + 5, 6 + 7 + 8, 0, 9]),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_sum,
+          rt_input=[[[1, 2], [], [3, 4, 5]], [[6, 7], [], [8]], [], [[9]]],
+          axis=[-3, -2, -1],
+          keepdims=False,
+          expected=sum([1, 2, 3, 4, 5, 6, 7, 8, 9])),
+      # Test case for GitHub issue 56222, small variance
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_variance,
+          rt_input=[[[0.214441], [0.214441], [0.214441], [0.214441], [0.214441],
+                     [0.214441], [0.214441]]],
+          axis=[1],
+          keepdims=False,
+          expected=[[0.0]],
+      ),
+      dict(
+          ragged_reduce_op=ragged_math_ops.reduce_std,
+          rt_input=[[[0.214441], [0.214441], [0.214441], [0.214441], [0.214441],
+                     [0.214441], [0.214441]]],
+          axis=[1],
+          keepdims=False,
+          expected=[[0.0]],
+      ),
+  )
+  def testReduce(self, ragged_reduce_op, rt_input, axis, keepdims, expected):
+    rt_input = ragged_factory_ops.constant(rt_input)
+    reduced = ragged_reduce_op(rt_input, axis, keepdims=keepdims)
+    self.assertAllEqual(reduced, expected)
+
+  def testReduceKeepsInnerDimensionShape(self):
+    # Test for bug [b/139823356].
+    rt = ragged_factory_ops.constant([[[[1, 1]]]], ragged_rank=2)
+    self.assertEqual(rt.shape.as_list(), [1, None, None, 2])
+    reduced = ragged_math_ops.reduce_sum(rt, axis=2)
+    self.assertEqual(reduced.shape.as_list(), [1, None, 2])
+
+  def assertEqualWithNan(self, actual, expected):
+    """Like assertEqual, but NaN==NaN."""
+    self.assertTrue(
+        ((actual == expected) | (np.isnan(actual) & np.isnan(expected))).all())
+
+  def testMeanNan(self):
+    rt_as_list = [[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]]
+    expected = (
+        np.array([0 + 1 + 2 + 3, 4, 0, 5 + 6, 7, 8 + 9]) /
+        np.array([4, 1, 0, 2, 1, 2]))
+    rt_input = ragged_factory_ops.constant(rt_as_list)
+    reduced = ragged_math_ops.reduce_mean(rt_input, axis=1)
+    self.assertEqualWithNan(self.evaluate(reduced), expected)
+
+  def testVarianceNan(self):
+    rt_as_list = [[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]]
+    expected = ([
+        variance(0, 1, 2, 3),
+        variance(4),
+        variance(),
+        variance(5, 6),
+        variance(7),
+        variance(8, 9)
+    ])
+    rt_input = ragged_factory_ops.constant(rt_as_list)
+    reduced = ragged_math_ops.reduce_variance(rt_input, axis=1)
+    self.assertEqualWithNan(self.evaluate(reduced), expected)
+
+  def testStdNan(self):
+    rt_as_list = [[0, 1, 1, 0], [4], [], [5, 6], [7], [8, 9]]
+    expected = ([std(0, 1, 1, 0), std(4), std(), std(5, 6), std(7), std(8, 9)])
+    rt_input = ragged_factory_ops.constant(rt_as_list)
+    reduced = ragged_math_ops.reduce_std(rt_input, axis=1)
+    self.assertEqualWithNan(self.evaluate(reduced), expected)
+
+  def testMeanWithTensorInputs(self):
+    tensor = [[1.0, 2.0, 3.0], [10.0, 20.0, 30.0]]
+    expected = [2.0, 20.0]
+    reduced = ragged_math_ops.reduce_mean(tensor, axis=1)
+    self.assertAllEqual(reduced, expected)
+
+  def testVarianceWithTensorInputs(self):
+    tensor = [[6.0, 9.0, 6.0], [60.0, 90.0, 60.0]]
+    expected = [2., 200.]
+    reduced = ragged_math_ops.reduce_variance(tensor, axis=1)
+    self.assertAllEqual(reduced, expected)
+
+  def testStdWithTensorInputs(self):
+    tensor = [[1.0, 2.0, 2.0, 1.0], [10.0, 20.0, 20.0, 10.0]]
+    expected = [0.5, 5.]
+    reduced = ragged_math_ops.reduce_std(tensor, axis=1)
+    self.assertAllEqual(reduced, expected)
+
+  def testErrors(self):
+    rt_input = ragged_factory_ops.constant([[1, 2, 3], [4, 5]])
+    axis = array_ops.placeholder_with_default(constant_op.constant([0]), None)
+
+    if not context.executing_eagerly():
+      self.assertRaisesRegex(ValueError,
+                             r'axis must be known at graph construction time.',
+                             ragged_math_ops.reduce_sum, rt_input, axis)
+    self.assertRaisesRegex(TypeError, r'axis must be an int; got str.*',
+                           ragged_math_ops.reduce_sum, rt_input, ['x'])
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_resize_image_op_test.py

@@ -0,0 +1,136 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Tests for RaggedTensor dispatch of tf.images.resize."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import image_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedResizeImageOpTest(test_util.TensorFlowTestCase,
+                              parameterized.TestCase):
+
+  def make_image_batch(self, sizes, channels):
+    if not sizes:
+      return ragged_tensor.RaggedTensor.from_tensor(
+          array_ops.zeros([0, 5, 5, channels]), ragged_rank=2)
+    images = [
+        array_ops.reshape(
+            math_ops.range(w * h * channels * 1.0), [w, h, channels])
+        for (w, h) in sizes
+    ]
+    return ragged_concat_ops.stack(images)
+
+  @parameterized.parameters([
+      dict(src_sizes=[], dst_size=(4, 4), v1=True),
+      dict(src_sizes=[], dst_size=(4, 4), v1=False),
+      dict(src_sizes=[(2, 2)], dst_size=(4, 4), v1=True),
+      dict(src_sizes=[(2, 2)], dst_size=(4, 4), v1=False),
+      dict(src_sizes=[(2, 8), (3, 5), (10, 10)], dst_size=(5, 5), v1=True),
+      dict(src_sizes=[(2, 8), (3, 5), (10, 10)], dst_size=(5, 5), v1=False),
+  ])
+  def testResize(self, src_sizes, dst_size, v1=False):
+    resize = image_ops.resize_images if v1 else image_ops.resize_images_v2
+
+    # Construct the input images.
+    channels = 3
+    images = self.make_image_batch(src_sizes, channels)
+    expected_shape = [len(src_sizes)] + list(dst_size) + [channels]
+
+    # Resize the ragged batch of images.
+    resized_images = resize(images, dst_size)
+    self.assertIsInstance(resized_images, tensor.Tensor)
+    self.assertEqual(resized_images.shape.as_list(), expected_shape)
+
+    # Check that results for each image matches what we'd get with the
+    # non-batch version of tf.images.resize.
+    for i in range(len(src_sizes)):
+      actual = resized_images[i]
+      expected = resize(images[i].to_tensor(), dst_size)
+      self.assertAllClose(actual, expected)
+
+  @parameterized.parameters([
+      dict(src_shape=[None, None, None, None], src_sizes=[], dst_size=(4, 4)),
+      dict(src_shape=[None, None, None, 3], src_sizes=[], dst_size=(4, 4)),
+      dict(src_shape=[0, None, None, None], src_sizes=[], dst_size=(4, 4)),
+      dict(src_shape=[0, None, None, 3], src_sizes=[], dst_size=(4, 4)),
+      dict(
+          src_shape=[None, None, None, None],
+          src_sizes=[(2, 2)],
+          dst_size=(4, 4)),
+      dict(
+          src_shape=[None, None, None, None],
+          src_sizes=[(2, 8), (3, 5), (10, 10)],
+          dst_size=(5, 5)),
+      dict(
+          src_shape=[None, None, None, 1],
+          src_sizes=[(2, 8), (3, 5), (10, 10)],
+          dst_size=(5, 5)),
+      dict(
+          src_shape=[3, None, None, 1],
+          src_sizes=[(2, 8), (3, 5), (10, 10)],
+          dst_size=(5, 5)),
+  ])
+  def testResizeWithPartialStaticShape(self, src_shape, src_sizes, dst_size):
+    channels = src_shape[-1] or 3
+    images = self.make_image_batch(src_sizes, channels)
+    rt_spec = ragged_tensor.RaggedTensorSpec(src_shape,
+                                             ragged_rank=images.ragged_rank)
+    expected_shape = [len(src_sizes)] + list(dst_size) + [channels]
+
+    # Use @tf.function to erase static shape information.
+    @def_function.function(input_signature=[rt_spec])
+    def do_resize(images):
+      return image_ops.resize_images_v2(images, dst_size)
+
+    resized_images = do_resize(images)
+    self.assertIsInstance(resized_images, tensor.Tensor)
+    self.assertTrue(resized_images.shape.is_compatible_with(expected_shape))
+
+    # Check that results for each image matches what we'd get with the
+    # non-batch version of tf.images.resize.
+    for i in range(len(src_sizes)):
+      actual = resized_images[i]
+      expected = image_ops.resize_images_v2(images[i].to_tensor(), dst_size)
+      self.assertAllClose(actual, expected)
+
+  def testSizeIsTensor(self):
+    @def_function.function
+    def do_resize(images, new_size):
+      return image_ops.resize_images_v2(images, new_size)
+
+    src_images = self.make_image_batch([[5, 8], [3, 2], [10, 4]], 3)
+    resized_images = do_resize(src_images, constant_op.constant([2, 2]))
+    self.assertIsInstance(resized_images, tensor.Tensor)
+    self.assertTrue(resized_images.shape.is_compatible_with([3, 2, 2, 3]))
+
+  def testBadRank(self):
+    rt = ragged_tensor.RaggedTensor.from_tensor(array_ops.zeros([5, 5, 3]))
+    with self.assertRaisesRegex(ValueError, 'rank must be 4'):
+      image_ops.resize_images_v2(rt, [10, 10])
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_reverse_op_test.py

@@ -0,0 +1,85 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.reverse."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedReverseOpTest(test_util.TensorFlowTestCase,
+                          parameterized.TestCase):
+
+  @parameterized.parameters([
+      dict(
+          descr='Docstring example 1',
+          data=[[[1, 2], [3, 4]],
+                [[5, 6]],
+                [[7, 8], [9, 10], [11, 12]]],
+          axis=[0, 2],
+          expected=[[[8, 7], [10, 9], [12, 11]],
+                    [[6, 5]],
+                    [[2, 1], [4, 3]]]),
+      dict(
+          descr='data.shape=[5, (D2)]; axis=[0]',
+          data=[[1, 2], [3, 4, 5, 6], [7, 8, 9], [], [1, 2, 3]],
+          axis=[0],
+          expected=[[1, 2, 3], [], [7, 8, 9], [3, 4, 5, 6], [1, 2]]),
+      dict(
+          descr='data.shape=[5, (D2)]; axis=[1]',
+          data=[[1, 2], [3, 4, 5, 6], [7, 8, 9], [], [1, 2, 3]],
+          axis=[1],
+          expected=[[2, 1], [6, 5, 4, 3], [9, 8, 7], [], [3, 2, 1]]),
+      dict(
+          descr='data.shape=[5, (D2), (D3)]; axis=[0, -1]',
+          data=[[[1], [2, 3]], [[4, 5], [6, 7]], [[8]]],
+          axis=[0, -1],
+          expected=[[[8]], [[5, 4], [7, 6]], [[1], [3, 2]]]),
+      dict(
+          descr='data.shape=[2, (D2), 2]; axis=[2]',
+          data=[[[1, 2], [3, 4]], [[5, 6]]],
+          axis=[2],
+          expected=[[[2, 1], [4, 3]], [[6, 5]]],
+          ragged_rank=1),
+      dict(
+          descr='data.shape=[2, (D2), (D3)]; axis=[-1]',
+          data=[[[1, 2], [3, 4]], [[5, 6]]],
+          axis=[-1],
+          expected=[[[2, 1], [4, 3]], [[6, 5]]]),
+      dict(
+          descr='data.shape=[2, (D2), (D3)]; axis=[]',
+          data=[[[1, 2], [3, 4]], [[5, 6]]],
+          axis=[],
+          expected=[[[1, 2], [3, 4]], [[5, 6]]])
+  ])  # pyformat: disable
+  def testReverse(self, descr, data, axis, expected, ragged_rank=None):
+    data = ragged_factory_ops.constant(data, ragged_rank=ragged_rank)
+    result = ragged_array_ops.reverse(data, axis)
+    expected = ragged_factory_ops.constant(expected, ragged_rank=ragged_rank)
+    self.assertAllClose(result, expected)
+
+  def testErrors(self):
+    self.assertRaisesRegex(
+        TypeError, '`axis` must be a list of int or a constant tensor *',
+        ragged_array_ops.reverse,
+        ragged_factory_ops.constant([[1], [2, 3]], ragged_rank=1), [0, None])
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_row_lengths_op_test.py

@@ -0,0 +1,142 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged.row_lengths."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedRowLengthsOp(test_util.TensorFlowTestCase,
+                         parameterized.TestCase):
+
+  @parameterized.parameters([
+      # Docstring Example
+      dict(
+          rt_input=[[[3, 1, 4], [1]], [], [[5, 9], [2]], [[6]], []],
+          expected=[2, 0, 2, 1, 0]),
+      dict(
+          rt_input=[[[3, 1, 4], [1]], [], [[5, 9], [2]], [[6]], []],
+          axis=2,
+          expected=[[3, 1], [], [2, 1], [1], []]),
+
+      # 2D Tensor (1 ragged dimension)
+      dict(
+          rt_input=[['a'], ['b', 'c', 'd'], ['e'], [], ['f']],
+          expected=[1, 3, 1, 0, 1]),
+      dict(
+          rt_input=[['a'], ['b', 'c', 'd'], ['e'], [], ['f']],
+          axis=0,
+          expected=5),
+      dict(
+          rt_input=[['a', 'b', 'c', 'd', 'e', 'f', 'g']],
+          expected=[7]),
+      dict(
+          rt_input=[[], ['a', 'b', 'c', 'd', 'e', 'f', 'g'], []],
+          expected=[0, 7, 0]),
+      dict(
+          rt_input=[],
+          ragged_rank=1,
+          expected=[]),
+      dict(
+          rt_input=[],
+          ragged_rank=1,
+          axis=0,
+          expected=0),
+
+      # 3D Tensor (1 ragged dimension)
+      dict(
+          rt_input=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10]]],
+          ragged_rank=1,
+          axis=0,
+          expected=2),
+      dict(
+          rt_input=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10]]],
+          ragged_rank=1,
+          axis=1,
+          expected=[3, 2]),
+      dict(
+          rt_input=[[[1, 2], [3, 4], [5, 6]], [[7, 8], [9, 10]]],
+          ragged_rank=1,
+          axis=2,
+          expected=[[2, 2, 2], [2, 2]],
+          expected_ragged_rank=1),
+
+      # 3D Tensor (2 ragged dimensions)
+      dict(
+          rt_input=[[[1, 2], [3, 4, 5], []], [[6, 7, 8, 9], [10]]],
+          axis=0,
+          expected=2),
+      dict(
+          rt_input=[[[1, 2], [3, 4, 5], []], [[6, 7, 8, 9], [10]]],
+          axis=-3,
+          expected=2),
+      dict(
+          rt_input=[[[1, 2], [3, 4, 5], []], [[6, 7, 8, 9], [10]]],
+          axis=1,
+          expected=[3, 2]),
+      dict(
+          rt_input=[[[1, 2], [3, 4, 5], []], [[6, 7, 8, 9], [10]]],
+          axis=-2,
+          expected=[3, 2]),
+      dict(
+          rt_input=[[[1, 2], [3, 4, 5], []], [[6, 7, 8, 9], [10]]],
+          axis=2,
+          expected=[[2, 3, 0], [4, 1]],
+          expected_ragged_rank=1),
+      dict(
+          rt_input=[[[1, 2], [3, 4, 5], []], [[6, 7, 8, 9], [10]]],
+          axis=-1,
+          expected=[[2, 3, 0], [4, 1]],
+          expected_ragged_rank=1),
+  ])  # pyformat: disable
+  def testRowLengths(self,
+                     rt_input,
+                     expected,
+                     axis=1,
+                     ragged_rank=None,
+                     expected_ragged_rank=None):
+    rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
+    lengths = rt.row_lengths(axis)
+    self.assertAllEqual(lengths, expected)
+    if expected_ragged_rank is not None:
+      if isinstance(lengths, ragged_tensor.RaggedTensor):
+        self.assertEqual(lengths.ragged_rank, expected_ragged_rank)
+      else:
+        self.assertEqual(0, expected_ragged_rank)
+
+  @parameterized.parameters([
+      dict(  # axis=2 out of bounds: expected -2<=axis<2.
+          rt_input=[[10, 20], [30]],
+          axis=2,
+          exception=(ValueError, errors.InvalidArgumentError)),
+      dict(  # axis=-3 out of bounds: expected -2<=axis<2.
+          rt_input=[[2, 3, 0], [4, 1, 2]],
+          axis=-3,
+          exception=(ValueError, errors.InvalidArgumentError)),
+  ])
+  def testErrors(self, rt_input, exception, message=None, axis=1):
+    rt = ragged_factory_ops.constant(rt_input)
+    with self.assertRaisesRegex(exception, message):
+      rt.row_lengths(axis)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_row_splits_to_segment_ids_op_test.py

@@ -0,0 +1,50 @@
+# 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.
+# ==============================================================================
+"""Tests for the segment_id_ops.row_splits_to_segment_ids() op."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import segment_id_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSplitsToSegmentIdsOpTest(test_util.TensorFlowTestCase):
+
+  def testDocStringExample(self):
+    splits = [0, 3, 3, 5, 6, 9]
+    expected = [0, 0, 0, 2, 2, 3, 4, 4, 4]
+    segment_ids = segment_id_ops.row_splits_to_segment_ids(splits)
+    self.assertAllEqual(segment_ids, expected)
+
+  def testEmptySplits(self):
+    # Note: the splits for an empty ragged tensor contains a single zero.
+    segment_ids = segment_id_ops.row_splits_to_segment_ids([0])
+    self.assertAllEqual(segment_ids, [])
+
+  def testErrors(self):
+    self.assertRaisesRegex(ValueError, r'Invalid row_splits: \[\]',
+                           segment_id_ops.row_splits_to_segment_ids, [])
+    self.assertRaisesRegex(ValueError, r'splits must have dtype int32 or int64',
+                           segment_id_ops.row_splits_to_segment_ids,
+                           constant_op.constant([0.5]))
+    self.assertRaisesRegex(ValueError, r'Shape \(\) must have rank 1',
+                           segment_id_ops.row_splits_to_segment_ids, 0)
+    self.assertRaisesRegex(ValueError, r'Shape \(1, 1\) must have rank 1',
+                           segment_id_ops.row_splits_to_segment_ids, [[0]])
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_segment_ids_to_row_splits_op_test.py

@@ -0,0 +1,67 @@
+# 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.
+# ==============================================================================
+"""Tests for the segment_id_ops.segment_ids_to_row_splits() op."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import segment_id_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSplitsToSegmentIdsOpTest(test_util.TensorFlowTestCase):
+
+  def testDocStringExample(self):
+    segment_ids = [0, 0, 0, 2, 2, 3, 4, 4, 4]
+    expected = [0, 3, 3, 5, 6, 9]
+    splits = segment_id_ops.segment_ids_to_row_splits(segment_ids)
+    self.assertAllEqual(splits, expected)
+
+  def testEmptySegmentIds(self):
+    # Note: the splits for an empty ragged tensor contains a single zero.
+    segment_ids = segment_id_ops.segment_ids_to_row_splits([])
+    self.assertAllEqual(segment_ids, [0])
+
+  def testErrors(self):
+    self.assertRaisesRegex(
+        TypeError,
+        r'Argument `tensor` \(name\: segment_ids\) must be of type integer.*',
+        segment_id_ops.segment_ids_to_row_splits, constant_op.constant([0.5]))
+    self.assertRaisesRegex(ValueError, r'Shape \(\) must have rank 1',
+                           segment_id_ops.segment_ids_to_row_splits, 0)
+    self.assertRaisesRegex(ValueError, r'Shape \(1, 1\) must have rank 1',
+                           segment_id_ops.segment_ids_to_row_splits, [[0]])
+
+  def testNumSegments(self):
+    segment_ids = [0, 0, 0, 2, 2, 3, 4, 4, 4]
+    num_segments = 7
+    expected = [0, 3, 3, 5, 6, 9, 9, 9]
+    splits = segment_id_ops.segment_ids_to_row_splits(segment_ids, num_segments)
+    self.assertAllEqual(splits, expected)
+
+  def testUnsortedSegmentIds(self):
+    # Segment ids are not required to be sorted.
+    segment_ids = [0, 4, 3, 2, 4, 4, 2, 0, 0]
+    splits1 = segment_id_ops.segment_ids_to_row_splits(segment_ids)
+    expected1 = [0, 3, 3, 5, 6, 9]
+
+    splits2 = segment_id_ops.segment_ids_to_row_splits(segment_ids, 7)
+    expected2 = [0, 3, 3, 5, 6, 9, 9, 9]
+    self.assertAllEqual(splits1, expected1)
+    self.assertAllEqual(splits2, expected2)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_segment_op_test.py

@@ -0,0 +1,217 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_range op."""
+
+import math
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+def prod(values):
+  val = 1
+  for v in values:
+    val *= v
+  return val
+  # return reduce(lambda x, y: x * y, values, 1)
+
+
+def mean(values):
+  return 1.0 * sum(values) / len(values)
+
+
+def sqrt_n(values):
+  return 1.0 * sum(values) / math.sqrt(len(values))
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSegmentOpsTest(test_util.TensorFlowTestCase,
+                           parameterized.TestCase):
+
+  def expected_value(self, data, segment_ids, num_segments, combiner):
+    """Find the expected value for a call to ragged_segment_<aggregate>.
+
+    Args:
+      data: The input RaggedTensor, expressed as a nested python list.
+      segment_ids: The segment ids, as a python list of ints.
+      num_segments: The number of segments, as a python int.
+      combiner: The Python function used to combine values.
+    Returns:
+      The expected value, as a nested Python list.
+    """
+    self.assertLen(data, len(segment_ids))
+
+    # Build an empty (num_segments x ncols) "grouped" matrix
+    ncols = max(len(row) for row in data)
+    grouped = [[[] for _ in range(ncols)] for row in range(num_segments)]
+
+    # Append values from data[row] to grouped[segment_ids[row]]
+    for row in range(len(data)):
+      for col in range(len(data[row])):
+        grouped[segment_ids[row]][col].append(data[row][col])
+
+    # Combine the values.
+    return [[combiner(values)
+             for values in grouped_row
+             if values]
+            for grouped_row in grouped]
+
+  @parameterized.parameters(
+      (ragged_math_ops.segment_sum, sum, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_sum, sum, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_sum, sum, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_sum, sum, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_prod, prod, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_prod, prod, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_prod, prod, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_prod, prod, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_min, min, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_min, min, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_min, min, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_min, min, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_max, max, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_max, max, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_max, max, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_max, max, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_mean, mean, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_mean, mean, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_mean, mean, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_mean, mean, [0, 0, 0, 10, 10, 10]),
+  )
+  def testRaggedSegment_Int(self, segment_op, combiner, segment_ids):
+    rt_as_list = [[0, 1, 2, 3], [4], [], [5, 6], [7], [8, 9]]
+    rt = ragged_factory_ops.constant(rt_as_list)
+    num_segments = max(segment_ids) + 1
+    expected = self.expected_value(rt_as_list, segment_ids, num_segments,
+                                   combiner)
+
+    segmented = segment_op(rt, segment_ids, num_segments)
+    self.assertAllEqual(segmented, expected)
+
+  @parameterized.parameters(
+      (ragged_math_ops.segment_sum, sum, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_sum, sum, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_sum, sum, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_sum, sum, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_prod, prod, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_prod, prod, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_prod, prod, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_prod, prod, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_min, min, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_min, min, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_min, min, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_min, min, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_max, max, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_max, max, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_max, max, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_max, max, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_mean, mean, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_mean, mean, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_mean, mean, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_mean, mean, [0, 0, 0, 10, 10, 10]),
+      (ragged_math_ops.segment_sqrt_n, sqrt_n, [0, 0, 1, 1, 2, 2]),
+      (ragged_math_ops.segment_sqrt_n, sqrt_n, [0, 0, 0, 1, 1, 1]),
+      (ragged_math_ops.segment_sqrt_n, sqrt_n, [5, 4, 3, 2, 1, 0]),
+      (ragged_math_ops.segment_sqrt_n, sqrt_n, [0, 0, 0, 10, 10, 10]),
+  )
+  def testRaggedSegment_Float(self, segment_op, combiner, segment_ids):
+    rt_as_list = [[0., 1., 2., 3.], [4.], [], [5., 6.], [7.], [8., 9.]]
+    rt = ragged_factory_ops.constant(rt_as_list)
+    num_segments = max(segment_ids) + 1
+    expected = self.expected_value(rt_as_list, segment_ids, num_segments,
+                                   combiner)
+
+    segmented = segment_op(rt, segment_ids, num_segments)
+    self.assertAllClose(segmented, expected)
+
+  def testRaggedRankTwo(self):
+    rt = ragged_factory_ops.constant([
+        [[111, 112, 113, 114], [121],],  # row 0
+        [],                              # row 1
+        [[], [321, 322], [331]],         # row 2
+        [[411, 412]]                     # row 3
+    ])  # pyformat: disable
+    segment_ids1 = [0, 2, 2, 2]
+    segmented1 = ragged_math_ops.segment_sum(rt, segment_ids1, 3)
+    expected1 = [[[111, 112, 113, 114], [121]],     # row 0
+                 [],                                # row 1
+                 [[411, 412], [321, 322], [331]]    # row 2
+                ]  # pyformat: disable
+    self.assertAllEqual(segmented1, expected1)
+
+    segment_ids2 = [1, 2, 1, 1]
+    segmented2 = ragged_math_ops.segment_sum(rt, segment_ids2, 3)
+    expected2 = [[],
+                 [[111+411, 112+412, 113, 114], [121+321, 322], [331]],
+                 []]  # pyformat: disable
+    self.assertAllEqual(segmented2, expected2)
+
+  def testRaggedSegmentIds(self):
+    rt = ragged_factory_ops.constant([
+        [[111, 112, 113, 114], [121],],  # row 0
+        [],                              # row 1
+        [[], [321, 322], [331]],         # row 2
+        [[411, 412]]                     # row 3
+    ])  # pyformat: disable
+    segment_ids = ragged_factory_ops.constant([[1, 2], [], [1, 1, 2], [2]])
+    segmented = ragged_math_ops.segment_sum(rt, segment_ids, 3)
+    expected = [[],
+                [111+321, 112+322, 113, 114],
+                [121+331+411, 412]]  # pyformat: disable
+    self.assertAllEqual(segmented, expected)
+
+  def testShapeMismatchError1(self):
+    dt = constant_op.constant([1, 2, 3, 4, 5, 6])
+    segment_ids = ragged_factory_ops.constant([[1, 2], []])
+    self.assertRaisesRegex(
+        ValueError, 'segment_ids.shape must be a prefix of data.shape, '
+        'but segment_ids is ragged and data is not.',
+        ragged_math_ops.segment_sum, dt, segment_ids, 3)
+
+  def testShapeMismatchError2(self):
+    rt = ragged_factory_ops.constant([
+        [[111, 112, 113, 114], [121]],  # row 0
+        [],                             # row 1
+        [[], [321, 322], [331]],        # row 2
+        [[411, 412]]                    # row 3
+    ])  # pyformat: disable
+    segment_ids = ragged_factory_ops.constant([[1, 2], [1], [1, 1, 2], [2]])
+
+    # Error is raised at graph-building time if we can detect it then.
+    self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        'segment_ids.shape must be a prefix of data.shape.*',
+        ragged_math_ops.segment_sum, rt, segment_ids, 3)
+
+    # Otherwise, error is raised when we run the graph.
+    segment_ids2 = ragged_tensor.RaggedTensor.from_row_splits(
+        array_ops.placeholder_with_default(segment_ids.values, None),
+        array_ops.placeholder_with_default(segment_ids.row_splits, None))
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        'segment_ids.shape must be a prefix of data.shape.*'):
+      self.evaluate(ragged_math_ops.segment_sum(rt, segment_ids2, 3))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_size_op_test.py

@@ -0,0 +1,43 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged.size."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSizeOpTest(test_util.TensorFlowTestCase,
+                       parameterized.TestCase):
+
+  @parameterized.parameters([
+      {'size': 1, 'test_input': 1},
+      {'size': 0, 'test_input': []},
+      {'size': 0, 'test_input': [], 'ragged_rank': 1},
+      {'size': 3, 'test_input': [1, 1, 1]},
+      {'size': 3, 'test_input': [[1, 1], [1]]},
+      {'size': 5, 'test_input': [[[1, 1, 1], [1]], [[1]]]},
+      {'size': 6, 'test_input': [[[1, 1], [1, 1]], [[1, 1]]], 'ragged_rank': 1},
+  ])
+  def testRaggedSize(self, test_input, size, ragged_rank=None):
+    input_rt = ragged_factory_ops.constant(test_input, ragged_rank=ragged_rank)
+    self.assertAllEqual(ragged_array_ops.size(input_rt), size)
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_split_op_test.py

@@ -0,0 +1,480 @@
+# Copyright 2021 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.split."""
+
+import itertools
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_spec
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSplitOpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Uniform splits.
+      #=========================================================================
+      dict(
+          descr='Uniform splits, rank-2 inputs, axis=0',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=2,
+          expected=[
+              [[1]],
+              [[2, 3, 4]]]),
+      dict(
+          descr='Uniform 3 splits, rank-2 inputs, axis=0',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 2, 1],  # shape=(3, None)
+          num_or_size_splits=3,
+          expected=[
+              [[1]],
+              [[2, 3]],
+              [[4]]]),
+      dict(
+          descr='Uniform 5 splits, rank-2 inputs, axis=0',
+          pylist=[1, 2, 3, 4, 5],
+          row_lengths=[1, 1, 1, 1, 1],  # shape=(5, None)
+          num_or_size_splits=5,
+          expected=[
+              [[1]],
+              [[2]],
+              [[3]],
+              [[4]],
+              [[5]]]),
+      dict(
+          descr='Uniform 2 splits, rank-2 inputs(empty), axis=0',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[4, 0],  # shape=(2, None)
+          num_or_size_splits=2,
+          expected=[
+              [[1, 2, 3, 4]],
+              [[]]]),
+      dict(
+          descr='Uniform 2 splits, rank-2 inputs(all empty), axis=0',
+          pylist=[],
+          row_lengths=[0, 0],  # shape=(2, None)
+          num_or_size_splits=2,
+          expected=[
+              [[]],
+              [[]]]),
+      dict(
+          descr='Uniform 1 split, rank-2 inputs, axis=0',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=1,
+          expected=[
+              [[1], [2, 3, 4]]]),
+      dict(
+          descr='Uniform 1 split, rank-2 inputs, axis=1',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=1,
+          axis=1,
+          expected=[
+              [[1], [2, 3, 4]]]),
+      dict(
+          descr='Uniform 2 split, rank-3 inputs, axis=0',
+          pylist=np.arange(4 * 2).reshape(4, 2),
+          row_lengths=[1, 3],  # shape=(2, None, 2)
+          num_or_size_splits=2,
+          expected=[
+              [[[0, 1]]],
+              [[[2, 3], [4, 5], [6, 7]]]]),
+      dict(
+          descr='Uniform 2 splits, rank-3 inputs, axis=2',
+          pylist=np.arange(4 * 2).reshape(4, 2),
+          row_lengths=[1, 3],  # shape=(2, None, 2)
+          num_or_size_splits=2,
+          axis=2,
+          expected=[
+              [[[0]], [[2], [4], [6]]],
+              [[[1]], [[3], [5], [7]]]]),
+      dict(
+          descr='Uniform 2 splits, rank-2 float inputs, axis=0',
+          pylist=[1.0, 2.0, 3.0, 4.0],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=2,
+          expected=[
+              [[1.0]],
+              [[2.0, 3.0, 4.0]]]),
+      dict(
+          descr='Uniform 2 splits, rank-2 string inputs, axis=0',
+          pylist=[b'a', b'bc', b'', b'd'],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=2,
+          expected=[
+              [[b'a']],
+              [[b'bc', b'', b'd']]]),
+      #=========================================================================
+      # Ragged splits.
+      #=========================================================================
+      dict(
+          descr='Ragged 2 splits, rank-2 inputs, axis=0',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=[1, 1],
+          expected=[
+              [[1]],
+              [[2, 3, 4]]]),
+      dict(
+          descr='Ragged 3 splits, rank-2 inputs, axis=0',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 2, 1],  # shape=(3, None)
+          num_or_size_splits=[1, 2],
+          expected=[
+              [[1]],
+              [[2, 3], [4]]]),
+      dict(
+          descr='Ragged 5 splits, rank-2 inputs(empty), axis=0',
+          pylist=[1, 2, 3, 4, 5],
+          row_lengths=[1, 1, 1, 1, 1],  # shape=(5, None)
+          num_or_size_splits=[1, 2, 2, 0],
+          expected=[
+              [[1]],
+              [[2], [3]],
+              [[4], [5]],
+              []]),
+      dict(
+          descr='Ragged 2 splits, rank-2 inputs(empty), axis=0',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[4, 0, 0, 0],  # shape=(2, None)
+          num_or_size_splits=[3, 1],
+          expected=[
+              [[1, 2, 3, 4], [], []],
+              [[]]]),
+      dict(
+          descr='Ragged 2 splits, rank-2 inputs(all empty), axis=0',
+          pylist=[],
+          row_lengths=[0, 0],  # shape=(2, None)
+          num_or_size_splits=[2, 0],
+          expected=[
+              [[], []],
+              []]),
+      dict(
+          descr='Ragged 1 split, rank-2 inputs, axis=0',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=[2],
+          expected=[
+              [[1], [2, 3, 4]]]),
+      dict(
+          descr='Ragged 2 split, rank-3 inputs, axis=0',
+          pylist=np.arange(4 * 2).reshape(4, 2),
+          row_lengths=[1, 3],  # shape=(2, None, 2)
+          num_or_size_splits=[1, 1],
+          expected=[
+              [[[0, 1]]],
+              [[[2, 3], [4, 5], [6, 7]]]]),
+      dict(
+          descr='Ragged 2 split, rank-3 inputs, axis=-3',
+          pylist=np.arange(4 * 2).reshape(4, 2),
+          row_lengths=[1, 3],  # shape=(2, None, 2)
+          num_or_size_splits=[1, 1],
+          expected=[
+              [[[0, 1]]],
+              [[[2, 3], [4, 5], [6, 7]]]]),
+      dict(
+          descr='Ragged 2 splits, rank-3 inputs, axis=2',
+          pylist=np.arange(4 * 3).reshape(4, 3),
+          row_lengths=[1, 3],  # shape=(2, None, 3)
+          num_or_size_splits=[2, 1],
+          axis=2,
+          expected=[
+              [[[0, 1]], [[3, 4], [6, 7], [9, 10]]],
+              [[[2]], [[5], [8], [11]]]]),
+      dict(
+          descr='Ragged 2 splits, rank-3 inputs, axis=-1',
+          pylist=np.arange(4 * 3).reshape(4, 3),
+          row_lengths=[1, 3],  # shape=(2, None, 3)
+          num_or_size_splits=[2, 1],
+          axis=2,
+          expected=[
+              [[[0, 1]], [[3, 4], [6, 7], [9, 10]]],
+              [[[2]], [[5], [8], [11]]]]),
+      dict(
+          descr='Ragged 3 splits, rank-2 float inputs, axis=0',
+          pylist=[1.0, 2.0, 3.0, 4.0],
+          row_lengths=[1, 2, 1],  # shape=(2, None)
+          num_or_size_splits=[2, 1],
+          expected=[
+              [[1.0], [2.0, 3.0]],
+              [[4.0]]]),
+      dict(
+          descr='Ragged 3 splits with name, rank-2 float inputs, axis=0',
+          pylist=[1.0, 2.0, 3.0, 4.0],
+          row_lengths=[1, 2, 1],  # shape=(2, None)
+          num_or_size_splits=[2, 1],
+          name='ragged_split',
+          expected=[
+              [[1.0], [2.0, 3.0]],
+              [[4.0]]]),
+      dict(
+          descr='Ragged 3 splits with num, rank-2 float inputs, axis=0',
+          pylist=[1.0, 2.0, 3.0, 4.0],
+          row_lengths=[1, 2, 1],  # shape=(2, None)
+          num_or_size_splits=[2, 1],
+          num=2,
+          expected=[
+              [[1.0], [2.0, 3.0]],
+              [[4.0]]]),
+      dict(
+          descr='Ragged 2 splits, rank-2 string inputs, axis=0',
+          pylist=[b'a', b'bc', b'', b'd'],
+          row_lengths=[1, 3, 0],  # shape=(2, None)
+          num_or_size_splits=[2, 1],
+          expected=[
+              [[b'a'], [b'bc', b'', b'd']],
+              [[]]]),
+  ])  # pyformat: disable
+  def testSplit(self,
+                descr,
+                pylist,
+                row_lengths,
+                num_or_size_splits,
+                expected,
+                axis=0,
+                num=None,
+                name=None):
+    rt = ragged_tensor.RaggedTensor.from_row_lengths(pylist, row_lengths)
+    result = ragged_array_ops.split(rt, num_or_size_splits, axis, num, name)
+    self.assertLen(result, len(expected))
+    for res, exp in zip(result, expected):
+      self.assertAllEqual(res, exp)
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Uniform splits errors.
+      #=========================================================================
+      dict(
+          descr='Uniform split, can not split',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=7,
+          exception=errors.InvalidArgumentError,
+          message='Cannot exactly split'),
+      dict(
+          descr='Uniform split, ragged dimension',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=2,
+          axis=1,
+          exception=ValueError,
+          message='ragged dimension'),
+      dict(
+          descr='Uniform split, zero split',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=0,
+          exception=ValueError,
+          message='must be >=1'),
+      #=========================================================================
+      # Ragged splits errors.
+      #=========================================================================
+      dict(
+          descr='Ragged split, 2 dimensional size_splits',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=[[1, 1]],
+          exception=TypeError,
+          message='Python list'),
+      dict(
+          descr='Ragged split, ragged dimension',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=[1, 1],
+          axis=1,
+          exception=ValueError,
+          message='ragged dimension'),
+      dict(
+          descr='Ragged split, cannot split',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=[1, 2],
+          exception=errors.InvalidArgumentError,
+          message='Cannot exactly split'),
+      dict(
+          descr='Ragged split, num does not match',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=[1, 1],
+          num=3,
+          exception=ValueError,
+          message='`num` does not match'),
+      dict(
+          descr='Ragged split, negative split',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=[1, -1, 2],
+          num=3,
+          exception=errors.InvalidArgumentError,
+          message='must be non-negative'),
+      dict(
+          descr='Ragged split, float splits',
+          pylist=[1, 2, 3, 4],
+          row_lengths=[1, 3],  # shape=(2, None)
+          num_or_size_splits=[1.0, 2.0],
+          num=2,
+          exception=TypeError,
+          message='integer'),
+  ])  # pyformat: disable
+  def testSplitError(self,
+                     descr,
+                     pylist,
+                     row_lengths,
+                     num_or_size_splits,
+                     exception,
+                     message,
+                     axis=0,
+                     num=None):
+    rt = ragged_tensor.RaggedTensor.from_row_lengths(pylist, row_lengths)
+    with self.assertRaises(exception):
+      result = ragged_array_ops.split(rt, num_or_size_splits, axis, num)
+      self.evaluate(result)
+
+  @parameterized.named_parameters([
+      ('int32', dtypes.int32),
+      ('int64', dtypes.int64)])
+  def testSplitTensorDtype(self, dtype):
+    rt = ragged_tensor.RaggedTensor.from_row_lengths([1.0, 2.0, 3.0, 4.0],
+                                                     [3, 1])
+    # split_lengths is a 1-D tensor
+    split_lengths = ops.convert_to_tensor([1, 1], dtype=dtype)
+    result = ragged_array_ops.split(rt, split_lengths)
+    expected = [
+        ragged_tensor.RaggedTensor.from_row_lengths([1.0, 2.0, 3.0], [3]),
+        ragged_tensor.RaggedTensor.from_row_lengths([4.0], [1])]
+    self.assertLen(result, len(expected))
+    for res, exp in zip(result, expected):
+      self.assertAllEqual(res, exp)
+
+  @parameterized.parameters([
+      dict(rt_shape=(2, None)),
+      dict(rt_shape=None),
+  ])
+  def testUniformSplitDynamicShape(self, rt_shape):
+    rt = ragged_tensor.RaggedTensor.from_row_lengths([1.0, 2.0, 3.0, 4.0],
+                                                     [3, 1])
+    rt_spec = ragged_tensor.RaggedTensorSpec(rt_shape, ragged_rank=1)
+    @def_function.function(input_signature=[rt_spec])
+    def split_tensors(rt):
+      return ragged_array_ops.split(rt, 2)
+
+    splited_rts = split_tensors(rt)
+    expected_rts = [
+        ragged_tensor.RaggedTensor.from_row_lengths([1.0, 2.0, 3.0], [3]),
+        ragged_tensor.RaggedTensor.from_row_lengths([4.0], [1])]
+    for splited_rt, expected_rt in zip(splited_rts, expected_rts):
+      self.assertAllEqual(splited_rt, expected_rt)
+
+  @parameterized.parameters([
+      dict(rt_shape=x, lengths_shape=y) for x, y in itertools.product(
+          [(2, None), None],
+          [(2,), (None,), None])
+  ])
+  def testRaggedSplitDynamicShape(self, rt_shape, lengths_shape):
+    rt_spec = ragged_tensor.RaggedTensorSpec(rt_shape, ragged_rank=1)
+    lengths_spec = tensor_spec.TensorSpec(lengths_shape, dtype=dtypes.int32)
+    @def_function.function(input_signature=[rt_spec, lengths_spec])
+    def split_tensors(rt, split_lengths):
+      return ragged_array_ops.split(rt, split_lengths, num=2)
+
+    rt = ragged_tensor.RaggedTensor.from_row_lengths([1.0, 2.0, 3.0, 4.0],
+                                                     [3, 1])
+    split_lengths = [1, 1]
+    # split_lengths matches num at runtime
+    splited_rts = split_tensors(rt, split_lengths)
+    expected_rts = [
+        ragged_tensor.RaggedTensor.from_row_lengths([1.0, 2.0, 3.0], [3]),
+        ragged_tensor.RaggedTensor.from_row_lengths([4.0], [1])]
+    for splited_rt, expected_rt in zip(splited_rts, expected_rts):
+      self.assertAllEqual(splited_rt, expected_rt)
+
+  @parameterized.parameters([
+      dict(
+          descr='lengths known rank, num and lengths mismatch',
+          rt_shape=(2, None),
+          lengths_shape=(None,),
+          lengths=[1, 1, 0],
+          num=2,
+          exception=errors.InvalidArgumentError,
+          message='inconsistent'),
+      dict(
+          descr='lengths unknown rank, num and lengths mismatch',
+          rt_shape=None,
+          lengths_shape=None,
+          lengths=[1, 1, 0],
+          num=2,
+          exception=errors.InvalidArgumentError,
+          message='inconsistent'),
+      dict(
+          descr='rt unknown rank, negative axis',
+          rt_shape=None,
+          lengths_shape=None,
+          lengths=[1, 1],
+          axis=-2,
+          num=2,
+          exception=ValueError,
+          message='negative'),
+      dict(
+          descr='lengths unknown rank, num is None',
+          rt_shape=None,
+          lengths_shape=None,
+          lengths=[1, 1],
+          exception=ValueError,
+          message='`num` must be specified'),
+      dict(
+          descr='lengths unknown rank, dynamic rank!=1',
+          rt_shape=None,
+          lengths_shape=None,
+          lengths=[[1, 1]],
+          num=2,
+          exception=(ValueError, errors.InvalidArgumentError)),
+  ])
+  def testRaggedSplitDynamicShapeError(self,
+                                       descr,
+                                       rt_shape,
+                                       lengths_shape,
+                                       lengths,
+                                       exception,
+                                       message='',
+                                       axis=0,
+                                       num=None):
+    rt_spec = ragged_tensor.RaggedTensorSpec(rt_shape, ragged_rank=1)
+    split_lengths_spec = tensor_spec.TensorSpec(lengths_shape,
+                                                dtype=dtypes.int32)
+    @def_function.function(input_signature=[rt_spec, split_lengths_spec])
+    def split_tensors(rt, split_lengths):
+      return ragged_array_ops.split(rt, split_lengths, axis=axis, num=num)
+
+    rt = ragged_tensor.RaggedTensor.from_row_lengths([1.0, 2.0, 3.0, 4.0],
+                                                     [3, 1])
+    with self.assertRaisesRegex(exception, message):
+      self.evaluate(split_tensors(rt=rt, split_lengths=lengths))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_squeeze_op.py

@@ -0,0 +1,133 @@
+# 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.
+# ==============================================================================
+"""Operator Squeeze for RaggedTensors."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
+
+
+@dispatch.dispatch_for_api(array_ops.squeeze_v2)
+def squeeze(input: ragged_tensor.Ragged, axis=None, name=None):  # pylint: disable=redefined-builtin
+  """Ragged compatible squeeze.
+
+  If `input` is a `tf.Tensor`, then this calls `tf.squeeze`.
+
+  If `input` is a `tf.RaggedTensor`, then this operation takes `O(N)` time,
+  where `N` is the number of elements in the squeezed dimensions.
+
+  Args:
+    input: A potentially ragged tensor. The input to squeeze.
+    axis: An optional list of ints. Defaults to `None`. If the `input` is
+      ragged, it only squeezes the dimensions listed. It fails if `input` is
+      ragged and axis is []. If `input` is not ragged it calls tf.squeeze. Note
+      that it is an error to squeeze a dimension that is not 1. It must be in
+      the range of [-rank(input), rank(input)).
+   name: A name for the operation (optional).
+
+  Returns:
+    A potentially ragged tensor. Contains the same data as input,
+    but has one or more dimensions of size 1 removed.
+  """
+  with ops.name_scope(name, 'RaggedSqueeze', [input]):
+    input = ragged_tensor.convert_to_tensor_or_ragged_tensor(input)
+    if isinstance(input, tensor.Tensor):
+      return array_ops.squeeze(input, axis, name)
+
+    if axis is None:
+      raise ValueError('Ragged.squeeze must have an axis argument.')
+    if isinstance(axis, int):
+      axis = [axis]
+    elif ((not isinstance(axis, (list, tuple))) or
+          (not all(isinstance(d, int) for d in axis))):
+      raise TypeError('Axis must be a list or tuple of integers.')
+
+    dense_dims = []
+    ragged_dims = []
+    # Normalize all the dims in axis to be positive
+    axis = [
+        array_ops.get_positive_axis(d, input.shape.ndims, 'axis[%d]' % i,
+                                    'rank(input)') for i, d in enumerate(axis)
+    ]
+    for dim in axis:
+      if dim > input.ragged_rank:
+        dense_dims.append(dim - input.ragged_rank)
+      else:
+        ragged_dims.append(dim)
+
+    # Make sure the specified ragged dimensions are squeezable.
+    assertion_list = []
+    scalar_tensor_one = constant_op.constant(1, dtype=input.row_splits.dtype)
+    for i, r in enumerate(input.nested_row_lengths()):
+      if i + 1 in ragged_dims:
+        assertion_list.append(
+            control_flow_assert.Assert(
+                math_ops.reduce_all(math_ops.equal(r, scalar_tensor_one)),
+                ['the given axis (axis = %d) is not squeezable!' % (i + 1)]))
+    if 0 in ragged_dims:
+      scalar_tensor_two = constant_op.constant(2, dtype=dtypes.int32)
+      assertion_list.append(
+          control_flow_assert.Assert(
+              math_ops.equal(
+                  array_ops.size(input.row_splits), scalar_tensor_two),
+              ['the given axis (axis = 0) is not squeezable!']))
+
+    # Till now, we are sure that the ragged dimensions are squeezable.
+    squeezed_rt = None
+    squeezed_rt = control_flow_ops.with_dependencies(assertion_list,
+                                                     input.flat_values)
+
+    if dense_dims:
+      # Gives error if the dense dimension is not squeezable.
+      squeezed_rt = array_ops.squeeze(squeezed_rt, dense_dims)
+
+    remaining_row_splits = []
+    remaining_row_splits = list()
+    for i, row_split in enumerate(input.nested_row_splits):
+      # each row_splits tensor is for dimension #(i+1) .
+      if (i + 1) not in ragged_dims:
+        remaining_row_splits.append(row_split)
+    # Take care of the first row if it is to be squeezed.
+    if remaining_row_splits and 0 in ragged_dims:
+      remaining_row_splits.pop(0)
+
+    squeezed_rt = RaggedTensor.from_nested_row_splits(squeezed_rt,
+                                                      remaining_row_splits)
+
+    # Corner case: when removing all the ragged dimensions and the output is
+    # a scalar tensor e.g. ragged.squeeze(ragged.constant([[[1]]])).
+    if set(range(0, input.ragged_rank + 1)).issubset(set(ragged_dims)):
+      squeezed_rt = array_ops.squeeze(squeezed_rt, [0], name)
+
+    return squeezed_rt
+
+
+@dispatch.dispatch_for_api(array_ops.squeeze)
+def _ragged_squeeze_v1(input: ragged_tensor.Ragged,  # pylint: disable=redefined-builtin
+                       axis=None,
+                       name=None,
+                       squeeze_dims=None):
+  axis = deprecation.deprecated_argument_lookup('axis', axis, 'squeeze_dims',
+                                                squeeze_dims)
+  return squeeze(input, axis, name)

tensorflow/python/ops/ragged/ragged_squeeze_op_test.py

@@ -0,0 +1,287 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged.squeeze."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_conversion_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_squeeze_op
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedSqueezeTest(test_util.TensorFlowTestCase,
+                        parameterized.TestCase):
+
+  @parameterized.parameters([
+      {
+          'input_list': []
+      },
+      {
+          'input_list': [[]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[[[], []], [[], []]]],
+          'squeeze_ranks': [0]
+      },
+  ])
+  def test_passing_empty(self, input_list, squeeze_ranks=None):
+    rt = ragged_squeeze_op.squeeze(
+        ragged_factory_ops.constant(input_list), squeeze_ranks)
+    dt = array_ops.squeeze(constant_op.constant(input_list), squeeze_ranks)
+    self.assertAllEqual(ragged_conversion_ops.to_tensor(rt), dt)
+
+  @parameterized.parameters([
+      {
+          'input_list': [[1]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[1]],
+          'squeeze_ranks': [0, 1]
+      },
+      {
+          'input_list': [[1, 2]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[1], [2]],
+          'squeeze_ranks': [1]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [1]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [3]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0, 3]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0, 1]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [1, 3]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0, 1, 3]
+      },
+      {
+          'input_list': [[[1], [2]], [[3], [4]]],
+          'squeeze_ranks': [2]
+      },
+      {
+          'input_list': [[1], [2]],
+          'squeeze_ranks': [-1]
+      },
+  ])
+  def test_passing_simple(self, input_list, squeeze_ranks=None):
+    rt = ragged_squeeze_op.squeeze(
+        ragged_factory_ops.constant(input_list), squeeze_ranks)
+    dt = array_ops.squeeze(constant_op.constant(input_list), squeeze_ranks)
+    self.assertAllEqual(ragged_conversion_ops.to_tensor(rt), dt)
+
+  @parameterized.parameters([
+      # ragged_conversion_ops.from_tensor does not work for this
+      # {'input_list': [1]},
+      {
+          'input_list': [[1]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[1, 2]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[1], [2]],
+          'squeeze_ranks': [1]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [1]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [3]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0, 3]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0, 1]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [1, 3]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0, 1, 3]
+      },
+      {
+          'input_list': [[[1], [2]], [[3], [4]]],
+          'squeeze_ranks': [2]
+      },
+  ])
+  def test_passing_simple_from_dense(self, input_list, squeeze_ranks=None):
+    dt = constant_op.constant(input_list)
+    rt = ragged_conversion_ops.from_tensor(dt)
+    rt_s = ragged_squeeze_op.squeeze(rt, squeeze_ranks)
+    dt_s = array_ops.squeeze(dt, squeeze_ranks)
+    self.assertAllEqual(ragged_conversion_ops.to_tensor(rt_s), dt_s)
+
+  @parameterized.parameters([
+      {
+          'input_list': [[[[[[1]], [[1, 2]]]], [[[[]], [[]]]]]],
+          'output_list': [[[1], [1, 2]], [[], []]],
+          'squeeze_ranks': [0, 2, 4]
+      },
+      {
+          'input_list': [[[[[[1]], [[1, 2]]]], [[[[]], [[]]]]]],
+          'output_list': [[[[[1]], [[1, 2]]]], [[[[]], [[]]]]],
+          'squeeze_ranks': [0]
+      },
+  ])
+  def test_passing_ragged(self, input_list, output_list, squeeze_ranks=None):
+    rt = ragged_factory_ops.constant(input_list)
+    rt_s = ragged_squeeze_op.squeeze(rt, squeeze_ranks)
+    ref = ragged_factory_ops.constant(output_list)
+    self.assertAllEqual(rt_s, ref)
+
+  def test_passing_text(self):
+    rt = ragged_factory_ops.constant([[[[[[[['H']], [['e']], [['l']], [['l']],
+                                           [['o']]],
+                                          [[['W']], [['o']], [['r']], [['l']],
+                                           [['d']], [['!']]]]],
+                                        [[[[['T']], [['h']], [['i']], [['s']]],
+                                          [[['i']], [['s']]],
+                                          [[['M']], [['e']], [['h']], [['r']],
+                                           [['d']], [['a']], [['d']]],
+                                          [[['.']]]]]]]])
+    output_list = [[['H', 'e', 'l', 'l', 'o'], ['W', 'o', 'r', 'l', 'd', '!']],
+                   [['T', 'h', 'i', 's'], ['i', 's'],
+                    ['M', 'e', 'h', 'r', 'd', 'a', 'd'], ['.']]]
+    ref = ragged_factory_ops.constant(output_list)
+    rt_s = ragged_squeeze_op.squeeze(rt, [0, 1, 3, 6, 7])
+    self.assertAllEqual(rt_s, ref)
+
+  @parameterized.parameters([
+      {
+          'input_list': [[]],
+          'squeeze_ranks': [1]
+      },
+      {
+          'input_list': [[1, 2]],
+          'squeeze_ranks': [1]
+      },
+      {
+          'input_list': [[1], [2]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0, 2]
+      },
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [2]
+      },
+      {
+          'input_list': [[[1], [2]], [[3], [4]]],
+          'squeeze_ranks': [0]
+      },
+      {
+          'input_list': [[[1], [2]], [[3], [4]]],
+          'squeeze_ranks': [1]
+      },
+      {
+          'input_list': [[], []],
+          'squeeze_ranks': [1]
+      },
+      {
+          'input_list': [[[], []], [[], []]],
+          'squeeze_ranks': [1]
+      },
+  ])
+  def test_failing_InvalidArgumentError(self, input_list, squeeze_ranks):
+    with self.assertRaises(errors.InvalidArgumentError):
+      self.evaluate(
+          ragged_squeeze_op.squeeze(
+              ragged_factory_ops.constant(input_list), squeeze_ranks))
+
+  @parameterized.parameters([
+      {
+          'input_list': [[]]
+      },
+      {
+          'input_list': [[1]]
+      },
+      {
+          'input_list': [[1, 2]]
+      },
+      {
+          'input_list': [[[1], [2]], [[3], [4]]]
+      },
+      {
+          'input_list': [[1]]
+      },
+      {
+          'input_list': [[[1], [2]], [[3], [4]]]
+      },
+      {
+          'input_list': [[[[12], [11]]]]
+      },
+  ])
+  def test_failing_no_squeeze_dim_specified(self, input_list):
+    with self.assertRaises(ValueError):
+      ragged_squeeze_op.squeeze(ragged_factory_ops.constant(input_list))
+
+  @parameterized.parameters([
+      {
+          'input_list': [[[[12], [11]]]],
+          'squeeze_ranks': [0, 1, 3]
+      },
+  ])
+  def test_failing_axis_is_not_a_list(self, input_list, squeeze_ranks):
+    with self.assertRaises(TypeError):
+      tensor_ranks = constant_op.constant(squeeze_ranks)
+      ragged_squeeze_op.squeeze(
+          ragged_factory_ops.constant(input_list), tensor_ranks)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_stack_op_test.py

@@ -0,0 +1,400 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_concat_ops.stack."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedStackOpTest(test_util.TensorFlowTestCase,
+                        parameterized.TestCase):
+
+  @parameterized.parameters(
+      dict(
+          descr='One rank-2 input (ragged_rank=1), axis=0',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21']],),   # shape=(3, None)
+          axis=0,
+          expected=[[[b'a00', b'a01'], [], [b'a20', b'a21']]]),
+      dict(
+          descr='One rank-2 input (ragged_rank=1), axis=1',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],),   # shape=(3, None)
+          axis=1,
+          expected=[
+              [[b'a00', b'a01']],
+              [[]],
+              [[b'a20', b'a21', b'a22']]]),
+      dict(
+          descr='One rank-2 input (ragged_rank=1), axis=2',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],),   # shape=(3, None)
+          axis=2,
+          expected=[
+              [[b'a00'], [b'a01']], [],
+              [[b'a20'], [b'a21'], [b'a22']]]),
+      dict(
+          descr='One rank-2 input (ragged_rank=1), axis=-3',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21']],),   # shape=(3, None)
+          axis=-3,
+          expected=[[[b'a00', b'a01'], [], [b'a20', b'a21']]]),
+      dict(
+          descr='One rank-2 input (ragged_rank=1), axis=-2',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],),   # shape=(3, None)
+          axis=-2,
+          expected=[
+              [[b'a00', b'a01']],
+              [[]],
+              [[b'a20', b'a21', b'a22']]]),
+      dict(
+          descr='One rank-2 input (ragged_rank=1), axis=-1',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],),  # shape=(3, None)
+          axis=-1,
+          expected=[
+              [[b'a00'], [b'a01']], [],
+              [[b'a20'], [b'a21'], [b'a22']]]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=0',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21']],   # shape=(3, None)
+              [['b00'], ['b10']]),                    # shape=(2, None)
+          axis=0,
+          expected=[[[b'a00', b'a01'], [], [b'a20', b'a21']], [[b'b00'],
+                                                               [b'b10']]]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=1',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00'], ['b10', 'b11', 'b12'], ['b20']]),    # shape=(3, None)
+          axis=1,
+          expected=[
+              [[b'a00', b'a01'], [b'b00']],
+              [[], [b'b10', b'b11', b'b12']],
+              [[b'a20', b'a21', b'a22'], [b'b20']]]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=2',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00', 'b01'], [], ['b20', 'b21', 'b22']]),  # shape=(3, None)
+          axis=2,
+          expected=[
+              [[b'a00', b'b00'], [b'a01', b'b01']], [],
+              [[b'a20', b'b20'], [b'a21', b'b21'], [b'a22', b'b22']]]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=-3',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21']],   # shape=(3, None)
+              [['b00'], ['b10']]),                    # shape=(2, None)
+          axis=-3,
+          expected=[[[b'a00', b'a01'], [], [b'a20', b'a21']], [[b'b00'],
+                                                               [b'b10']]]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=-2',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00'], ['b10', 'b11', 'b12'], ['b20']]),    # shape=(3, None)
+          axis=-2,
+          expected=[
+              [[b'a00', b'a01'], [b'b00']],
+              [[], [b'b10', b'b11', b'b12']],
+              [[b'a20', b'a21', b'a22'], [b'b20']]]),
+      dict(
+          descr='Two rank-2 inputs (ragged_rank=1), axis=-1',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00', 'b01'], [], ['b20', 'b21', 'b22']]),  # shape=(3, None)
+          axis=-1,
+          expected=[
+              [[b'a00', b'b00'], [b'a01', b'b01']], [],
+              [[b'a20', b'b20'], [b'a21', b'b21'], [b'a22', b'b22']]]),
+      dict(
+          descr='Three rank-2 inputs (ragged_rank=1), axis=0',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00'], ['b10']],                            # shape=(2, None)
+              [['c00'], ['c10', 'c11'], ['c21']]),           # shape=(3, None)
+          axis=0,
+          expected=[[[b'a00', b'a01'], [], [b'a20', b'a21', b'a22']],
+                    [[b'b00'], [b'b10']],
+                    [[b'c00'], [b'c10', b'c11'], [b'c21']]]),
+      dict(
+          descr='Three rank-2 inputs (ragged_rank=1), axis=1',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00'], ['b10', 'b11', 'b12'], ['b20']],     # shape=(3, None)
+              [[], ['c10', 'c11'], ['c20', 'c21']]),         # shape=(3, None)
+          axis=1,
+          expected=[
+              [[b'a00', b'a01'], [b'b00'], []],
+              [[], [b'b10', b'b11', b'b12'], [b'c10', b'c11']],
+              [[b'a20', b'a21', b'a22'], [b'b20'], [b'c20', b'c21']]],
+          expected_shape=[3, None, None]),
+      dict(
+          descr='Three rank-2 inputs (ragged_rank=1), axis=2',
+          rt_inputs=(
+              [['a00', 'a01'], [], ['a20', 'a21', 'a22']],   # shape=(3, None)
+              [['b00', 'b01'], [], ['b20', 'b21', 'b22']],   # shape=(3, None)
+              [['c00', 'c01'], [], ['c20', 'c21', 'c22']]),  # shape=(3, None)
+          axis=2,
+          expected=[
+              [[b'a00', b'b00', b'c00'], [b'a01', b'b01', b'c01']], [],
+              [[b'a20', b'b20', b'c20'], [b'a21', b'b21', b'c21'],
+               [b'a22', b'b22', b'c22']]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=0',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']],
+               [['a100', 'a101', 'a102'], ['a110', 'a111']]],
+              [[['b000']], [['b100', 'b101'], ['b110']]],
+              [[], [['c100', 'c101', 'c102', 'c103']], [[], ['c210', 'c211']]]),
+          axis=0,
+          expected=[
+              [[[b'a000', b'a001'], [b'a010']],
+               [[b'a100', b'a101', b'a102'], [b'a110', b'a111']]],
+              [[[b'b000']],
+               [[b'b100', b'b101'], [b'b110']]],
+              [[],
+               [[b'c100', b'c101', b'c102', b'c103']],
+               [[], [b'c210', b'c211']]]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=1',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']],
+               [['a100', 'a101', 'a102'], ['a110', 'a111']]],
+              [[['b000']], [['b100', 'b101'], ['b110']]],
+              [[], [[], ['c110', 'c111']]]),
+          axis=1,
+          expected=[
+              [[[b'a000', b'a001'], [b'a010']], [[b'b000']], []],
+              [[[b'a100', b'a101', b'a102'], [b'a110', b'a111']],
+               [[b'b100', b'b101'], [b'b110']],
+               [[], [b'c110', b'c111']]]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=2',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']],
+               [['a100', 'a101', 'a102'], ['a110', 'a111']]],
+              [[[], ['b010', 'b011']], [['b100', 'b101'], ['b110']]],
+              [[['c000'], ['c010']], [[], ['c110', 'c111']]]),
+          axis=2,
+          expected=[
+              [[[b'a000', b'a001'], [], [b'c000']],
+               [[b'a010'], [b'b010', b'b011'], [b'c010']]],
+              [[[b'a100', b'a101', b'a102'], [b'b100', b'b101'], []],
+               [[b'a110', b'a111'], [b'b110'], [b'c110', b'c111']]]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=3',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']]],
+              [[['b000', 'b001'], ['b010']]],
+              [[['c000', 'c001'], ['c010']]]),
+          axis=3,
+          expected=[[
+              [[b'a000', b'b000', b'c000'], [b'a001', b'b001', b'c001']],
+              [[b'a010', b'b010', b'c010']]]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=-2',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']],
+               [['a100', 'a101', 'a102'], ['a110', 'a111']]],
+              [[[], ['b010', 'b011']], [['b100', 'b101'], ['b110']]],
+              [[['c000'], ['c010']], [[], ['c110', 'c111']]]),
+          axis=-2,
+          expected=[
+              [[[b'a000', b'a001'], [], [b'c000']],
+               [[b'a010'], [b'b010', b'b011'], [b'c010']]],
+              [[[b'a100', b'a101', b'a102'], [b'b100', b'b101'], []],
+               [[b'a110', b'a111'], [b'b110'], [b'c110', b'c111']]]]),
+      dict(
+          descr='Three rank-3 inputs (ragged_rank=2), axis=-1',
+          rt_inputs=(
+              [[['a000', 'a001'], ['a010']]],
+              [[['b000', 'b001'], ['b010']]],
+              [[['c000', 'c001'], ['c010']]]),
+          axis=-1,
+          expected=[[
+              [[b'a000', b'b000', b'c000'], [b'a001', b'b001', b'c001']],
+              [[b'a010', b'b010', b'c010']]]]),
+      dict(
+          descr='ragged_stack([uniform, ragged, uniform], axis=1)',
+          ragged_ranks=[0, 1, 0],
+          rt_inputs=(
+              [['0('], ['1('], ['2(']],                   # shape=(3, 1)
+              [['b00'], ['b10', 'b11', 'b12'], ['b20']],  # shape=(3, None)
+              [[')0'], [')1'], [')2']]),                  # shape=(3, 1)
+          axis=1,
+          expected=[
+              [[b'0('], [b'b00'], [b')0']],
+              [[b'1('], [b'b10', b'b11', b'b12'], [b')1']],
+              [[b'2('], [b'b20'], [b')2']]]),
+      dict(
+          descr='ragged_stack([uniform, uniform], axis=0)',
+          ragged_ranks=[0, 0],
+          rt_inputs=(
+              [['a00', 'a01'], ['a10', 'a11'], ['a20', 'a21']],  # shape=(3, 2)
+              [['b00', 'b01', 'b02'], ['b10', 'b11', 'b12']]),   # shape=(2, 3)
+          axis=0,
+          expected=[
+              [[b'a00', b'a01'], [b'a10', b'a11'], [b'a20', b'a21']],
+              [[b'b00', b'b01', b'b02'], [b'b10', b'b11', b'b12']]]),
+      dict(
+          descr='ragged_stack([1D, 1D], axis=0)',
+          ragged_ranks=[0, 0],
+          rt_inputs=(['a', 'b'], ['c', 'd', 'e']),
+          axis=0,
+          expected=[[b'a', b'b'], [b'c', b'd', b'e']]),
+      dict(
+          descr='ragged_stack([uniform, ragged], axis=0)',
+          ragged_ranks=[0, 1],
+          rt_inputs=(
+              [['a00', 'a01'], ['a10', 'a11'], ['a20', 'a21']],  # shape=(3, 2)
+              [['b00', 'b01', 'b02'], ['b10', 'b11', 'b12']]),   # shape=(2, 3)
+          axis=0,
+          expected=[
+              [[b'a00', b'a01'], [b'a10', b'a11'], [b'a20', b'a21']],
+              [[b'b00', b'b01', b'b02'], [b'b10', b'b11', b'b12']]]),
+      dict(
+          descr='ragged_stack([uniform, ragged], axis=0) with rank-3 inputs',
+          ragged_ranks=[0, 2],
+          rt_inputs=(
+              [[[0, 1], [2, 3]], [[4, 5], [6, 7]]],  # shape = (2, 2, 2)
+              [[[8], [8, 8]]]),                      # shape = (2, None, None)
+          axis=0,
+          expected=[[[[0, 1], [2, 3]], [[4, 5], [6, 7]]], [[[8], [8, 8]]]]),
+      dict(
+          descr='Two rank-3 inputs with ragged_rank=1, axis=-1',
+          ragged_ranks=[1, 1],
+          rt_inputs=(
+              [[[0, 1], [2, 3], [4, 5]], [], [[6, 7], [8, 9]]],
+              [[[9, 8], [7, 6], [5, 4]], [], [[3, 2], [1, 0]]]),
+          axis=-1,
+          expected=[
+              [[[0, 9], [1, 8]], [[2, 7], [3, 6]], [[4, 5], [5, 4]]],
+              [],
+              [[[6, 3], [7, 2]], [[8, 1], [9, 0]]]],
+          expected_shape=[3, None, 2, 2]),
+      dict(
+          descr='Two rank-3 inputs with ragged_rank=1, axis=-2',
+          ragged_ranks=[1, 1],
+          rt_inputs=(
+              [[[0, 1], [2, 3], [4, 5]], [], [[6, 7], [8, 9]]],
+              [[[9, 8], [7, 6], [5, 4]], [], [[3, 2], [1, 0]]]),
+          axis=-2,
+          expected=[
+              [[[0, 1], [9, 8]], [[2, 3], [7, 6]], [[4, 5], [5, 4]]], [],
+              [[[6, 7], [3, 2]], [[8, 9], [1, 0]]]]),
+      dict(
+          descr='ragged_stack([vector, vector], axis=0)',
+          ragged_ranks=[0, 0],
+          rt_inputs=([1, 2, 3], [4, 5, 6]),
+          axis=0,
+          expected=[[1, 2, 3], [4, 5, 6]]),
+      dict(
+          descr='One input (so just adds an outer dimension)',
+          rt_inputs=([['a00', 'a01'], [], ['a20', 'a21']],),
+          axis=0,
+          expected=[[[b'a00', b'a01'], [], [b'a20', b'a21']]]),
+      dict(
+          descr='One input (uniform 0D)',
+          rt_inputs=(1,),
+          ragged_ranks=[0],
+          axis=0,
+          expected=[1]),
+      dict(
+          descr='One input (uniform 1D)',
+          rt_inputs=([1, 2],),
+          ragged_ranks=[0],
+          axis=0,
+          expected=[[1, 2]],
+          expected_ragged_rank=1),
+      dict(
+          descr='One input (uniform 2D)',
+          rt_inputs=([[1, 2], [3, 4], [5, 6]],),
+          ragged_ranks=[0],
+          axis=0,
+          expected=[[[1, 2], [3, 4], [5, 6]]],
+          expected_ragged_rank=2),
+  )   # pyformat: disable
+  def testRaggedStack(self,
+                      descr,
+                      rt_inputs,
+                      axis,
+                      expected,
+                      ragged_ranks=None,
+                      expected_ragged_rank=None,
+                      expected_shape=None):
+    if ragged_ranks is None:
+      ragged_ranks = [None] * len(rt_inputs)
+    rt_inputs = [
+        ragged_factory_ops.constant(rt_input, ragged_rank=rrank)  # pylint: disable=g-long-ternary
+        if rrank != 0 else constant_op.constant(rt_input)
+        for (rt_input, rrank) in zip(rt_inputs, ragged_ranks)
+    ]
+    stacked = ragged_concat_ops.stack(rt_inputs, axis)
+    if expected_ragged_rank is not None:
+      self.assertEqual(stacked.ragged_rank, expected_ragged_rank)
+    if expected_shape is not None:
+      self.assertEqual(stacked.shape.as_list(), expected_shape)
+    self.assertAllEqual(stacked, expected)
+
+  @parameterized.parameters(
+      dict(
+          rt_inputs=(),
+          axis=0,
+          error=ValueError,
+          message=r'rt_inputs may not be empty\.'),
+      dict(
+          rt_inputs=([[1, 2]], [[3, 4]]),
+          axis=r'foo',
+          error=TypeError,
+          message='axis must be an int'),
+      dict(
+          rt_inputs=([[1, 2]], [[3, 4]]),
+          axis=-4,
+          error=ValueError,
+          message='axis=-4 out of bounds: expected -3<=axis<3'),
+      dict(
+          rt_inputs=([[1, 2]], [[3, 4]]),
+          axis=3,
+          error=ValueError,
+          message='axis=3 out of bounds: expected -3<=axis<3'),
+  )
+  def testError(self, rt_inputs, axis, error, message):
+    self.assertRaisesRegex(error, message, ragged_concat_ops.stack, rt_inputs,
+                           axis)
+
+  def testSingleTensorInput(self):
+    """Tests ragged_stack with a single tensor input.
+
+    Usually, we pass a list of values in for rt_inputs.  However, you can
+    also pass in a single value (as with tf.stack), in which case it is
+    equivalent to expand_dims(axis=0).  This test exercises that path.
+    """
+    rt_inputs = ragged_factory_ops.constant([[1, 2], [3, 4]])
+    stacked = ragged_concat_ops.stack(rt_inputs, 0)
+    self.assertAllEqual(stacked, [[[1, 2], [3, 4]]])
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_string_ops.py

@@ -0,0 +1,950 @@
+# Copyright 2015 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.
+# ==============================================================================
+"""Ragged operations for working with string Tensors."""
+
+import typing
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import cond
+from tensorflow.python.ops import gen_string_ops
+from tensorflow.python.ops import map_fn as map_fn_lib
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_math_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.util import compat as util_compat
+from tensorflow.python.util import deprecation
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export("strings.bytes_split")
+@dispatch.add_dispatch_support
+def string_bytes_split(input, name=None):  # pylint: disable=redefined-builtin
+  """Split string elements of `input` into bytes.
+
+  Examples:
+
+  >>> tf.strings.bytes_split('hello').numpy()
+  array([b'h', b'e', b'l', b'l', b'o'], dtype=object)
+  >>> tf.strings.bytes_split(['hello', '123'])
+  <tf.RaggedTensor [[b'h', b'e', b'l', b'l', b'o'], [b'1', b'2', b'3']]>
+
+  Note that this op splits strings into bytes, not unicode characters.  To
+  split strings into unicode characters, use `tf.strings.unicode_split`.
+
+  See also: `tf.io.decode_raw`, `tf.strings.split`, `tf.strings.unicode_split`.
+
+  Args:
+    input: A string `Tensor` or `RaggedTensor`: the strings to split.  Must
+      have a statically known rank (`N`).
+    name: A name for the operation (optional).
+
+  Returns:
+    A `RaggedTensor` of rank `N+1`: the bytes that make up the source strings.
+  """
+  with ops.name_scope(name, "StringsByteSplit", [input]):
+    input = ragged_tensor.convert_to_tensor_or_ragged_tensor(input,
+                                                             name="input")
+    if isinstance(input, ragged_tensor.RaggedTensor):
+      return input.with_flat_values(string_bytes_split(input.flat_values))
+
+    rank = input.shape.ndims
+    if rank is None:
+      raise ValueError("input must have a statically-known rank.")
+
+    if rank == 0:
+      return string_bytes_split(array_ops_stack.stack([input]))[0]
+    elif rank == 1:
+      indices, values, shape = gen_string_ops.string_split(
+          input, delimiter="", skip_empty=False)
+      return ragged_tensor.RaggedTensor.from_value_rowids(
+          values=values, value_rowids=indices[:, 0], nrows=shape[0],
+          validate=False)
+    else:
+      return string_bytes_split(ragged_tensor.RaggedTensor.from_tensor(input))
+
+
+# pylint: disable=redefined-builtin
+@tf_export("strings.unicode_encode")
+@dispatch.add_dispatch_support
+def unicode_encode(input,
+                   output_encoding,
+                   errors="replace",
+                   replacement_char=65533,
+                   name=None):
+  r"""Encodes each sequence of Unicode code points in `input` into a string.
+
+  `result[i1...iN]` is the string formed by concatenating the Unicode
+  codepoints `input[1...iN, :]`, encoded using `output_encoding`.
+
+  Args:
+    input: An `N+1` dimensional potentially ragged integer tensor with shape
+      `[D1...DN, num_chars]`.
+    output_encoding: Unicode encoding that should be used to encode each
+      code point sequence.  Can be `"UTF-8"`, `"UTF-16-BE"`, or `"UTF-32-BE"`.
+    errors: Specifies the response when a code point that is not a [Unicode
+      scalar value](https://www.unicode.org/glossary/#unicode_scalar_value)
+      is encountered (optional). One of:
+            * `'replace'`: Replace non-scalar-value with the
+              `replacement_char`. (default)
+            * `'ignore'`: Skip non-scalar-value .
+            * `'strict'`: Raise an exception for any non-scalar-value.
+    replacement_char: The replacement character to be used in place of
+      any invalid input when `errors='replace'`. Any Unicode scalar value may
+      be used. The default value is the default Unicode replacement character
+      which is 0xFFFD (U+65533).
+    name: A name for the operation (optional).
+
+  Returns:
+    A `N` dimensional `string` tensor with shape `[D1...DN]`.
+
+  #### Example:
+
+  >>> input = tf.ragged.constant(
+  ...     [[71, 246, 246, 100, 110, 105, 103, 104, 116], [128522]])
+  >>> print(unicode_encode(input, 'UTF-8'))
+  tf.Tensor([b'G\xc3\xb6\xc3\xb6dnight' b'\xf0\x9f\x98\x8a'],
+            shape=(2,), dtype=string)
+  """
+  with ops.name_scope(name, "UnicodeEncode", [input]):
+    input_tensor = ragged_tensor.convert_to_tensor_or_ragged_tensor(input)
+    if input_tensor.shape.ndims is None:
+      raise ValueError("Rank of input_tensor must be statically known.")
+    if ragged_tensor.is_ragged(input_tensor):
+      if input_tensor.flat_values.shape.ndims > 1:
+        # If the flat_values of our ragged tensor is multi-dimensional, we can
+        # process it separately and our output will have the same nested splits
+        # as our input.
+        return input_tensor.with_flat_values(
+            unicode_encode(input_tensor.flat_values, output_encoding, errors,
+                           replacement_char))
+      elif input_tensor.ragged_rank > 1:
+        # Recursively process the values of the ragged tensor.
+        return input_tensor.with_values(
+            unicode_encode(input_tensor.values, output_encoding, errors,
+                           replacement_char))
+      else:
+        # Our ragged tensor is of the correct shape (rank 1 flat_values tensor
+        # with ragged_rank of 1) so we can process it as normal.
+        return gen_string_ops.unicode_encode(
+            input_values=input_tensor.values,
+            input_splits=input_tensor.row_splits,
+            output_encoding=output_encoding,
+            errors=errors,
+            replacement_char=replacement_char)
+    else:
+      if input_tensor.shape.ndims == 2:
+        # The input tensor is of the correct 2-D shape, it's just not ragged.
+        return unicode_encode(
+            ragged_tensor.RaggedTensor.from_tensor(input_tensor),
+            output_encoding, errors, replacement_char)
+      elif input_tensor.shape.ndims > 2:
+        # We need to initially flatten the input tensor to 2-D, and then can
+        # reshape the output of our processed flattened tensor.
+        flat_input_tensor = array_ops.reshape(
+            input_tensor,
+            array_ops_stack.stack([-1, array_ops.shape(input_tensor)[-1]]))
+        flat_output_tensor = unicode_encode(flat_input_tensor, output_encoding,
+                                            errors, replacement_char)
+        return array_ops.reshape(flat_output_tensor, input_tensor.shape[:-1])
+      elif input_tensor.shape.ndims == 0:
+        raise ValueError("input_tensor's rank must be at least 1.")
+      else:
+        # Our input tensor is rank 1, so we create a ragged tensor with an added
+        # dimension to create the correct input shape & type, and then remove
+        # the additional dimension from the output and return the string scalar.
+        ragged_input_tensor = ragged_tensor.RaggedTensor.from_row_splits(
+            input_tensor,
+            array_ops_stack.stack(
+                [0, array_ops.shape(input_tensor, out_type=dtypes.int32)[0]]),
+            validate=False)
+        output_tensor = unicode_encode(ragged_input_tensor, output_encoding,
+                                       errors, replacement_char)
+        return array_ops.reshape(output_tensor, [])
+
+
+# pylint: disable=redefined-builtin
+@tf_export("strings.unicode_decode")
+@dispatch.add_dispatch_support
+def unicode_decode(input,
+                   input_encoding,
+                   errors="replace",
+                   replacement_char=0xFFFD,
+                   replace_control_characters=False,
+                   name=None):
+  r"""Decodes each string in `input` into a sequence of Unicode code points.
+
+  `result[i1...iN, j]` is the Unicode codepoint for the `j`th character in
+  `input[i1...iN]`, when decoded using `input_encoding`.
+
+  Args:
+    input: An `N` dimensional potentially ragged `string` tensor with shape
+      `[D1...DN]`.  `N` must be statically known.
+    input_encoding: String name for the unicode encoding that should be used to
+      decode each string.
+    errors: Specifies the response when an input string can't be converted
+      using the indicated encoding. One of:
+      * `'strict'`: Raise an exception for any illegal substrings.
+      * `'replace'`: Replace illegal substrings with `replacement_char`.
+      * `'ignore'`: Skip illegal substrings.
+    replacement_char: The replacement codepoint to be used in place of invalid
+      substrings in `input` when `errors='replace'`; and in place of C0 control
+      characters in `input` when `replace_control_characters=True`.
+    replace_control_characters: Whether to replace the C0 control characters
+      `(U+0000 - U+001F)` with the `replacement_char`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `N+1` dimensional `int32` tensor with shape `[D1...DN, (num_chars)]`.
+    The returned tensor is a `tf.Tensor` if `input` is a scalar, or a
+    `tf.RaggedTensor` otherwise.
+
+  #### Example:
+
+  >>> input = [s.encode('utf8') for s in (u'G\xf6\xf6dnight', u'\U0001f60a')]
+  >>> tf.strings.unicode_decode(input, 'UTF-8').to_list()
+  [[71, 246, 246, 100, 110, 105, 103, 104, 116], [128522]]
+  """
+  with ops.name_scope(name, "UnicodeDecode", [input]):
+    return _unicode_decode(input, input_encoding, errors, replacement_char,
+                           replace_control_characters, with_offsets=False)
+
+
+@tf_export("strings.unicode_decode_with_offsets")
+@dispatch.add_dispatch_support
+def unicode_decode_with_offsets(input,
+                                input_encoding,
+                                errors="replace",
+                                replacement_char=0xFFFD,
+                                replace_control_characters=False,
+                                name=None):
+  r"""Decodes each string into a sequence of code points with start offsets.
+
+  This op is similar to `tf.strings.decode(...)`, but it also returns the
+  start offset for each character in its respective string.  This information
+  can be used to align the characters with the original byte sequence.
+
+  Returns a tuple `(codepoints, start_offsets)` where:
+
+  * `codepoints[i1...iN, j]` is the Unicode codepoint for the `j`th character
+    in `input[i1...iN]`, when decoded using `input_encoding`.
+  * `start_offsets[i1...iN, j]` is the start byte offset for the `j`th
+    character in `input[i1...iN]`, when decoded using `input_encoding`.
+
+  Args:
+    input: An `N` dimensional potentially ragged `string` tensor with shape
+      `[D1...DN]`.  `N` must be statically known.
+    input_encoding: String name for the unicode encoding that should be used to
+      decode each string.
+    errors: Specifies the response when an input string can't be converted
+      using the indicated encoding. One of:
+      * `'strict'`: Raise an exception for any illegal substrings.
+      * `'replace'`: Replace illegal substrings with `replacement_char`.
+      * `'ignore'`: Skip illegal substrings.
+    replacement_char: The replacement codepoint to be used in place of invalid
+      substrings in `input` when `errors='replace'`; and in place of C0 control
+      characters in `input` when `replace_control_characters=True`.
+    replace_control_characters: Whether to replace the C0 control characters
+      `(U+0000 - U+001F)` with the `replacement_char`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A tuple of `N+1` dimensional tensors `(codepoints, start_offsets)`.
+
+    * `codepoints` is an `int32` tensor with shape `[D1...DN, (num_chars)]`.
+    * `offsets` is an `int64` tensor with shape `[D1...DN, (num_chars)]`.
+
+    The returned tensors are `tf.Tensor`s if `input` is a scalar, or
+    `tf.RaggedTensor`s otherwise.
+
+  #### Example:
+
+  >>> input = [s.encode('utf8') for s in (u'G\xf6\xf6dnight', u'\U0001f60a')]
+  >>> result = tf.strings.unicode_decode_with_offsets(input, 'UTF-8')
+  >>> result[0].to_list()  # codepoints
+  [[71, 246, 246, 100, 110, 105, 103, 104, 116], [128522]]
+  >>> result[1].to_list()  # offsets
+  [[0, 1, 3, 5, 6, 7, 8, 9, 10], [0]]
+
+  """
+  with ops.name_scope(name, "UnicodeDecodeWithOffsets", [input]):
+    return _unicode_decode(input, input_encoding, errors, replacement_char,
+                           replace_control_characters, with_offsets=True)
+
+
+@tf_export("strings.unicode_split")
+@dispatch.add_dispatch_support
+def unicode_split(input,
+                  input_encoding,
+                  errors="replace",
+                  replacement_char=0xFFFD,
+                  name=None):
+  r"""Splits each string in `input` into a sequence of Unicode code points.
+
+  `result[i1...iN, j]` is the substring of `input[i1...iN]` that encodes its
+  `j`th character, when decoded using `input_encoding`.
+
+  Args:
+    input: An `N` dimensional potentially ragged `string` tensor with shape
+      `[D1...DN]`.  `N` must be statically known.
+    input_encoding: String name for the unicode encoding that should be used to
+      decode each string.
+    errors: Specifies the response when an input string can't be converted
+      using the indicated encoding. One of:
+      * `'strict'`: Raise an exception for any illegal substrings.
+      * `'replace'`: Replace illegal substrings with `replacement_char`.
+      * `'ignore'`: Skip illegal substrings.
+    replacement_char: The replacement codepoint to be used in place of invalid
+      substrings in `input` when `errors='replace'`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A `N+1` dimensional `int32` tensor with shape `[D1...DN, (num_chars)]`.
+    The returned tensor is a `tf.Tensor` if `input` is a scalar, or a
+    `tf.RaggedTensor` otherwise.
+
+  #### Example:
+
+  >>> input = [s.encode('utf8') for s in (u'G\xf6\xf6dnight', u'\U0001f60a')]
+  >>> tf.strings.unicode_split(input, 'UTF-8').to_list()
+  [[b'G', b'\xc3\xb6', b'\xc3\xb6', b'd', b'n', b'i', b'g', b'h', b't'],
+   [b'\xf0\x9f\x98\x8a']]
+  """
+  with ops.name_scope(name, "UnicodeSplit", [input]):
+    codepoints = _unicode_decode(input, input_encoding, errors,
+                                 replacement_char, False, with_offsets=False)
+    return unicode_encode(
+        ragged_array_ops.expand_dims(codepoints, -1),
+        output_encoding=input_encoding,
+        errors=errors,
+        replacement_char=replacement_char)
+
+
+@tf_export("strings.unicode_split_with_offsets")
+@dispatch.add_dispatch_support
+def unicode_split_with_offsets(input,
+                               input_encoding,
+                               errors="replace",
+                               replacement_char=0xFFFD,
+                               name=None):
+  r"""Splits each string into a sequence of code points with start offsets.
+
+  This op is similar to `tf.strings.decode(...)`, but it also returns the
+  start offset for each character in its respective string.  This information
+  can be used to align the characters with the original byte sequence.
+
+  Returns a tuple `(chars, start_offsets)` where:
+
+  * `chars[i1...iN, j]` is the substring of `input[i1...iN]` that encodes its
+    `j`th character, when decoded using `input_encoding`.
+  * `start_offsets[i1...iN, j]` is the start byte offset for the `j`th
+    character in `input[i1...iN]`, when decoded using `input_encoding`.
+
+  Args:
+    input: An `N` dimensional potentially ragged `string` tensor with shape
+      `[D1...DN]`.  `N` must be statically known.
+    input_encoding: String name for the unicode encoding that should be used to
+      decode each string.
+    errors: Specifies the response when an input string can't be converted
+      using the indicated encoding. One of:
+      * `'strict'`: Raise an exception for any illegal substrings.
+      * `'replace'`: Replace illegal substrings with `replacement_char`.
+      * `'ignore'`: Skip illegal substrings.
+    replacement_char: The replacement codepoint to be used in place of invalid
+      substrings in `input` when `errors='replace'`.
+    name: A name for the operation (optional).
+
+  Returns:
+    A tuple of `N+1` dimensional tensors `(codepoints, start_offsets)`.
+
+    * `codepoints` is an `int32` tensor with shape `[D1...DN, (num_chars)]`.
+    * `offsets` is an `int64` tensor with shape `[D1...DN, (num_chars)]`.
+
+    The returned tensors are `tf.Tensor`s if `input` is a scalar, or
+    `tf.RaggedTensor`s otherwise.
+
+  #### Example:
+
+  >>> input = [s.encode('utf8') for s in (u'G\xf6\xf6dnight', u'\U0001f60a')]
+  >>> result = tf.strings.unicode_split_with_offsets(input, 'UTF-8')
+  >>> result[0].to_list()  # character substrings
+  [[b'G', b'\xc3\xb6', b'\xc3\xb6', b'd', b'n', b'i', b'g', b'h', b't'],
+   [b'\xf0\x9f\x98\x8a']]
+  >>> result[1].to_list()  # offsets
+  [[0, 1, 3, 5, 6, 7, 8, 9, 10], [0]]
+
+  """
+  with ops.name_scope(name, "UnicodeSplitWithOffsets", [input]):
+    codepoints, offsets = _unicode_decode(input, input_encoding, errors,
+                                          replacement_char, False,
+                                          with_offsets=True)
+    chars = unicode_encode(
+        ragged_array_ops.expand_dims(codepoints, -1),
+        output_encoding=input_encoding,
+        errors=errors,
+        replacement_char=replacement_char)
+    return chars, offsets
+
+
+def _unicode_decode(input, input_encoding, errors, replacement_char,
+                    replace_control_characters, with_offsets):
+  """Decodes each string into a sequence of codepoints."""
+  input = ragged_tensor.convert_to_tensor_or_ragged_tensor(input, name="input")
+  input_ndims = input.shape.ndims
+  if input_ndims is None:
+    raise ValueError("Rank of `input` must be statically known.")
+
+  if input_ndims > 1:
+    # Convert to a ragged tensor with ragged_rank = input_ndims - 1.
+    if not ragged_tensor.is_ragged(input):
+      input = ragged_tensor.RaggedTensor.from_tensor(
+          input, ragged_rank=input_ndims - 1)
+    elif input.ragged_rank < input_ndims - 1:
+      input = input.with_flat_values(
+          ragged_tensor.RaggedTensor.from_tensor(
+              input.flat_values,
+              ragged_rank=input_ndims - input.ragged_rank - 1))
+
+  # Reshape the input to a flat vector, and apply the gen_string_ops op.
+  if ragged_tensor.is_ragged(input):
+    flat_input = array_ops.reshape(input.flat_values, [-1])
+  else:
+    flat_input = array_ops.reshape(input, [-1])
+
+  if with_offsets:
+    decode_op = gen_string_ops.unicode_decode_with_offsets
+  else:
+    decode_op = gen_string_ops.unicode_decode
+  flat_result = decode_op(
+      input=flat_input,
+      input_encoding=input_encoding,
+      errors=errors,
+      replacement_char=replacement_char,
+      replace_control_characters=replace_control_characters)
+
+  if input_ndims == 0:
+    codepoints = flat_result.char_values
+    if with_offsets:
+      offsets = flat_result.char_to_byte_starts
+  else:
+    codepoints = ragged_tensor.RaggedTensor.from_row_splits(
+        flat_result.char_values, flat_result.row_splits, validate=False)
+    if input_ndims > 1:
+      codepoints = input.with_flat_values(codepoints)
+    if with_offsets:
+      offsets = ragged_tensor.RaggedTensor.from_row_splits(
+          flat_result.char_to_byte_starts, flat_result.row_splits,
+          validate=False)
+      if input_ndims > 1:
+        offsets = input.with_flat_values(offsets)
+
+  if with_offsets:
+    return codepoints, offsets
+  else:
+    return codepoints
+
+
+@tf_export("strings.split", v1=[])
+@dispatch.add_dispatch_support
+def string_split_v2(input, sep=None, maxsplit=-1, name=None):  # pylint: disable=redefined-builtin
+  """Split elements of `input` based on `sep` into a `RaggedTensor`.
+
+  Let N be the size of `input` (typically N will be the batch size). Split each
+  element of `input` based on `sep` and return a `RaggedTensor` containing the
+  split tokens. Empty tokens are ignored.
+
+  Example:
+
+  >>> tf.strings.split('hello world').numpy()
+   array([b'hello', b'world'], dtype=object)
+  >>> tf.strings.split(['hello world', 'a b c'])
+  <tf.RaggedTensor [[b'hello', b'world'], [b'a', b'b', b'c']]>
+
+  If `sep` is given, consecutive delimiters are not grouped together and are
+  deemed to delimit empty strings. For example, `input` of `"1<>2<><>3"` and
+  `sep` of `"<>"` returns `["1", "2", "", "3"]`. If `sep` is None or an empty
+  string, consecutive whitespace are regarded as a single separator, and the
+  result will contain no empty strings at the start or end if the string has
+  leading or trailing whitespace.
+
+  Note that the above mentioned behavior matches python's str.split.
+
+  Args:
+    input: A string `Tensor` of rank `N`, the strings to split.  If
+      `rank(input)` is not known statically, then it is assumed to be `1`.
+    sep: `0-D` string `Tensor`, the delimiter string.
+    maxsplit: An `int`. If `maxsplit > 0`, limit of the split of the result.
+    name: A name for the operation (optional).
+
+  Raises:
+    ValueError: If sep is not a string.
+
+  Returns:
+    A `RaggedTensor` of rank `N+1`, the strings split according to the
+    delimiter.
+  """
+  with ops.name_scope(name, "StringSplit", [input]):
+    input = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        input, dtype=dtypes.string, name="input")
+    if isinstance(input, ragged_tensor.RaggedTensor):
+      return input.with_flat_values(
+          string_split_v2(input.flat_values, sep, maxsplit))
+
+    rank = input.shape.ndims
+    if rank == 0:
+      return string_split_v2(array_ops_stack.stack([input]), sep, maxsplit)[0]
+    elif rank == 1 or rank is None:
+      sparse_result = string_ops.string_split_v2(
+          input, sep=sep, maxsplit=maxsplit)
+      return ragged_tensor.RaggedTensor.from_value_rowids(
+          values=sparse_result.values,
+          value_rowids=sparse_result.indices[:, 0],
+          nrows=sparse_result.dense_shape[0],
+          validate=False)
+    else:
+      return string_split_v2(
+          ragged_tensor.RaggedTensor.from_tensor(input), sep, maxsplit)
+
+
+@tf_export(v1=["string_split"])
+@dispatch.add_dispatch_support
+@deprecation.deprecated_args(None,
+                             "delimiter is deprecated, please use sep instead.",
+                             "delimiter")
+def string_split(source, sep=None, skip_empty=True, delimiter=None,
+                 result_type="SparseTensor", name=None):  # pylint: disable=invalid-name
+  """Split elements of `source` based on `delimiter`.
+
+  Let N be the size of `source` (typically N will be the batch size). Split each
+  element of `source` based on `delimiter` and return a `SparseTensor`
+  or `RaggedTensor` containing the split tokens. Empty tokens are ignored.
+
+  If `sep` is an empty string, each element of the `source` is split
+  into individual strings, each containing one byte. (This includes splitting
+  multibyte sequences of UTF-8.) If delimiter contains multiple bytes, it is
+  treated as a set of delimiters with each considered a potential split point.
+
+  Examples:
+
+  >>> print(tf.compat.v1.string_split(['hello world', 'a b c']))
+  SparseTensor(indices=tf.Tensor( [[0 0] [0 1] [1 0] [1 1] [1 2]], ...),
+               values=tf.Tensor([b'hello' b'world' b'a' b'b' b'c'], ...),
+               dense_shape=tf.Tensor([2 3], shape=(2,), dtype=int64))
+
+  >>> print(tf.compat.v1.string_split(['hello world', 'a b c'],
+  ...     result_type="RaggedTensor"))
+  <tf.RaggedTensor [[b'hello', b'world'], [b'a', b'b', b'c']]>
+
+  Args:
+    source: `1-D` string `Tensor`, the strings to split.
+    sep: `0-D` string `Tensor`, the delimiter character, the string should
+      be length 0 or 1. Default is ' '.
+    skip_empty: A `bool`. If `True`, skip the empty strings from the result.
+    delimiter: deprecated alias for `sep`.
+    result_type: The tensor type for the result: one of `"RaggedTensor"` or
+      `"SparseTensor"`.
+    name: A name for the operation (optional).
+
+  Raises:
+    ValueError: If delimiter is not a string.
+
+  Returns:
+    A `SparseTensor` or `RaggedTensor` of rank `2`, the strings split according
+    to the delimiter.  The first column of the indices corresponds to the row
+    in `source` and the second column corresponds to the index of the split
+    component in this row.
+  """
+  with ops.name_scope(name, "StringSplit", [source]):
+    sparse_result = string_ops.string_split(
+        source, sep=sep, skip_empty=skip_empty, delimiter=delimiter)
+    if result_type == "SparseTensor":
+      return sparse_result
+    elif result_type == "RaggedTensor":
+      return ragged_tensor.RaggedTensor.from_value_rowids(
+          values=sparse_result.values,
+          value_rowids=sparse_result.indices[:, 0],
+          nrows=sparse_result.dense_shape[0],
+          validate=False)
+    else:
+      raise ValueError("result_type must be 'RaggedTensor' or 'SparseTensor'.")
+
+
+# In TensorFlow 1.x, "tf.strings.split" uses the new signature (with maxsplit),
+# but we need to add the result_type argument.
+@tf_export(v1=["strings.split"])
+@dispatch.add_dispatch_support
+def strings_split_v1(input=None, sep=None, maxsplit=-1,  # pylint: disable=redefined-builtin
+                     result_type="SparseTensor", source=None, name=None):
+  """Split elements of `input` based on `sep`.
+
+  Let N be the size of `input` (typically N will be the batch size). Split each
+  element of `input` based on `sep` and return a `SparseTensor` or
+  `RaggedTensor` containing the split tokens. Empty tokens are ignored.
+
+  Examples:
+
+  >>> print(tf.compat.v1.strings.split(['hello world', 'a b c']))
+  SparseTensor(indices=tf.Tensor( [[0 0] [0 1] [1 0] [1 1] [1 2]], ...),
+               values=tf.Tensor([b'hello' b'world' b'a' b'b' b'c'], ...),
+               dense_shape=tf.Tensor([2 3], shape=(2,), dtype=int64))
+
+  >>> print(tf.compat.v1.strings.split(['hello world', 'a b c'],
+  ...     result_type="RaggedTensor"))
+  <tf.RaggedTensor [[b'hello', b'world'], [b'a', b'b', b'c']]>
+
+  If `sep` is given, consecutive delimiters are not grouped together and are
+  deemed to delimit empty strings. For example, `input` of `"1<>2<><>3"` and
+  `sep` of `"<>"` returns `["1", "2", "", "3"]`. If `sep` is None or an empty
+  string, consecutive whitespace are regarded as a single separator, and the
+  result will contain no empty strings at the start or end if the string has
+  leading or trailing whitespace.
+
+  Note that the above mentioned behavior matches python's str.split.
+
+  Args:
+    input: A string `Tensor` of rank `N`, the strings to split.  If
+      `rank(input)` is not known statically, then it is assumed to be `1`.
+    sep: `0-D` string `Tensor`, the delimiter character.
+    maxsplit: An `int`. If `maxsplit > 0`, limit of the split of the result.
+    result_type: The tensor type for the result: one of `"RaggedTensor"` or
+      `"SparseTensor"`.
+    source: alias for "input" argument.
+    name: A name for the operation (optional).
+
+  Raises:
+    ValueError: If sep is not a string.
+
+  Returns:
+    A `SparseTensor` or `RaggedTensor` of rank `N+1`, the strings split
+    according to the delimiter.
+  """
+  input = deprecation.deprecated_argument_lookup(
+      "input", input, "source", source)
+  with ops.name_scope(name, "StringSplit", [input]):
+    input = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        input, dtype=dtypes.string, name="input")
+
+    if input.shape.rank == 0:
+      input = array_ops.expand_dims(input, 0)
+
+    if result_type == "SparseTensor":
+      if input.shape.rank == 1:
+        return string_ops.string_split_v2(input, sep=sep, maxsplit=maxsplit)
+      else:
+        return string_split_v2(input, sep=sep, maxsplit=maxsplit).to_sparse()
+    elif result_type == "RaggedTensor":
+      return string_split_v2(input, sep=sep, maxsplit=maxsplit)
+    else:
+      raise ValueError("result_type must be 'RaggedTensor' or 'SparseTensor'.")
+
+
+@dispatch.dispatch_for_api(string_ops.reduce_join_v2)
+def reduce_join(inputs: ragged_tensor.Ragged,
+                axis=None,
+                keepdims=None,
+                separator="",
+                name=None):
+  """For docs, see: _RAGGED_REDUCE_DOCSTRING."""
+  return ragged_math_ops.ragged_reduce_aggregate(
+      string_ops.reduce_join, string_ops.unsorted_segment_join, inputs, axis,
+      keepdims, separator, name or "RaggedSegmentJoin")
+
+
+@tf_export("strings.ngrams")
+@dispatch.add_dispatch_support
+def ngrams(data,
+           ngram_width,
+           separator=" ",
+           pad_values=None,
+           padding_width=None,
+           preserve_short_sequences=False,
+           name=None):
+  """Create a tensor of n-grams based on `data`.
+
+  Creates a tensor of n-grams based on `data`. The n-grams are created by
+  joining windows of `width` adjacent strings from the inner axis of `data`
+  using `separator`.
+
+  The input data can be padded on both the start and end of the sequence, if
+  desired, using the `pad_values` argument. If set, `pad_values` should contain
+  either a tuple of strings or a single string; the 0th element of the tuple
+  will be used to pad the left side of the sequence and the 1st element of the
+  tuple will be used to pad the right side of the sequence. The `padding_width`
+  arg controls how many padding values are added to each side; it defaults to
+  `ngram_width-1`.
+
+  If this op is configured to not have padding, or if it is configured to add
+  padding with `padding_width` set to less than ngram_width-1, it is possible
+  that a sequence, or a sequence plus padding, is smaller than the ngram
+  width. In that case, no ngrams will be generated for that sequence. This can
+  be prevented by setting `preserve_short_sequences`, which will cause the op
+  to always generate at least one ngram per non-empty sequence.
+
+  Examples:
+
+  >>> tf.strings.ngrams(["A", "B", "C", "D"], 2).numpy()
+  array([b'A B', b'B C', b'C D'], dtype=object)
+  >>> tf.strings.ngrams(["TF", "and", "keras"], 1).numpy()
+  array([b'TF', b'and', b'keras'], dtype=object)
+
+  Args:
+    data: A Tensor or RaggedTensor containing the source data for the ngrams.
+    ngram_width: The width(s) of the ngrams to create. If this is a list or
+      tuple, the op will return ngrams of all specified arities in list order.
+      Values must be non-Tensor integers greater than 0.
+    separator: The separator string used between ngram elements. Must be a
+      string constant, not a Tensor.
+    pad_values: A tuple of (left_pad_value, right_pad_value), a single string,
+      or None. If None, no padding will be added; if a single string, then that
+      string will be used for both left and right padding. Values must be Python
+      strings. Should be set when `padding_width` is not None.
+    padding_width: If set, `padding_width` pad values will be added to both
+      sides of each sequence. Defaults to `ngram_width`-1. Must be greater than
+      0. (Note that 1-grams are never padded, regardless of this value.)
+    preserve_short_sequences: If true, then ensure that at least one ngram is
+      generated for each input sequence.  In particular, if an input sequence is
+      shorter than `min(ngram_width) + 2*pad_width`, then generate a single
+      ngram containing the entire sequence.  If false, then no ngrams are
+      generated for these short input sequences.
+    name: The op name.
+
+  Returns:
+    A RaggedTensor of ngrams. If `data.shape=[D1...DN, S]`, then
+    `output.shape=[D1...DN, NUM_NGRAMS]`, where
+    `NUM_NGRAMS=S-ngram_width+1+2*padding_width`.
+
+  Raises:
+    TypeError: if `pad_values` is set to an invalid type.
+    ValueError: if `pad_values`, `padding_width`, or `ngram_width` is set to an
+      invalid value.
+    ValueError: if `padding_width` is not None and `pad_values` is None.
+  """
+
+  with ops.name_scope(name, "StringNGrams", [data]):
+    if pad_values is None:
+      left_pad = ""
+      right_pad = ""
+    elif isinstance(pad_values, (list, tuple)):
+      if (not isinstance(pad_values[0], util_compat.bytes_or_text_types) or
+          not isinstance(pad_values[1], util_compat.bytes_or_text_types)):
+        raise TypeError(
+            "pad_values must be a string, tuple of strings, or None.")
+      left_pad = pad_values[0]
+      right_pad = pad_values[1]
+    else:
+      if not isinstance(pad_values, util_compat.bytes_or_text_types):
+        raise TypeError(
+            "pad_values must be a string, tuple of strings, or None.")
+      left_pad = pad_values
+      right_pad = pad_values
+
+    if padding_width is not None and padding_width < 1:
+      raise ValueError("padding_width must be greater than 0.")
+
+    if padding_width is not None and pad_values is None:
+      raise ValueError("pad_values must be provided if padding_width is set.")
+
+    data = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        data, name="data", dtype=dtypes.string)
+
+    # preserve the shape of the data if it is a tensor
+    to_tensor = False
+    if isinstance(data, tensor_lib.Tensor):
+      dense_shape = array_ops.concat([array_ops.shape(data)[:-1], [-1]], axis=0)
+      to_tensor = True
+
+    if not isinstance(data, ragged_tensor.RaggedTensor):
+      if data.shape.ndims is None:
+        raise ValueError("Rank of data must be known.")
+      elif data.shape.ndims == 0:
+        raise ValueError("Data must have rank>0")
+      elif data.shape.ndims == 1:
+        rt = ragged_tensor.RaggedTensor.from_row_starts(
+            data, [0], validate=False)
+        return ngrams(rt, ngram_width, separator, pad_values, padding_width,
+                      preserve_short_sequences, name)[0]
+      else:
+        data = ragged_tensor.RaggedTensor.from_tensor(
+            data, ragged_rank=data.shape.ndims - 1)
+
+    if data.ragged_rank > 1:
+      output = data.with_values(
+          ngrams(data.values, ngram_width, separator, pad_values, padding_width,
+                 preserve_short_sequences, name))
+      return array_ops.reshape(output.flat_values,
+                               dense_shape) if to_tensor else output
+
+    if pad_values is None:
+      padding_width = 0
+
+    if pad_values is not None and padding_width is None:
+      padding_width = -1
+
+    if not isinstance(ngram_width, (list, tuple)):
+      ngram_widths = [ngram_width]
+    else:
+      ngram_widths = ngram_width
+    for width in ngram_widths:
+      if width < 1:
+        raise ValueError("All ngram_widths must be greater than 0. Got %s" %
+                         ngram_width)
+
+    output, output_splits = gen_string_ops.string_n_grams(
+        data=data.flat_values,
+        data_splits=data.row_splits,
+        separator=separator,
+        ngram_widths=ngram_widths,
+        left_pad=left_pad,
+        right_pad=right_pad,
+        pad_width=padding_width,
+        preserve_short_sequences=preserve_short_sequences)
+
+    # if the input is Dense tensor, the output should also be a dense tensor
+    output = ragged_tensor.RaggedTensor.from_row_splits(
+        values=output, row_splits=output_splits, validate=False)
+    return array_ops.reshape(output.flat_values,
+                             dense_shape) if to_tensor else output
+
+
+@dispatch.dispatch_for_api(string_ops.string_format)
+def string_format(
+    template: str,
+    inputs: typing.Union[ragged_tensor.Ragged,
+                         typing.List[ragged_tensor.RaggedOrDense]],
+    placeholder="{}",
+    summarize=3,
+    name=None):
+  """Version of tf.strings.format that handles RaggedTensors."""
+  if tensor_util.is_tf_type(inputs) or ragged_tensor.is_ragged(inputs):
+    inputs = [inputs]
+
+  split_template = template.split(placeholder)
+  if len(inputs) != len(split_template) - 1:
+    raise ValueError("num placeholders in template and num inputs must match"
+                     ": {} vs {}".format(len(split_template) - 1, len(inputs)))
+
+  with ops.name_scope(name, "StringFormat", [inputs]):
+    output_pieces = [constant_op.constant(split_template[0])]
+    for i, input in enumerate(inputs):
+      if ragged_tensor.is_ragged(input):
+        output_pieces.append(ragged_tensor_to_string(input, summarize))
+      else:
+        output_pieces.append(string_ops.string_format(
+            "{}", [input], summarize=summarize))
+      output_pieces.append(constant_op.constant(split_template[i + 1]))
+    if len(output_pieces) == 1:
+      return output_pieces[0]
+    else:
+      return string_ops.reduce_join(output_pieces)
+
+
+def ragged_tensor_to_string(rt, summarize=None):
+  """Returns a scalar string tensor with the contents of a RaggedTensor.
+
+  Requires that `rt.shape.rank` is not `None`.
+
+  Note: this converts the entire `RaggedTensor` into a single string scalar.
+  If you want to convert individual elements, use `tf.strings.as_string(rt)`.
+
+  >>> rt1 = tf.ragged.constant([[1, 2, 3], [4, 5]])
+  >>> ragged_tensor_to_string(rt1).numpy()
+  b'[[1, 2, 3], [4, 5]]'
+
+  >>> rt2 = tf.ragged.constant([[['a'], ['b', 'c']], [['d', 'e', 'f'], []]])
+  >>> ragged_tensor_to_string(rt2).numpy()
+  b"[[['a'], ['b', 'c']], [['d', 'e', 'f'], []]]"
+
+  >>> rt3 = tf.ragged.constant([[1], [2, 3, 4, 5, 6], [], [], [7], [8, 9]])
+  >>> ragged_tensor_to_string(rt3, summarize=2).numpy()
+  b'[[1], [2, 3, ..., 5, 6], ..., [7], [8, 9]]'
+
+  Args:
+    rt: The RaggedTensor that should be converted to a string.
+    summarize: If specified, then only the first and last `summarize` elements
+      within each dimension are included in the string. If `-1` or `None`, then
+      all elements are included.
+  """
+  if (summarize is not None and summarize != -1 and
+      not (isinstance(summarize, int) and summarize > 0)):
+    raise ValueError("Expected summarize to be -1 or a positive int, got %r" %
+                     summarize)
+  with ops.name_scope(None, "AsString", [rt]):
+    rt = ragged_tensor.convert_to_tensor_or_ragged_tensor(rt)
+    if rt.shape.rank is None:
+      raise ValueError("RaggedTensor to_string requires that rt.shape.rank "
+                       "is not None.")
+    # Convert all elements of `rt` to strings.
+    if rt.dtype == dtypes.string:
+      escaped = string_ops.regex_replace(rt.flat_values, r"(['\\])", r"\\\1")
+      str_t = rt.with_flat_values("'" + escaped + "'")
+    else:
+      str_t = rt.with_flat_values(string_ops.as_string(rt.flat_values))
+
+    return _ragged_tensor_to_string(str_t, summarize)
+
+
+def _ragged_tensor_to_string(string_tensor, summarize):
+  """Returns a scalar string tensor with the contents of `string_tensor`.
+
+  Args:
+    string_tensor: A potentially ragged tensor with dtype=string.
+    summarize: Include only the first and last `summarize` elements of each
+      dimension.  If `-1` or `None`, then include all elements.
+
+  Returns:
+    A scalar string Tensor.
+  """
+  if string_tensor.shape.rank == 1:
+    pieces = string_tensor
+  else:
+    pieces = map_fn_lib.map_fn(
+        lambda s: _ragged_tensor_to_string(s, summarize),
+        string_tensor,
+        fn_output_signature=tensor_lib.TensorSpec(None, dtypes.string))
+  if summarize not in (-1, None):
+    pieces = cond.cond(
+        _nrows(string_tensor) <= 2 * summarize,
+        lambda: pieces,
+        lambda: array_ops.concat(  # pylint: disable=g-long-lambda
+            [pieces[:summarize], ["..."], pieces[-summarize:]],
+            axis=0))
+  return "[" + string_ops.reduce_join(pieces, separator=", ") + "]"
+
+
+def _nrows(tensor, out_type=dtypes.int32):
+  if isinstance(tensor, ragged_tensor.RaggedTensor):
+    return tensor.nrows(out_type=out_type)
+  else:
+    return array_ops.shape(tensor, out_type=out_type)[0]
+
+
+@dispatch.dispatch_for_api(string_ops.string_join)
+def string_join(inputs: typing.List[ragged_tensor.RaggedOrDense],
+                separator="",
+                name=None):
+  """RaggedTensor implementation for tf.strings.join."""
+  if len(inputs) < 0:
+    raise ValueError("tf.strings.join: expected at least one input.")
+  with ops.name_scope(name, "RaggedStringJoin", inputs):
+    return ragged_functional_ops.map_flat_values(string_ops.string_join, inputs,
+                                                 separator)

tensorflow/python/ops/ragged/ragged_tensor_bounding_shape_op_test.py

@@ -0,0 +1,139 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged.bounding_shape."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorBoundingShapeOp(test_util.TensorFlowTestCase,
+                                  parameterized.TestCase):
+
+  @parameterized.named_parameters([
+      # rank = 2
+      dict(testcase_name='docstring_example',
+           rt=[[1, 2, 3, 4], [5], [], [6, 7, 8, 9], [10]],
+           expected=[5, 4]),
+      dict(testcase_name='shape_5_3',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           expected=[5, 3]),
+      dict(testcase_name='shape_1_7',
+           rt=[['a', 'b', 'c', 'd', 'e', 'f', 'g']],
+           expected=[1, 7]),
+      dict(testcase_name='shape_3_7',
+           rt=[[], ['a', 'b', 'c', 'd', 'e', 'f', 'g'], []],
+           expected=[3, 7]),
+      dict(testcase_name='shape_5_3_row_splits_int32',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           rt_row_splits_dtype=dtypes.int32,
+           expected=[5, 3]),
+      dict(testcase_name='shape_0_0',
+           rt=[],
+           rt_ragged_rank=1,
+           expected=[0, 0]),
+      dict(testcase_name='shape_3_0',
+           rt=[[], [], []],
+           expected=[3, 0]),
+      # rank = 3
+      dict(testcase_name='shape_5_3_2',
+           rt=[[[0, 1], [2]], [[3, 4], [], [5, 6]], [[7]], [], [[8, 9]]],
+           expected=[5, 3, 2]),
+      dict(testcase_name='shape_1_7_2',
+           rt=[[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13]]],
+           expected=[1, 7, 2]),
+      dict(testcase_name='shape_3_7_4',
+           rt=[[], [[0, 1], [2], [], [3], [4], [5, 6, 7, 8], [9]], []],
+           expected=[3, 7, 4]),
+      dict(testcase_name='shape_1_7_2_ragged_rank_1',
+           rt=[[[0, 1], [2, 3], [4, 5], [6, 7], [8, 9], [10, 11], [12, 13]]],
+           rt_ragged_rank=1,
+           expected=[1, 7, 2]),
+      # axis != None
+      dict(testcase_name='shape_5_3_axis_0',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           axis=0,
+           expected=5),
+      dict(testcase_name='shape_5_3_axis_1',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           axis=1,
+           expected=3),
+      dict(testcase_name='shape_5_3_axis_1_0',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           axis=[1, 0],
+           expected=[3, 5]),
+      # out_type != None
+      dict(testcase_name='shape_5_3_row_splits_int64_out_type_int64',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           rt_row_splits_dtype=dtypes.int64,
+           out_type=dtypes.int64,
+           expected=[5, 3]),
+      dict(testcase_name='shape_5_3_row_splits_int32_out_type_int32',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           rt_row_splits_dtype=dtypes.int32,
+           out_type=dtypes.int32,
+           expected=[5, 3]),
+      dict(testcase_name='shape_5_3_row_splits_int64_out_type_int32',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           rt_row_splits_dtype=dtypes.int64,
+           out_type=dtypes.int32,
+           expected=[5, 3]),
+      dict(testcase_name='shape_5_3_row_splits_int32_out_type_int64',
+           rt=[['a', 'b'], ['c', 'd', 'e'], ['f'], [], ['g']],
+           rt_row_splits_dtype=dtypes.int32,
+           out_type=dtypes.int64,
+           expected=[5, 3]),
+      dict(testcase_name='shape_1_3_axis_1_row_splits_int64_out_type_int32',
+           rt=[[1, 2, 3]],
+           rt_row_splits_dtype=dtypes.int64,
+           axis=1,
+           out_type=dtypes.int32,
+           expected=3)
+  ])  # pyformat: disable
+  def testBoundingShape(self,
+                        rt,
+                        expected,
+                        axis=None,
+                        out_type=None,
+                        rt_row_splits_dtype=dtypes.int64,
+                        rt_ragged_rank=None):
+    rt = ragged_factory_ops.constant(
+        rt, ragged_rank=rt_ragged_rank, row_splits_dtype=rt_row_splits_dtype)
+    bounding_shape = rt.bounding_shape(axis=axis, out_type=out_type)
+    self.assertAllEqual(bounding_shape, expected)
+    if out_type is not None:
+      self.assertEqual(bounding_shape.dtype, out_type)
+    else:
+      self.assertEqual(bounding_shape.dtype, rt_row_splits_dtype)
+
+    # If we're testing a configuration that uses `axis`, then make sure
+    # that it also works if `axis` is a tensor.
+    if axis is not None:
+      bounding_shape = rt.bounding_shape(
+          axis=constant_op.constant(axis), out_type=out_type)
+      self.assertAllEqual(bounding_shape, expected)
+      if out_type is not None:
+        self.assertEqual(bounding_shape.dtype, out_type)
+      else:
+        self.assertEqual(bounding_shape.dtype, rt_row_splits_dtype)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_tensor_shape.py

@@ -0,0 +1,628 @@
+# 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.
+# ==============================================================================
+"""Shapes & broadcasting for RaggedTensors."""
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import array_ops_stack
+from tensorflow.python.ops import control_flow_assert
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_config
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_util
+
+
+class RaggedTensorDynamicShape:
+  """A collection of tensors encoding the shape of a potentially ragged tensor.
+
+  Each `RaggedTensorDynamicShape` consists of an ordered list of dimension
+  sizes.  There are two dimension types:
+
+    * "Uniform dimensions" are dimensions where all slices have the same
+      length.  `RaggedTensorDynamicShape` records the size of each uniform
+      dimension using a single scalar integer.
+
+    * "Ragged dimensions" are dimensions whose slices may have different
+      lengths.  `RaggedTensorDynamicShape` records the size of each ragged
+      dimension using an integer vector containing the slice lengths for all
+      the slices across that dimension.
+
+  Furthermore, there are two ways a dimension might be encoded:
+
+    * "Partitioned dimensions" are dimensions that are encoded using a
+      `RaggedTensor`'s `nested_row_splits`.  The outermostmost partitioned
+      dimension must be uniform, and the innermost partitioned dimension must
+      be ragged.
+
+    * "Inner dimensions" are dimensions that are encoded using a
+      `RaggedTensor`'s `flat_values`.  Inner dimensions are always uniform.
+
+  The sizes of partitioned dimensions are recorded using `partitioned_dim_sizes`
+  and `inner_dim_sizes`:
+
+    * `partitioned_dim_sizes` is a list of tensors (one for each partitioned
+      dimension).
+
+      * For uniform dimensions, the tensor is an integer scalar specifying the
+        size of all slices across that dimension.
+      * For ragged dimensions, the tensor is an integer vector specifying the
+        size of each slice across that dimension.
+
+    * `inner_dim_sizes` is a single integer vector, where each element
+      specifies the size of a single inner dimension.
+
+  Examples:
+
+  Tensor                         | Ragged | Partitioned Dim Sizes  | Inner Dim
+                                 : Rank   :                        : Sizes
+  ------------------------------ | ------ | ---------------------- | ----------
+  `[[1, 2, 3], [4, 5, 6]]`       |      0 |                        | `2, 3`
+  `[[1, 2], [], [3, 4, 5]]`      |      1 | `3, (2, 0, 3)`         |
+  `[[[1, 2], [3, 4]], [[5, 6]]]` |      1 | `2, (2, 1)`            | 2
+  `[[[1, 2], [3]], [[4, 5]]]`    |      2 | `2, (2, 1), (2, 1, 2)` |
+  """
+
+  def __init__(self, partitioned_dim_sizes, inner_dim_sizes,
+               dim_size_dtype=None):
+    """Creates a RaggedTensorDynamicShape.
+
+    Args:
+      partitioned_dim_sizes: A `list` of 0-D or 1-D integer `Tensor`, one for
+        each partitioned dimension.  If dimension `d` is uniform, then
+        `partitioned_dim_sizes[d]` must be an integer scalar, specifying the
+        size of all slices across dimension `d`.  If dimension `d` is ragged,
+        then `partitioned_dim_sizes[d]` must be an integer vector, specifying
+        the size of each slice across dimension `d`.
+      inner_dim_sizes: A 1-D integer `Tensor`, whose length is equal to the
+        number of inner dimensions.  `inner_dim_sizes[n]` is the size of all
+        slices across the `n`th inner dimension (which is the
+        `(len(partitioned_dim_sizes)+n)`th dimension in the overall tensor.
+      dim_size_dtype: dtype for dimension sizes.  If not specified, then it
+        is chosen based on the dtypes of `partitioned_dim_sizes` and
+        `inner_dim_sizes`.
+    """
+    assert isinstance(partitioned_dim_sizes, (list, tuple))
+
+    with ops.name_scope(None, 'RaggedTensorDynamicShape',
+                        (partitioned_dim_sizes, inner_dim_sizes)):
+      partitioned_dim_sizes = tuple(
+          ops.convert_to_tensor(size, name='partitioned_dimension_size_%d' % i)
+          for (i, size) in enumerate(partitioned_dim_sizes))
+      inner_dim_sizes = ops.convert_to_tensor(
+          inner_dim_sizes, name='inner_dim_sizes')
+
+      # Validate shapes.
+      if partitioned_dim_sizes:
+        for axis, dimension_size in enumerate(partitioned_dim_sizes):
+          if dimension_size.shape.ndims is None:
+            raise ValueError(
+                'rank of partitioned_dim_sizes[%d] is unknown' % axis)
+          dimension_size.shape.with_rank_at_most(1)
+        if partitioned_dim_sizes[0].shape.ndims == 1:
+          raise ValueError('outermost partitioned dimension must be uniform')
+        if partitioned_dim_sizes[-1].shape.ndims == 0:
+          raise ValueError('innermost partitioned dimension must be ragged')
+      inner_dim_sizes.shape.assert_has_rank(1)
+
+      # Convert dimension size tensors to a single dtype.
+      if dim_size_dtype is None:
+        dim_size_dtypes = set(
+            p.dtype for p in partitioned_dim_sizes if p.shape.ndims == 1)
+        if not dim_size_dtypes:
+          dim_size_dtype = dtypes.int64
+        elif len(dim_size_dtypes) == 1:
+          dim_size_dtype = dim_size_dtypes.pop()
+        else:
+          if not ragged_config.auto_cast_partition_dtype():
+            raise ValueError('partitioned_dim_sizes must have matching dtypes')
+          dim_size_dtype = dtypes.int64
+      partitioned_dim_sizes = tuple(math_ops.cast(p, dim_size_dtype)
+                                    for p in partitioned_dim_sizes)
+      inner_dim_sizes = math_ops.cast(inner_dim_sizes, dim_size_dtype)
+
+      self._partitioned_dim_sizes = partitioned_dim_sizes
+      self._inner_dim_sizes = inner_dim_sizes
+
+  def __repr__(self):
+    return ('RaggedTensorDynamicShape'
+            '(partitioned_dim_sizes=%r, inner_dim_sizes=%r)' %
+            (self._partitioned_dim_sizes, self._inner_dim_sizes))
+
+  @staticmethod
+  def from_dim_sizes(dim_sizes):
+    """Constructs a ragged shape from a list of dimension sizes.
+
+    This list contains a single tensor for each dimension, where the tensor
+    is a scalar if the dimension is uniform, or a vector if the dimension is
+    ragged.
+
+    Args:
+      dim_sizes: List of int32 or int64 scalars or vectors.
+
+    Returns:
+      A RaggedTensorDynamicShape.
+    """
+    with ops.name_scope(None, 'RaggedTensorDynamicShapeFromDimensionSizes',
+                        [dim_sizes]):
+      dim_sizes = tuple(
+          ops.convert_to_tensor(size, preferred_dtype=dtypes.int64,
+                                name='dim_sizes') for size in dim_sizes)
+      # Split the dimensions into partitioned & inner dimensions.
+      inner_split = 0
+      for dim, dim_size in enumerate(dim_sizes):
+        if dim_size.shape.ndims == 1:
+          inner_split = dim + 1
+        elif dim_size.shape.ndims != 0:
+          raise ValueError('Each dim_size must be a scalar or a vector')
+      return RaggedTensorDynamicShape(dim_sizes[:inner_split],
+                                      dim_sizes[inner_split:])
+
+  @classmethod
+  def from_tensor(cls, rt_input, dim_size_dtype=None):
+    """Constructs a ragged shape for a potentially ragged tensor."""
+    with ops.name_scope(None, 'RaggedTensorDynamicShapeFromTensor', [rt_input]):
+      rt_input = ragged_tensor.convert_to_tensor_or_ragged_tensor(rt_input)
+      if not ragged_tensor.is_ragged(rt_input):
+        return cls([], array_ops.shape(rt_input), dim_size_dtype=dim_size_dtype)
+      else:
+        partitioned_dim_sizes = (
+            (rt_input.nrows(),) + rt_input.nested_row_lengths())
+        return RaggedTensorDynamicShape(
+            partitioned_dim_sizes,
+            array_ops.shape(rt_input.flat_values)[1:],
+            dim_size_dtype=dim_size_dtype)
+
+  def dimension_size(self, axis):
+    """Returns the size of slices across the specified dimension."""
+    if not isinstance(axis, int):
+      raise TypeError('axis must be an integer')
+    partitioned_ndims = len(self._partitioned_dim_sizes)
+    if axis < partitioned_ndims:
+      return self._partitioned_dim_sizes[axis]
+    else:
+      return self._inner_dim_sizes[axis - partitioned_ndims]
+
+  def is_ragged(self, axis):
+    """Returns true if the indicated dimension is ragged."""
+    if not isinstance(axis, int):
+      raise TypeError('axis must be an integer')
+    rank = self.rank
+    if axis < 0:
+      raise ValueError('Negative axis values are not supported')
+    elif rank is not None and axis >= rank:
+      raise ValueError('Expected axis=%s < rank=%s' % (axis, rank))
+    else:
+      return (axis > 0 and axis < len(self._partitioned_dim_sizes) and
+              self._partitioned_dim_sizes[axis].shape.ndims == 1)
+
+  @property
+  def rank(self):
+    """The number of dimensions in this shape, or None if unknown."""
+    inner_ndims = tensor_shape.dimension_value(self._inner_dim_sizes.shape[0])
+    if inner_ndims is None:
+      return None
+    else:
+      return len(self._partitioned_dim_sizes) + inner_ndims
+
+  @property
+  def partitioned_dim_sizes(self):
+    """The partitioned dimension sizes for this shape.
+
+    Returns:
+      A `list` of 0-D or 1-D integer `Tensor`.
+    """
+    return self._partitioned_dim_sizes
+
+  @property
+  def inner_dim_sizes(self):
+    """The inner dimension sizes for this shape.
+
+    Returns:
+      A 1-D integer `Tensor`.
+    """
+    return self._inner_dim_sizes
+
+  @property
+  def num_partitioned_dimensions(self):
+    """The number of partitioned dimensions in this shape."""
+    return len(self._partitioned_dim_sizes)
+
+  @property
+  def num_inner_dimensions(self):
+    """The number of inner dimensions, or `None` if not statically known."""
+    return tensor_shape.dimension_value(self._inner_dim_sizes.shape[0])
+
+  @property
+  def dim_size_dtype(self):
+    """DType used by this shape for dimension sizes."""
+    return self._inner_dim_sizes.dtype
+
+  def broadcast_to_rank(self, rank):
+    """Adds leading size-1 dimensions to broadcast `self` to the given rank.
+
+    E.g., if `shape1` is `[3, (D2), 4]`, then `shape1.broadcast_to_rank(5)`
+    is `[1, 1, 3, (D2), 4]`.
+
+    Args:
+      rank: The rank for the returned shape.
+
+    Returns:
+      A RaggedTensorDynamicShape with `rank` dimensions, whose inner dimensions
+      have the same size as `self` and whose outer dimensions have size `1`.
+
+    Raises:
+      ValueError: If `self.rank` is unknown or greater than `rank`.
+    """
+    if self.rank is None:
+      raise ValueError('Unable to broadcast: self.rank is unknown')
+    dims_to_add = rank - self.rank
+    if dims_to_add < 0:
+      raise ValueError('Unable to broadcast: rank=%d must be greater than '
+                       'self.rank=%d.' % (rank, self.rank))
+    elif dims_to_add == 0:
+      return self
+    elif self._partitioned_dim_sizes:
+      partitioned_dims = (1,) * dims_to_add + self._partitioned_dim_sizes
+      return RaggedTensorDynamicShape(partitioned_dims, self.inner_dim_sizes,
+                                      self.dim_size_dtype)
+    else:
+      inner_dims = array_ops.concat(
+          [array_ops.ones([dims_to_add], self.dim_size_dtype),
+           self.inner_dim_sizes],
+          axis=0)
+      return RaggedTensorDynamicShape([], inner_dims, self.dim_size_dtype)
+
+  def broadcast_dimension(self, axis, lengths):
+    """Returns a shape that is broadcast-compatible with self & lengths.
+
+    * If dimension[axis] is uniform and lengths is a scalar, the check
+      that either lengths==1 or axis==1 or lengths==axis, and tile
+      dimension[axis] with tf.where(lengths==axis, 1, axis) repeats.
+
+    * If dimension[axis] is uniform and lengths is a vector, then check
+      that dimension[axis]==1, and raggedly tile dimension[axis] with
+      lengths repeats.  (we can skip tiling if we statically know that
+      slice_lengths == 1??)
+
+    * If dimension[axis] is ragged and lengths is a scalar, then check
+      that lengths==1.
+
+    * If dimension[axis] is ragged and lengths is a vector, then check
+      that self.dimension_size(axis) == lengths.
+
+    Args:
+      axis: `int`.  The dimension to broadcast.
+      lengths: 0-D or 1-D integer `Tensor`.
+
+    Returns:
+      A `RaggedTensorDynamicShape`.
+    """
+    lengths = ragged_util.convert_to_int_tensor(
+        lengths, name='lengths', dtype=self.dim_size_dtype)
+    # Check whether lengths is a scalar (for uniform dimensions) or
+    # vector (for ragged dimensions).
+    if lengths.shape.ndims is None:
+      raise ValueError('lengths must have a known rank.')
+    elif lengths.shape.ndims > 1:
+      raise ValueError('lengths must be a scalar or vector')
+    else:
+      lengths_is_scalar = (lengths.shape.ndims == 0)
+
+    # Verify that the shapes are compatible.
+    if self.is_ragged(axis):
+      if lengths_is_scalar:
+        condition = math_ops.equal(lengths, 1)
+      else:
+        condition = math_ops.reduce_all(
+            math_ops.equal(lengths, self.dimension_size(axis)))
+    else:
+      axis_dim_size = self.dimension_size(axis)
+      if lengths_is_scalar:
+        condition = (
+            math_ops.equal(lengths, 1) | math_ops.equal(axis_dim_size, 1)
+            | math_ops.equal(axis_dim_size, lengths))
+      else:
+        condition = math_ops.equal(axis_dim_size, 1)
+    broadcast_err = [
+        'Unable to broadcast: dimension size mismatch in dimension', axis,
+        'lengths=', lengths, 'dim_size=',
+        self.dimension_size(axis)
+    ]
+    broadcast_check = control_flow_assert.Assert(
+        condition, data=broadcast_err, summarize=10)
+
+    with ops.control_dependencies([broadcast_check]):
+      # Partitioned dimensions:
+      if axis < self.num_partitioned_dimensions:
+        if self.is_ragged(axis):
+          # Use an identity op to make sure the check actually gets run.
+          return RaggedTensorDynamicShape(
+              self._partitioned_dim_sizes,
+              array_ops.identity(self.inner_dim_sizes), self.dim_size_dtype)
+        else:
+          return self._broadcast_uniform_partitioned_dimension(axis, lengths)
+
+      # Inner dimensions:
+      else:
+        if lengths_is_scalar:
+          return self._broadcast_inner_dimension_to_uniform(axis, lengths)
+        else:
+          if axis == 0:
+            raise ValueError('Unable to broadcast: '
+                             'outermost dimension must be uniform.')
+          return self._broadcast_inner_dimension_to_ragged(axis, lengths)
+
+  def num_slices_in_dimension(self, axis):
+    """Returns the total number of slices across the indicated dimension."""
+    if axis < 0:
+      return constant_op.constant(1, dtype=self.dim_size_dtype)
+    elif self.is_ragged(axis):
+      return math_ops.reduce_sum(self._partitioned_dim_sizes[axis])
+    else:
+      return self.dimension_size(axis) * self.num_slices_in_dimension(axis - 1)
+
+  def _broadcast_uniform_partitioned_dimension(self, axis, lengths):
+    """Broadcasts the partitioned dimension `axis` to match `lengths`."""
+    axis_dim_size = self.dimension_size(axis)
+    partitioned_sizes = list(self._partitioned_dim_sizes[:axis])
+
+    if lengths.shape.ndims == 0:
+      lengths = array_ops.where(
+          math_ops.equal(axis_dim_size, 1), lengths, axis_dim_size)
+      repeats = array_ops.where(math_ops.equal(axis_dim_size, 1), lengths, 1)
+      splits = array_ops_stack.stack([0, self.num_slices_in_dimension(axis)])
+    else:
+      splits = math_ops.range(
+          array_ops.size(lengths, out_type=self.dim_size_dtype) + 1)
+      repeats = lengths
+
+    partitioned_sizes.append(lengths)
+
+    for dim_size in self._partitioned_dim_sizes[axis + 1:]:
+      if dim_size.shape.ndims == 0:
+        partitioned_sizes.append(dim_size)
+        splits *= dim_size
+      else:
+        partitioned_sizes.append(
+            ragged_util.repeat_ranges(dim_size, splits, repeats))
+        splits = array_ops.gather(
+            ragged_util.lengths_to_splits(dim_size), splits)
+    inner_sizes = self._inner_dim_sizes
+    return RaggedTensorDynamicShape(partitioned_sizes, inner_sizes,
+                                    self.dim_size_dtype)
+
+  def _broadcast_inner_dimension_to_uniform(self, axis, length):
+    """Broadcasts the inner dimension `axis` to match `lengths`."""
+    dim_size = self.dimension_size(axis)
+    axis_in_inner_dims = axis - self.num_partitioned_dimensions
+    partitioned_sizes = self._partitioned_dim_sizes
+    inner_sizes = array_ops.concat([
+        self._inner_dim_sizes[:axis_in_inner_dims],
+        [array_ops.where(math_ops.equal(dim_size, 1), length, dim_size)],
+        self._inner_dim_sizes[axis_in_inner_dims + 1:]
+    ],
+                                   axis=0)
+    return RaggedTensorDynamicShape(partitioned_sizes, inner_sizes,
+                                    self.dim_size_dtype)
+
+  def _broadcast_inner_dimension_to_ragged(self, axis, lengths):
+    axis_in_inner_dims = axis - self.num_partitioned_dimensions
+    partitioned_sizes = (
+        self._partitioned_dim_sizes + tuple([
+            self._inner_dim_sizes[i] for i in range(axis_in_inner_dims)
+        ]) + (lengths,))
+    inner_sizes = self._inner_dim_sizes[axis_in_inner_dims + 1:]
+    return RaggedTensorDynamicShape(partitioned_sizes, inner_sizes)
+
+  def with_dim_size_dtype(self, dtype):
+    if dtype not in (dtypes.int32, dtypes.int64):
+      raise ValueError('dtype must be int32 or int64')
+    if self.dim_size_dtype == dtype:
+      return self
+    return RaggedTensorDynamicShape(
+        [math_ops.cast(p, dtype) for p in self._partitioned_dim_sizes],
+        math_ops.cast(self._inner_dim_sizes, dtype))
+
+
+def broadcast_dynamic_shape(shape_x, shape_y):
+  """Returns the shape formed by broadcasting two shapes to be compatible.
+
+  Args:
+    shape_x: A `RaggedTensorDynamicShape`
+    shape_y: A `RaggedTensorDynamicShape`
+
+  Returns:
+    A `RaggedTensorDynamicShape`.
+  Raises:
+    ValueError: If `shape_x` and `shape_y` are not broadcast-compatible.
+  """
+  if not isinstance(shape_x, RaggedTensorDynamicShape):
+    raise TypeError('shape_x must be a RaggedTensorDynamicShape')
+  if not isinstance(shape_y, RaggedTensorDynamicShape):
+    raise TypeError('shape_y must be a RaggedTensorDynamicShape')
+
+  # Broadcast both shapes to have the same rank.
+  if shape_x.rank is None or shape_y.rank is None:
+    raise ValueError('Unable to broadcast: unknown rank')
+  broadcast_rank = max(shape_x.rank, shape_y.rank)
+  shape_x = shape_x.broadcast_to_rank(broadcast_rank)
+  shape_y = shape_y.broadcast_to_rank(broadcast_rank)
+
+  # Broadcast dimensions one at a time, starting from the outermost dimension.
+  for axis in range(broadcast_rank):
+    shape_x = shape_x.broadcast_dimension(axis, shape_y.dimension_size(axis))
+    shape_y = shape_y.broadcast_dimension(axis, shape_x.dimension_size(axis))
+
+  return shape_x
+
+
+def broadcast_to(rt_input, shape, broadcast_inner_dimensions=True):
+  """Broadcasts a potentially ragged tensor to a ragged shape.
+
+  Tiles `rt_input` as necessary to match the given shape.
+
+  Behavior is undefined if `rt_input` is not broadcast-compatible with `shape`.
+
+  Args:
+    rt_input: The potentially ragged tensor to broadcast.
+    shape: A `RaggedTensorDynamicShape`
+    broadcast_inner_dimensions: If false, then inner dimensions will not be
+      tiled.
+
+  Returns:
+    A potentially ragged tensor whose values are taken from
+    `rt_input`, and whose shape matches `shape`.
+  """
+  if not isinstance(shape, RaggedTensorDynamicShape):
+    raise TypeError('shape must be a RaggedTensorDynamicShape')
+  rt_input = ragged_tensor.convert_to_tensor_or_ragged_tensor(rt_input)
+
+  # Broadcasting to a uniform shape.
+  if shape.num_partitioned_dimensions == 0:
+    return _broadcast_to_uniform_shape(rt_input, shape,
+                                       broadcast_inner_dimensions)
+  else:
+    return _broadcast_to_ragged_shape(rt_input, shape,
+                                      broadcast_inner_dimensions)
+
+
+def _broadcast_to_uniform_shape(rt_input, shape, broadcast_inner_dimensions):
+  """Broadcasts rt_input to the uniform shape `shape`."""
+  if isinstance(rt_input, ragged_tensor.RaggedTensor):
+    raise ValueError('Incompatible with shape: ragged rank mismatch')
+  if broadcast_inner_dimensions:
+    return array_ops.broadcast_to(rt_input, shape.inner_dim_sizes)
+  else:
+    return rt_input
+
+
+def _broadcast_to_ragged_shape(rt_input, dst_shape, broadcast_inner_dimensions):
+  """Broadcasts rt_input to the ragged shape `dst_shape`."""
+  # Check that rt_input and dst_shape have the same row_splits dtype.
+  if (isinstance(rt_input, ragged_tensor.RaggedTensor) and
+      rt_input.row_splits.dtype != dst_shape.dim_size_dtype):
+    if not ragged_config.auto_cast_partition_dtype():
+      raise ValueError('rt_input and dst_shape have different row_split '
+                       'dtypes; use RaggedTensor.with_row_splits_dtype() or '
+                       'RaggedTensorDynamicShape.with_dim_size_dtype() to '
+                       'convert to a compatible dtype.')
+    rt_input = rt_input.with_row_splits_dtype(dtypes.int64)
+    dst_shape = dst_shape.with_dim_size_dtype(dtypes.int64)
+
+  # dst_shape's rank and ragged_rank must be greater than or equal to rt_input's
+  if rt_input.shape.ndims is None or dst_shape.rank is None:
+    raise ValueError('Unable to broadcast: unknown rank')
+  if rt_input.shape.ndims > dst_shape.rank:
+    raise ValueError('Incompatible with shape: rank mismatch')
+  if (isinstance(rt_input, ragged_tensor.RaggedTensor) and
+      rt_input.ragged_rank >= dst_shape.num_partitioned_dimensions):
+    raise ValueError('Incompatible with shape: ragged rank mismatch')
+
+  src_shape = RaggedTensorDynamicShape.from_tensor(rt_input)
+  src_shape = src_shape.broadcast_to_rank(dst_shape.rank)
+
+  # Add dimensions to rt_input so its rank and ragged_rank matches dst_shape.
+  if dst_shape.rank > rt_input.shape.ndims:
+    if rt_input.shape.ndims < dst_shape.num_inner_dimensions + 1:
+      rt_input = array_ops.reshape(
+          rt_input, array_ops.concat([[-1], dst_shape.inner_dim_sizes], axis=0))
+    for _ in range(dst_shape.rank - rt_input.shape.ndims):
+      if ragged_tensor.is_ragged(rt_input):
+        nrows = rt_input.nrows()
+      else:
+        nrows = array_ops.shape(rt_input,
+                                out_type=dst_shape.dim_size_dtype)[0]
+      rt_input = ragged_tensor.RaggedTensor.from_row_lengths(rt_input, [nrows],
+                                                             validate=False)
+
+  # Add ragged dimensions to match dst_shape.
+  if ragged_tensor.is_ragged(rt_input):
+    inner_rank_diff = (
+        rt_input.flat_values.shape.ndims - 1 - dst_shape.num_inner_dimensions)
+    if inner_rank_diff > 0:
+      rt_input = rt_input.with_flat_values(
+          ragged_tensor.RaggedTensor.from_tensor(
+              rt_input.flat_values, ragged_rank=inner_rank_diff,
+              row_splits_dtype=dst_shape.dim_size_dtype))
+  else:
+    rt_input = ragged_tensor.RaggedTensor.from_tensor(
+        rt_input, ragged_rank=dst_shape.num_partitioned_dimensions - 1,
+        row_splits_dtype=dst_shape.dim_size_dtype)
+
+  # Do broadcasting for any dimensions that will remain uniform.  We can do
+  # these all at once, since they're independent of one another.
+  multiples = [1] * dst_shape.rank
+  for axis in range(dst_shape.num_partitioned_dimensions):
+    if not src_shape.is_ragged(axis) and not dst_shape.is_ragged(axis):
+      src_size = src_shape.dimension_size(axis)
+      dst_size = dst_shape.dimension_size(axis)
+      if ((tensor_util.constant_value(src_size) in (1, None)) and
+          (tensor_util.constant_value(dst_size) != 1)):
+        multiples[axis] = array_ops.where(
+            math_ops.equal(src_size, 1), dst_size, 1)
+  if not all(isinstance(v, int) and v == 1 for v in multiples):
+    multiples = array_ops_stack.stack(multiples, axis=0)
+    rt_input = ragged_array_ops.tile(rt_input, multiples)
+
+  if broadcast_inner_dimensions:
+    new_shape = array_ops.broadcast_dynamic_shape(
+        array_ops.shape(
+            rt_input.flat_values, out_type=dst_shape.dim_size_dtype),
+        array_ops.concat([[1], dst_shape.inner_dim_sizes], axis=0))
+    rt_input = rt_input.with_flat_values(
+        array_ops.broadcast_to(rt_input.flat_values, new_shape))
+
+  # Do broadcasting for dimensions that become ragged.  We must do these from
+  # outermost to innermost.
+  for axis in range(dst_shape.num_partitioned_dimensions):
+    if not src_shape.is_ragged(axis) and dst_shape.is_ragged(axis):
+      dst_size = dst_shape.dimension_size(axis)
+      rt_input = _ragged_tile_axis(rt_input, axis, dst_size,
+                                   dst_shape.dim_size_dtype)
+
+  return rt_input
+
+
+def _ragged_tile_axis(rt_input, axis, repeats, row_splits_dtype):
+  """Tile a dimension of a RaggedTensor to match a ragged shape."""
+  assert axis > 0  # Outermost dimension may not be ragged.
+
+  if not ragged_tensor.is_ragged(rt_input):
+    rt_input = ragged_tensor.RaggedTensor.from_tensor(
+        rt_input, ragged_rank=1, row_splits_dtype=row_splits_dtype)
+
+  if axis > 1:
+    return rt_input.with_values(
+        _ragged_tile_axis(rt_input.values, axis - 1, repeats,
+                          row_splits_dtype))
+  else:
+    src_row_splits = rt_input.nested_row_splits
+    src_row_lengths = rt_input.nested_row_lengths()
+    splits = src_row_splits[0]
+
+    dst_row_lengths = [repeats]
+    for i in range(1, len(src_row_lengths)):
+      dst_row_lengths.append(
+          ragged_util.repeat_ranges(src_row_lengths[i], splits, repeats))
+      splits = array_ops.gather(src_row_splits[i], splits)
+    dst_values = ragged_util.repeat_ranges(rt_input.flat_values, splits,
+                                           repeats)
+    return ragged_tensor.RaggedTensor.from_nested_row_lengths(
+        dst_values, dst_row_lengths, validate=False)

tensorflow/python/ops/ragged/ragged_tensor_shape_test.py

@@ -0,0 +1,483 @@
+# 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.
+# ==============================================================================
+"""Tests for tf.ragged.ragged_tensor_shape."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_shape
+from tensorflow.python.ops.ragged.ragged_tensor_shape import RaggedTensorDynamicShape
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorShapeTest(test_util.TensorFlowTestCase,
+                            parameterized.TestCase):
+
+  def assertShapeEq(self, x, y):
+    assert isinstance(x, RaggedTensorDynamicShape)
+    assert isinstance(y, RaggedTensorDynamicShape)
+    self.assertLen(x.partitioned_dim_sizes, len(y.partitioned_dim_sizes))
+    for x_dims, y_dims in zip(x.partitioned_dim_sizes, y.partitioned_dim_sizes):
+      self.assertAllEqual(x_dims, y_dims)
+    self.assertAllEqual(x.inner_dim_sizes, y.inner_dim_sizes)
+
+  @parameterized.parameters([
+      dict(value='x', expected_dim_sizes=[]),
+      dict(value=['a', 'b', 'c'], expected_dim_sizes=[3]),
+      dict(value=[['a', 'b', 'c'], ['d', 'e', 'f']], expected_dim_sizes=[2, 3]),
+      dict(
+          value=[[['a', 'b', 'c'], ['d', 'e', 'f']]],
+          expected_dim_sizes=[1, 2, 3]),
+      dict(
+          value=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d',
+                                                                     'e']]),
+          expected_dim_sizes=[2, [3, 2]]),
+      dict(
+          value=ragged_factory_ops.constant_value([[['a', 'b', 'c'], ['d',
+                                                                      'e']]]),
+          expected_dim_sizes=[1, [2], [3, 2]]),
+      dict(
+          value=ragged_factory_ops.constant_value(
+              [[['a', 'b', 'c'], ['d', 'e', 'f']]], ragged_rank=1),
+          expected_dim_sizes=[1, [2], 3]),
+      dict(
+          value=ragged_factory_ops.constant_value(
+              [[[[1], [2]], [[3], [4]]], [[[5], [6]]]], ragged_rank=1),
+          expected_dim_sizes=[2, [2, 1], 2, 1]),
+      dict(
+          value=ragged_factory_ops.constant_value([[10, 20], [30]]),
+          expected_dim_sizes=[2, [2, 1]]),
+      # Docstring examples:
+      dict(value=[[1, 2, 3], [4, 5, 6]], expected_dim_sizes=[2, 3]),
+      dict(
+          value=ragged_factory_ops.constant_value([[1, 2], [], [3, 4, 5]]),
+          expected_dim_sizes=[3, [2, 0, 3]]),
+      dict(
+          value=ragged_factory_ops.constant_value([[[1, 2], [3, 4]], [[5, 6]]],
+                                                  ragged_rank=1),
+          expected_dim_sizes=[2, [2, 1], 2]),
+      dict(
+          value=ragged_factory_ops.constant_value([[[1, 2], [3]], [[4, 5]]]),
+          expected_dim_sizes=[2, [2, 1], [2, 1, 2]]),
+  ])
+  def testFromTensor(self, value, expected_dim_sizes):
+    shape = RaggedTensorDynamicShape.from_tensor(value)
+    expected = RaggedTensorDynamicShape.from_dim_sizes(expected_dim_sizes)
+    self.assertShapeEq(shape, expected)
+
+  @parameterized.parameters([
+      dict(dim_sizes=[], rank=0, expected_dim_sizes=[]),
+      dict(dim_sizes=[], rank=3, expected_dim_sizes=[1, 1, 1]),
+      dict(dim_sizes=[3], rank=1, expected_dim_sizes=[3]),
+      dict(dim_sizes=[3], rank=3, expected_dim_sizes=[1, 1, 3]),
+      dict(dim_sizes=[2, 3], rank=3, expected_dim_sizes=[1, 2, 3]),
+      dict(dim_sizes=[3, [3, 2, 4]], rank=2, expected_dim_sizes=[3, [3, 2, 4]]),
+      dict(
+          dim_sizes=[3, [3, 2, 4]],
+          rank=4,
+          expected_dim_sizes=[1, 1, 3, [3, 2, 4]]),
+      dict(
+          dim_sizes=[3, [3, 2, 4], 2, 3],
+          rank=5,
+          expected_dim_sizes=[1, 3, [3, 2, 4], 2, 3]),
+  ])
+  def testBroadcastToRank(self, dim_sizes, rank, expected_dim_sizes):
+    shape = RaggedTensorDynamicShape.from_dim_sizes(dim_sizes)
+    expected = RaggedTensorDynamicShape.from_dim_sizes(expected_dim_sizes)
+    broadcasted_shape = shape.broadcast_to_rank(rank)
+    self.assertShapeEq(broadcasted_shape, expected)
+    self.assertEqual(broadcasted_shape.rank, rank)
+
+  @parameterized.parameters([
+      #=========================================================================
+      # dimension[axis] is uniform inner; and row_lengths is a scalar
+      #=========================================================================
+      # shape: [BROADCAST(UNIFORM), UNIFORM, UNIFORM]
+      dict(axis=0,
+           row_length=3,
+           original_dim_sizes=[1, 4, 5],
+           broadcast_dim_sizes=[3, 4, 5]),
+
+      # shape: [UNIFORM, UNIFORM, BROADCAST(UNIFORM)]
+      dict(axis=2,
+           row_length=5,
+           original_dim_sizes=[3, 4, 1],
+           broadcast_dim_sizes=[3, 4, 5]),
+
+      # shape: [UNIFORM, RAGGED, BROADCAST(UNIFORM)]
+      dict(axis=2,
+           row_length=5,
+           original_dim_sizes=[3, [3, 2, 8], 1],
+           broadcast_dim_sizes=[3, [3, 2, 8], 5]),
+
+      # shape: [UNIFORM, RAGGED, RAGGED, UNIFORM, UNIFORM, BROADCAST(UNIFORM)]
+      dict(axis=5,
+           row_length=5,
+           original_dim_sizes=[2, [2, 1], [3, 2, 8], 3, 4, 1],
+           broadcast_dim_sizes=[2, [2, 1], [3, 2, 8], 3, 4, 5]),
+
+      #=========================================================================
+      # dimension[axis] is uniform inner; and row_lengths is a vector
+      #=========================================================================
+      # shape: [UNIFORM, BROADCAST(UNIFORM)]
+      dict(axis=1,
+           row_length=[2, 0, 1],
+           original_dim_sizes=[3, 1],
+           broadcast_dim_sizes=[3, [2, 0, 1]]),
+      # shape: [UNIFORM, BROADCAST(UNIFORM), UNIFORM]
+      dict(axis=1,
+           row_length=[2, 0, 1],
+           original_dim_sizes=[3, 1, 5],
+           broadcast_dim_sizes=[3, [2, 0, 1], 5]),
+
+      # shape: [UNIFORM, UNIFORM, BROADCAST(UNIFORM)]
+      dict(axis=2,
+           row_length=[2, 0, 1, 3, 8, 2, 3, 4, 1, 8, 7, 0],
+           original_dim_sizes=[4, 3, 1],
+           broadcast_dim_sizes=[4, 3, [2, 0, 1, 3, 8, 2, 3, 4, 1, 8, 7, 0]]),
+
+      # shape: [UNIFORM, RAGGED, BROADCAST(UNIFORM)]
+      dict(axis=2,
+           row_length=[2, 5, 3],
+           original_dim_sizes=[2, [2, 1], 1],
+           broadcast_dim_sizes=[2, [2, 1], [2, 5, 3]]),
+
+      # shape: [UNIFORM, RAGGED, UNIFORM, UNIFORM, BROADCAST(UNIFORM), UNIFORM]
+      dict(axis=4,
+           row_length=list(range(18)),
+           original_dim_sizes=[2, [2, 1], 3, 2, 1, 8],
+           broadcast_dim_sizes=[2, [2, 1], 3, 2, list(range(18)), 8]),
+
+      #=========================================================================
+      # dimension[axis] is uniform partitioned; and row_lengths is a scalar
+      #=========================================================================
+      # shape: [BROADCAST(UNIFORM), RAGGED]
+      dict(axis=0,
+           row_length=3,
+           original_dim_sizes=[1, [5]],
+           broadcast_dim_sizes=[3, [5, 5, 5]]),
+
+      # shape: [BROADCAST(UNIFORM), UNIFORM, RAGGED]
+      dict(axis=0,
+           row_length=2,
+           original_dim_sizes=[1, 3, [3, 0, 2]],
+           broadcast_dim_sizes=[2, 3, [3, 0, 2, 3, 0, 2]]),
+
+      # shape: [BROADCAST(UNIFORM), RAGGED, RAGGED, UNIFORM, UNIFORM]
+      dict(axis=0,
+           row_length=3,
+           original_dim_sizes=[1, [3], [3, 5, 2], 9, 4, 5],
+           broadcast_dim_sizes=[3, [3, 3, 3], [3, 5, 2, 3, 5, 2, 3, 5, 2],
+                                9, 4, 5]),
+
+      # shape: [BROADCAST(UNIFORM), UNIFORM, RAGGED, UNIFORM]
+      dict(axis=0,
+           row_length=2,
+           original_dim_sizes=[1, 2, [2, 1], [3, 5, 2], 2],
+           broadcast_dim_sizes=[2, 2, [2, 1, 2, 1], [3, 5, 2, 3, 5, 2], 2]),
+
+      # shape: [UNIFORM, BROADCAST(UNIFORM), RAGGED, UNIFORM]
+      dict(axis=1,
+           row_length=2,
+           original_dim_sizes=[3, 1, [4, 0, 2], 5],
+           broadcast_dim_sizes=[3, 2, [4, 0, 2, 4, 0, 2], 5]),
+
+      # shape: [UNIFORM, BROADCAST(UNIFORM), RAGGED]
+      dict(axis=1,
+           row_length=1,
+           original_dim_sizes=[2, 3, (1, 2, 3, 4, 5, 6)],
+           broadcast_dim_sizes=[2, 3, (1, 2, 3, 4, 5, 6)]),
+
+      #=========================================================================
+      # dimension[axis] is uniform partitioned; and row_lengths is a vector
+      #=========================================================================
+      # shape: [UNIFORM, BROADCAST(UNIFORM), RAGGED, UNIFORM]
+      dict(axis=1,
+           row_length=[4, 1, 2],
+           original_dim_sizes=[
+               3,                          # axis=0
+               1,                          # axis=1 (broadcast)
+               [3, 1, 2],                  # axis=2
+               5],                         # axis=3
+           broadcast_dim_sizes=[
+               3,                          # axis=0
+               [4, 1, 2],                  # axis=1 (broadcast)
+               [3, 3, 3, 3, 1, 2, 2],      # axis=2
+               5]),                        # axis=3
+
+      # shape: [UNIFORM, BROADCAST(UNIFORM), RAGGED, RAGGED]
+      dict(axis=1,
+           row_length=[2, 0, 3],
+           original_dim_sizes=[
+               3,                                         # axis=0
+               1,                                         # axis=1 (broadcast)
+               [3, 1, 2],                                 # axis=2
+               [3, 1, 4, 1, 5, 9]],                       # axis=3
+           broadcast_dim_sizes=[
+               3,                                         # axis=0
+               [2, 0, 3],                                 # axis=1 (broadcast)
+               [3, 3, 2, 2, 2],                           # axis=2
+               [3, 1, 4, 3, 1, 4, 5, 9, 5, 9, 5, 9]]),    # axis=3
+
+      # shape: [UNIFORM, RAGGED, BROADCAST(UNIFORM), RAGGED, RAGGED, UNIFORM]
+      dict(axis=2,
+           row_length=[4, 1, 2],
+           original_dim_sizes=[
+               3,                                         # axis=0
+               [2, 0, 1],                                 # axis=1
+               1,                                         # axis=2 (broadcast)
+               [3, 2, 1],                                 # axis=3
+               [1, 0, 1, 0, 2, 3],                        # axis=4
+               5],                                        # axis=5
+           broadcast_dim_sizes=[
+               3,                                         # axis=0
+               [2, 0, 1],                                 # axis=2
+               [4, 1, 2],                                 # axis=2 (broadcast)
+               [3, 3, 3, 3, 2, 1, 1],                     # axis=3
+               [1, 0, 1, 1, 0, 1, 1, 0, 1, 1, 0, 1, 0,    # axis=4
+                2, 3, 3],
+               5]),                                       # axis=5
+
+      dict(axis=0,
+           row_length=2,
+           original_dim_sizes=[1, 1, 2, (2, 1)],
+           broadcast_dim_sizes=[2, 1, 2, (2, 1, 2, 1)]),
+      dict(axis=1,
+           row_length=(2, 1),
+           original_dim_sizes=[2, 1, 2, (2, 1, 2, 1)],
+           broadcast_dim_sizes=[2, (2, 1), 2, (2, 1, 2, 1, 2, 1)]),
+      dict(axis=2,
+           row_length=2,
+           original_dim_sizes=[2, (2, 1), 2, (2, 1, 2, 1, 2, 1)],
+           broadcast_dim_sizes=[2, (2, 1), 2, (2, 1, 2, 1, 2, 1)]),
+      dict(axis=3,
+           row_length=(2, 1, 2, 1, 2, 1),
+           original_dim_sizes=[2, (2, 1), 2, 1],
+           broadcast_dim_sizes=[2, (2, 1), 2, (2, 1, 2, 1, 2, 1)]),
+  ])  # pyformat: disable
+  def testBroadcastDimension(self, axis, row_length, original_dim_sizes,
+                             broadcast_dim_sizes):
+    """Tests for the broadcast_dimension method.
+
+    Verifies that:
+
+    * `original.broadcast_dimension(axis, row_length) == broadcast`
+    * `broadcast.broadcast_dimension(axis, row_length) == broadcast`
+    * `broadcast.broadcast_dimension(axis, 1) == broadcast`
+
+    Args:
+      axis: The axis to broadcast
+      row_length: The slice lengths to broadcast to.
+      original_dim_sizes: The dimension sizes before broadcasting.
+        original_dim_sizes[axis] should be equal to `1` or `row_length`.
+      broadcast_dim_sizes: THe dimension sizes after broadcasting.
+    """
+    original_shape = RaggedTensorDynamicShape.from_dim_sizes(original_dim_sizes)
+    bcast_shape = RaggedTensorDynamicShape.from_dim_sizes(broadcast_dim_sizes)
+    self.assertEqual(original_shape.rank, bcast_shape.rank)
+    # shape[axis].value == 1 and row_length > 1:
+    bcast1 = original_shape.broadcast_dimension(axis, row_length)
+    # shape[axis].value > 1 and row_length == shape[axis].value:
+    bcast2 = bcast_shape.broadcast_dimension(axis, row_length)
+    # shape[axis].value > 1 and row_length == 1:
+    bcast3 = bcast_shape.broadcast_dimension(axis, 1)
+
+    self.assertShapeEq(bcast1, bcast_shape)
+    self.assertShapeEq(bcast2, bcast_shape)
+    self.assertShapeEq(bcast3, bcast_shape)
+
+  @parameterized.parameters(
+      [
+          # Broadcast scalar
+          dict(x_dims=[], y_dims=[], expected_dims=[]),
+          dict(x_dims=[], y_dims=[2], expected_dims=[2]),
+          dict(x_dims=[], y_dims=[2, 3], expected_dims=[2, 3]),
+          dict(
+              x_dims=[],
+              y_dims=[2, (2, 3), (5, 7, 2, 0, 9)],
+              expected_dims=[2, (2, 3), (5, 7, 2, 0, 9)]),
+          # Broadcast vector
+          dict(x_dims=[3], y_dims=[4, 2, 3], expected_dims=[4, 2, 3]),
+          dict(x_dims=[1], y_dims=[4, 2, 3], expected_dims=[4, 2, 3]),
+          dict(x_dims=[3], y_dims=[4, 2, 1], expected_dims=[4, 2, 3]),
+          dict(
+              x_dims=[3],
+              y_dims=[3, (2, 3, 1), 1],
+              expected_dims=[3, (2, 3, 1), 3]),
+          dict(x_dims=[1], y_dims=[3, (2, 1, 3)], expected_dims=[3, (2, 1, 3)]),
+          dict(
+              x_dims=[1],
+              y_dims=[3, (2, 1, 3), 8],
+              expected_dims=[3, (2, 1, 3), 8]),
+          dict(
+              x_dims=[1],
+              y_dims=[2, (2, 3), (5, 7, 2, 0, 9)],
+              expected_dims=[2, (2, 3), (5, 7, 2, 0, 9)]),
+          # Mixed broadcasting
+          dict(
+              x_dims=[
+                  1,  # axis=0
+                  3,  # axis=1
+                  (3, 0, 2),  # axis=2
+                  1,  # axis=3
+                  2,  # axis=4
+              ],
+              y_dims=[
+                  2,  # axis=0
+                  1,  # axis=1
+                  1,  # axis=2
+                  (7, 2),  # axis=3
+                  1,  # axis=4
+              ],
+              expected_dims=[
+                  2,  # axis=0
+                  3,  # axis=1
+                  (3, 0, 2, 3, 0, 2),  # axis=2
+                  (7, 7, 7, 7, 7, 2, 2, 2, 2, 2),  # axis=3
+                  2,  # axis=4
+              ]),
+          dict(
+              x_dims=[2, (2, 1), 2, 1],
+              y_dims=[1, 1, 2, (2, 1)],
+              expected_dims=[2, (2, 1), 2, (2, 1, 2, 1, 2, 1)]),
+      ])
+  def testBroadcastDynamicShape(self, x_dims, y_dims, expected_dims):
+    x_shape = RaggedTensorDynamicShape.from_dim_sizes(x_dims)
+    y_shape = RaggedTensorDynamicShape.from_dim_sizes(y_dims)
+    expected = RaggedTensorDynamicShape.from_dim_sizes(expected_dims)
+    result1 = ragged_tensor_shape.broadcast_dynamic_shape(x_shape, y_shape)
+    result2 = ragged_tensor_shape.broadcast_dynamic_shape(y_shape, x_shape)
+    self.assertShapeEq(expected, result1)
+    self.assertShapeEq(expected, result2)
+
+  def testRepr(self):
+    shape = RaggedTensorDynamicShape.from_dim_sizes([2, (2, 1), 2, 1])
+    self.assertRegex(
+        repr(shape), r'RaggedTensorDynamicShape\('
+        r'partitioned_dim_sizes=\(<[^>]+>, <[^>]+>\), '
+        r'inner_dim_sizes=<[^>]+>\)')
+
+  @parameterized.parameters([
+      dict(
+          x=[[10], [20], [30]],  # shape=[3, 1]
+          dim_sizes=[3, 2],
+          expected=[[10, 10], [20, 20], [30, 30]]),
+      dict(
+          x=[[10], [20], [30]],  # shape=[3, 1]
+          dim_sizes=[3, [3, 0, 2]],
+          expected=ragged_factory_ops.constant_value(
+              [[10, 10, 10], [], [30, 30]], dtype=np.int32)),
+      dict(
+          x=[[[1, 2, 3]], [[4, 5, 6]]],  # shape = [2, 1, 3]
+          dim_sizes=[2, [2, 3], 3],
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 2, 3], [1, 2, 3]], [[4, 5, 6], [4, 5, 6], [4, 5, 6]]],
+              dtype=np.int32,
+              ragged_rank=1)),
+      dict(
+          x=[[[1]], [[2]]],  # shape = [2, 1, 1]
+          dim_sizes=[2, [2, 3], [0, 2, 1, 2, 0]],
+          expected=ragged_factory_ops.constant_value(
+              [[[], [1, 1]], [[2], [2, 2], []]], dtype=np.int32,
+              ragged_rank=2)),
+      dict(
+          x=10,
+          dim_sizes=[3, [3, 0, 2]],
+          expected=ragged_factory_ops.constant_value([[10, 10, 10], [],
+                                                      [10, 10]])),
+      dict(
+          x=ragged_factory_ops.constant_value([[[1], [2]], [[3]]],
+                                              ragged_rank=1),
+          dim_sizes=[2, [2, 1], 2],
+          expected=ragged_factory_ops.constant_value(
+              [[[1, 1], [2, 2]], [[3, 3]]], ragged_rank=1)),
+  ])
+  def testRaggedBroadcastTo(self, x, dim_sizes, expected):
+    shape = RaggedTensorDynamicShape.from_dim_sizes(dim_sizes)
+    result = ragged_tensor_shape.broadcast_to(x, shape)
+    self.assertEqual(
+        getattr(result, 'ragged_rank', 0), getattr(expected, 'ragged_rank', 0))
+    self.assertAllEqual(result, expected)
+
+  @parameterized.parameters(
+      [
+          dict(
+              doc='x.shape=[3, (D1)]; y.shape=[3, 1]; bcast.shape=[3, (D1)]',
+              x=ragged_factory_ops.constant_value([[1, 2, 3], [], [4, 5]],
+                                                  dtype=np.int32),
+              y=[[10], [20], [30]],
+              expected=ragged_factory_ops.constant_value([[11, 12, 13], [],
+                                                          [34, 35]])),
+          dict(
+              doc='x.shape=[3, (D1)]; y.shape=[]; bcast.shape=[3, (D1)]',
+              x=ragged_factory_ops.constant_value([[1, 2, 3], [], [4, 5]],
+                                                  dtype=np.int32),
+              y=10,
+              expected=ragged_factory_ops.constant_value([[11, 12, 13], [],
+                                                          [14, 15]])),
+          dict(
+              doc='x.shape=[1, (D1)]; y.shape=[3, 1]; bcast.shape=[3, (D1)]',
+              x=ragged_factory_ops.constant_value([[1, 2, 3]], dtype=np.int32),
+              y=[[10], [20], [30]],
+              expected=ragged_factory_ops.constant_value(
+                  [[11, 12, 13], [21, 22, 23], [31, 32, 33]], dtype=np.int32)),
+          dict(
+              doc=('x.shape=[2, (D1), 1]; y.shape=[1, (D2)]; '
+                   'bcast.shape=[2, (D1), (D2)]'),
+              x=ragged_factory_ops.constant_value([[[1], [2], [3]], [[4]]],
+                                                  ragged_rank=1),
+              y=ragged_factory_ops.constant_value([[10, 20, 30]]),
+              expected=ragged_factory_ops.constant_value([[[11, 21, 31],
+                                                           [12, 22, 32],
+                                                           [13, 23, 33]],
+                                                          [[14, 24, 34]]])),
+          dict(
+              doc=('x.shape=[2, (D1), 1]; y.shape=[1, 1, 4]; '
+                   'bcast.shape=[2, (D1), 4]'),
+              x=ragged_factory_ops.constant_value([[[10], [20]], [[30]]],
+                                                  ragged_rank=1),
+              y=[[[1, 2, 3, 4]]],
+              expected=ragged_factory_ops.constant_value(
+                  [[[11, 12, 13, 14], [21, 22, 23, 24]], [[31, 32, 33, 34]]],
+                  ragged_rank=1)),
+          dict(
+              doc=('x.shape=[2, (D1), 2, 1]; y.shape=[2, (D2)]; '
+                   'bcast.shape=[2, (D1), (2), (D2)'),
+              x=ragged_factory_ops.constant_value(
+                  [[[[1], [2]], [[3], [4]]], [[[5], [6]]]], ragged_rank=1),
+              y=ragged_factory_ops.constant_value([[10, 20], [30]]),
+              expected=ragged_factory_ops.constant_value([[[[11, 21], [32]],
+                                                           [[13, 23], [34]]],
+                                                          [[[15, 25], [36]]]])),
+      ])
+  def testRaggedAddWithBroadcasting(self, x, y, expected, doc):
+    expected_rrank = getattr(expected, 'ragged_rank', 0)
+    x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x, dtype=dtypes.int32)
+    y = ragged_tensor.convert_to_tensor_or_ragged_tensor(y, dtype=dtypes.int32)
+    result = x + y
+    result_rrank = getattr(result, 'ragged_rank', 0)
+    self.assertEqual(expected_rrank, result_rrank)
+    if hasattr(expected, 'tolist'):
+      expected = expected.tolist()
+    self.assertAllEqual(result, expected)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_tensor_supported_values_test.py

@@ -0,0 +1,594 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""Tests for RaggedTensor supported value types."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import extension_type
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework import type_spec
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import clip_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import string_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_test_ops as test_ops
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensorSpec
+from tensorflow.python.platform import googletest
+from tensorflow.python.util import dispatch
+
+
+class WrappedTensor(extension_type.BatchableExtensionType):
+  value: tensor.Tensor
+
+  @property
+  def shape(self):
+    return self.value.shape
+
+  @property
+  def dtype(self):
+    return self.value.dtype
+
+  def __getitem__(self, idx):
+    return WrappedTensor(self.value.__getitem__(idx))
+
+  def set_shape(self, shape):
+    return self.value.set_shape(shape)
+
+  class Spec(type_spec.TypeSpec):
+
+    @property
+    def shape(self):
+      return self.value.shape
+
+    @property
+    def dtype(self):
+      return self.value.dtype
+
+
+class WrappedTensorOpDispatcher(dispatch.GlobalOpDispatcher):
+  """Global op dispatcher for WrappedTensor."""
+
+  # For these ops, just return plain Tensors (not WrappedTensors).
+  OPS_THAT_RETURN_UNTRACED_RESULTS = (array_ops.shape, array_ops.shape_v2,
+                                      check_ops.assert_rank_at_least)
+
+  def call_op(self, op, *args, **kwargs):
+    return op(*args, **kwargs)
+
+  def handle(self, op, args, kwargs):
+    # Dispatcher only applies if at least one arg is a WrappedTensor.
+    if not (any(self.is_wrapped_tensor_arg(x) for x in args) or
+            any(self.is_wrapped_tensor_arg(x) for x in kwargs.values())):
+      return self.NOT_SUPPORTED
+
+    args = [self.unwrap(v) for v in args]
+    kwargs = dict([(k, self.unwrap(v)) for (k, v) in kwargs.items()])
+    value = self.call_op(op, *args, **kwargs)
+    if op in self.OPS_THAT_RETURN_UNTRACED_RESULTS:
+      return value
+    else:
+      return WrappedTensor(value)
+
+  def unwrap(self, value):
+    if isinstance(value, WrappedTensor):
+      return value.value
+    elif isinstance(value, (list, tuple)):
+      return type(value)([self.unwrap(v) for v in value])
+    else:
+      return value
+
+  def is_wrapped_tensor_arg(self, value):
+    if isinstance(value, WrappedTensor):
+      return True
+    if isinstance(value, (list, tuple)):
+      if any(isinstance(x, WrappedTensor) for x in value):
+        return True
+    return False
+
+
+WrappedTensorOpDispatcher().register()
+ragged_tensor._add_supported_value_type(WrappedTensor)
+
+
+# pylint: disable=g-complex-comprehension
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorSupportedValuesTest(test_util.TensorFlowTestCase,
+                                      parameterized.TestCase):
+
+  def assertAllTensorsEqual(self, list1, list2):
+    self.assertLen(list1, len(list2))
+    for (t1, t2) in zip(list1, list2):
+      self.assertAllEqual(t1, t2)
+
+  def testConstruction(self):
+    tensor_values = constant_op.constant(
+        ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'])
+    values = WrappedTensor(tensor_values)
+
+    row_splits = constant_op.constant([0, 2, 2, 5, 6, 8], dtypes.int64)
+    rt = RaggedTensor.from_row_splits(values, row_splits)
+    self.assertIsInstance(rt.values, WrappedTensor)
+    self.assertAllEqual(rt.values.value, tensor_values)
+    self.assertAllEqual(rt.row_splits, row_splits)
+
+    row_starts = constant_op.constant([0, 2, 2, 5, 6], dtypes.int64)
+    rt = RaggedTensor.from_row_starts(values, row_starts)
+    self.assertIsInstance(rt.values, WrappedTensor)
+    self.assertAllEqual(rt.values.value, tensor_values)
+    self.assertAllEqual(rt.row_starts(), row_starts)
+
+    row_limits = constant_op.constant([2, 2, 5, 6, 8], dtypes.int64)
+    rt = RaggedTensor.from_row_limits(values, row_limits)
+    self.assertIsInstance(rt.values, WrappedTensor)
+    self.assertAllEqual(rt.values.value, tensor_values)
+    self.assertAllEqual(rt.row_limits(), row_limits)
+
+    row_lengths = constant_op.constant([2, 0, 3, 1, 2], dtypes.int64)
+    rt = RaggedTensor.from_row_lengths(values, row_lengths)
+    self.assertIsInstance(rt.values, WrappedTensor)
+    self.assertAllEqual(rt.values.value, tensor_values)
+    self.assertAllEqual(rt.row_lengths(), row_lengths)
+
+    rt = RaggedTensor.from_uniform_row_length(values, 4)
+    self.assertIsInstance(rt.values, WrappedTensor)
+    self.assertAllEqual(rt.values.value, tensor_values)
+    self.assertAllEqual(rt.uniform_row_length, 4)
+
+  def testWithValues(self):
+    tensor_values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
+    values = WrappedTensor(tensor_values)
+    nested_row_splits = [[0, 2, 5], [0, 2, 2, 5, 6, 7]]
+    rt = RaggedTensor.from_nested_row_splits(values, nested_row_splits)
+
+    tensor_int = constant_op.constant([1, 2, 3, 4, 5])
+    rt_int = rt.with_values(tensor_int)
+    self.assertAllEqual(rt_int.values, tensor_int)
+
+    rt_wrapped_int = rt.with_values(WrappedTensor(tensor_int))
+    self.assertIsInstance(rt_wrapped_int.values, WrappedTensor)
+    self.assertAllEqual(rt_wrapped_int.values.value, tensor_int)
+
+  def testWithFlatValues(self):
+    tensor_values = constant_op.constant(['a', 'b', 'c', 'd', 'e', 'f', 'g'])
+    values = WrappedTensor(tensor_values)
+    nested_row_splits = [[0, 2, 5], [0, 2, 2, 5, 6, 7]]
+    rt = RaggedTensor.from_nested_row_splits(values, nested_row_splits)
+
+    tensor_int = constant_op.constant([1, 2, 3, 4, 5, 6, 7])
+    rt_int = rt.with_flat_values(tensor_int)
+    self.assertAllEqual(rt_int.flat_values, tensor_int)
+
+    rt_wrapped_int = rt.with_flat_values(WrappedTensor(tensor_int))
+    self.assertIsInstance(rt_wrapped_int.flat_values, WrappedTensor)
+    self.assertAllEqual(rt_wrapped_int.flat_values.value, tensor_int)
+
+  @parameterized.parameters(
+      #=========================================================================
+      # Test each unary op.
+      #=========================================================================
+      [{'x': ([[-2.0, 3.0], [-3.0]]), 'op': op}
+       for op in test_ops.UNARY_FLOAT_OPS] +
+      [{'x': ([[True, False], [True]]),
+        'op': op}
+       for op in test_ops.UNARY_BOOL_OPS] +
+      [{'x': [[18, 512], [12412]],
+        'x_dtype': dtypes.int32,
+        'op': op}
+       for op in test_ops.UNARY_INT_OPS] +
+      [{'x': ([['abcd', 'efgh'], ['aabbccdd']]),
+        'op': op}
+       for op in test_ops.UNARY_STRING_OPS] +
+      [
+          {'op': clip_ops.clip_by_value,
+           'x': ([[-2.0, 3.0], [-3.0]]),
+           'clip_value_min': 0.1, 'clip_value_max': 4.0},
+          {'op': math_ops.cast,
+           'x': ([[-2.0, 3.0], [-3.0]]),
+           'dtype': dtypes.int32},
+          {'op': math_ops.saturate_cast,
+           'x': ([[-2.0, 3.0], [-3.0]]),
+           'dtype': dtypes.int32},
+          {'op': string_ops.string_to_hash_bucket,
+           'x': (
+               [['abcd', 'efgh'], ['aabbccdd']]),
+           'num_buckets': 1000},
+          {'op': string_ops.string_to_hash_bucket_fast,
+           'x': (
+               [['abcd', 'efgh'], ['aabbccdd']]),
+           'num_buckets': 1000},
+          {'op': string_ops.string_to_hash_bucket_strong,
+           'x': (
+               [['abcd', 'efgh'], ['aabbccdd']]),
+           'num_buckets': 1000,
+           'key': [1231, 12512]},
+          {'op': string_ops.string_to_number,
+           'x': ([['-2.0', '3.0'], ['-3.0']])},
+          {'op': string_ops.regex_full_match,
+           'x': ([['hello', '123'], ['1+1']]),
+           'pattern': r'\w+'},
+          {'op': string_ops.regex_replace,
+           'x': ([['hello', '123'], ['1+1']]),
+           'pattern': r'\d',
+           'rewrite': '#'},
+          {'op': string_ops.substr,
+           'x': ([['hello', '123'], ['1+1']]),
+           'pos': 2, 'len': 3},
+          {'op': array_ops.check_numerics,
+           'x': ([[-2.0, 3.0], [-3.0]]),
+           'message': 'check-numerics'},
+          {'op': nn_ops.dropout,
+           'x': ([[-2.0, 3.0], [-3.0]]),
+           'rate': 0.5,
+           'seed': 1},
+          {'op': array_ops.expand_dims_v2,
+           'x': ([[-2.0, 3.0], [-3.0]]),
+           'axis': -1},
+      ])  # pyformat: disable
+  def testUnaryElementwiseOp(self,
+                             x,
+                             x_dtype=None,
+                             op=math_ops.abs,
+                             **extra_args):
+    x = ragged_factory_ops.constant(x, x_dtype)
+    wrapped_x = ragged_tensor.RaggedTensor.from_nested_row_splits(
+        WrappedTensor(x.flat_values), x.nested_row_splits)
+    test_util.random_seed.set_seed(1234)
+    res = op(x, **extra_args)
+    test_util.random_seed.set_seed(1234)
+    wrapped_res = op(wrapped_x, **extra_args)
+    self.assertIsInstance(wrapped_res.flat_values, WrappedTensor)
+    self.assertAllEqual(wrapped_res.flat_values.value, res.flat_values)
+    self.assertAllTensorsEqual(wrapped_res.nested_row_splits,
+                               res.nested_row_splits)
+
+  @parameterized.parameters(
+      #=========================================================================
+      # Test each binary op.
+      #=========================================================================
+      [{'x': [[-2.0, 3.0], [-3.0]],
+        'y': [[5.0, 1.0], [12.0]],
+        'op': op}
+       for op in test_ops.BINARY_FLOAT_OPS] +
+      [{'x': [[-2, 3], [-3]],
+        'y': [[5, 1], [12]],
+        'op': op}
+       for op in test_ops.BINARY_INT_OPS] +
+      [{'x': [[True, True], [False]],
+        'y': [[False, True], [False]],
+        'op': op}
+       for op in test_ops.BINARY_BOOL_OPS]
+      )  # pyformat: disable
+  def testBinaryElementwiseOp(self, x, y, op=math_ops.add):
+    x = ragged_factory_ops.constant(x)
+    y = ragged_factory_ops.constant(y)
+    wrapped_x = ragged_tensor.RaggedTensor.from_nested_row_splits(
+        WrappedTensor(x.flat_values), x.nested_row_splits)
+    wrapped_y = ragged_tensor.RaggedTensor.from_nested_row_splits(
+        WrappedTensor(y.flat_values), y.nested_row_splits)
+    res = op(x, y)
+    wrapped_res = op(wrapped_x, wrapped_y)
+    self.assertIsInstance(wrapped_res.flat_values, WrappedTensor)
+    self.assertAllEqual(wrapped_res.flat_values.value, res.flat_values)
+    self.assertAllTensorsEqual(wrapped_res.nested_row_splits,
+                               res.nested_row_splits)
+
+  @parameterized.parameters(
+      (lambda x: x[1:], True),
+      (lambda x: x[0], False),
+      (lambda x: x[:], True),
+      (lambda x: x[0, :], False),
+      (lambda x: x[...], True),
+      (lambda x: x[1:2, ...], True),
+      (lambda x: x[..., 1:2], True),
+      (lambda x: x[0:2, ::2], True),
+  )
+  def testSlicing(self, slice_fn, is_ragged_output):
+    tensor_values = constant_op.constant([[1.0, 2], [3, 4], [5, 6], [7, 8]])
+    row_splits = constant_op.constant([0, 2, 3, 4], dtypes.int32)
+    raw_rt = RaggedTensor.from_row_splits(tensor_values, row_splits)
+
+    values = WrappedTensor(tensor_values)
+    rt = RaggedTensor.from_row_splits(values, row_splits)
+
+    res = slice_fn(rt)
+    raw_res = slice_fn(raw_rt)
+    if is_ragged_output:
+      self.assertIsInstance(res, RaggedTensor)
+      self.assertIsInstance(res.flat_values, WrappedTensor)
+      self.assertAllEqual(res.flat_values.value, raw_res.flat_values)
+      self.assertAllTensorsEqual(res.nested_row_splits,
+                                 raw_res.nested_row_splits)
+    else:
+      self.assertIsInstance(res, WrappedTensor)
+      self.assertAllEqual(res.value, raw_res)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorSpecSupportedValuesTest(test_util.TensorFlowTestCase,
+                                          parameterized.TestCase):
+
+  def assertAllTensorsEqual(self, list1, list2):
+    self.assertLen(list1, len(list2))
+    for (t1, t2) in zip(list1, list2):
+      self.assertAllEqual(t1, t2)
+
+  def testConstruction(self):
+    flat_values_spec = WrappedTensor.Spec(
+        tensor.TensorSpec(shape=(None, 5), dtype=dtypes.float32))
+    spec1 = RaggedTensorSpec(
+        shape=None,
+        dtype=dtypes.float32,
+        ragged_rank=1,
+        row_splits_dtype=dtypes.int64,
+        flat_values_spec=flat_values_spec)
+    self.assertIsNone(spec1._shape.rank)
+    self.assertEqual(spec1._dtype, dtypes.float32)
+    self.assertEqual(spec1._row_splits_dtype, dtypes.int64)
+    self.assertEqual(spec1._ragged_rank, 1)
+    self.assertEqual(spec1._flat_values_spec, flat_values_spec)
+
+    self.assertIsNone(spec1.shape.rank)
+    self.assertEqual(spec1.dtype, dtypes.float32)
+    self.assertEqual(spec1.row_splits_dtype, dtypes.int64)
+    self.assertEqual(spec1.ragged_rank, 1)
+    self.assertEqual(spec1.flat_values_spec, flat_values_spec)
+
+    with self.assertRaisesRegex(
+        ValueError, 'dtype must be the same as flat_values_spec.dtype'
+    ):
+      spec1 = RaggedTensorSpec(
+          shape=None,
+          dtype=dtypes.float64,
+          ragged_rank=1,
+          row_splits_dtype=dtypes.int64,
+          flat_values_spec=flat_values_spec,
+      )
+
+  @parameterized.parameters([
+      (
+          RaggedTensorSpec(
+              ragged_rank=1,
+              flat_values_spec=tensor.TensorSpec(None, dtypes.float32),
+          ),
+          (
+              tensor_shape.TensorShape(None),
+              dtypes.float32,
+              1,
+              dtypes.int64,
+              tensor.TensorSpec(None, dtypes.float32),
+          ),
+      ),
+      (
+          RaggedTensorSpec(
+              shape=(5, None, 5),
+              ragged_rank=1,
+              dtype=dtypes.float64,
+              flat_values_spec=tensor.TensorSpec((5,), dtypes.float64),
+          ),
+          (
+              tensor_shape.TensorShape((5, None, 5)),
+              dtypes.float64,
+              1,
+              dtypes.int64,
+              tensor.TensorSpec((5,), dtypes.float64),
+          ),
+      ),
+  ])
+  def testSerialize(self, rt_spec, expected):
+    serialization = rt_spec._serialize()
+    # TensorShape has an unconventional definition of equality, so we can't use
+    # assertEqual directly here.  But repr() is deterministic and lossless for
+    # the expected values, so we can use that instead.
+    self.assertEqual(repr(serialization), repr(expected))
+
+  @parameterized.parameters([
+      (RaggedTensorSpec(
+          ragged_rank=0,
+          shape=[5, 3],
+          flat_values_spec=WrappedTensor.Spec(
+              tensor.TensorSpec([5, 3], dtypes.float32))),
+       [WrappedTensor.Spec(tensor.TensorSpec([5, 3], dtypes.float32))]),
+      (RaggedTensorSpec(
+          ragged_rank=1,
+          flat_values_spec=WrappedTensor.Spec(
+              tensor.TensorSpec([None, 3], dtypes.float32))),
+       [
+           WrappedTensor.Spec(
+               tensor.TensorSpec([None, 3], dtypes.float32)),
+           tensor.TensorSpec([None], dtypes.int64),
+       ]),
+      (RaggedTensorSpec(
+          ragged_rank=2,
+          dtype=dtypes.float64,
+          flat_values_spec=WrappedTensor.Spec(
+              tensor.TensorSpec([None, 3], dtypes.float64))),
+       [
+           WrappedTensor.Spec(
+               tensor.TensorSpec([None, 3], dtypes.float64)),
+           tensor.TensorSpec([None], dtypes.int64),
+           tensor.TensorSpec([None], dtypes.int64),
+       ]),
+      (RaggedTensorSpec(
+          shape=[5, None, None],
+          dtype=dtypes.string,
+          flat_values_spec=WrappedTensor.Spec(
+              tensor.TensorSpec([None, 3], dtypes.string))),
+       [
+           WrappedTensor.Spec(tensor.TensorSpec([None, 3], dtypes.string)),
+           tensor.TensorSpec([6], dtypes.int64),
+           tensor.TensorSpec([None], dtypes.int64),
+       ]),
+  ])
+  def testComponentSpecs(self, rt_spec, expected):
+    self.assertEqual(rt_spec._component_specs, expected)
+
+  @parameterized.parameters([
+      {
+          'rt_spec':
+              RaggedTensorSpec(
+                  shape=[3, None, None],
+                  ragged_rank=1,
+                  flat_values_spec=WrappedTensor.Spec(
+                      tensor.TensorSpec(None, dtype=dtypes.float32))),
+          'flat_values': [[1.0, 2.0], [3.0, 4.0]],
+          'nested_row_splits': [[0, 1, 1, 2]],
+      },
+      {
+          'rt_spec':
+              RaggedTensorSpec(
+                  shape=[2, None, None],
+                  flat_values_spec=WrappedTensor.Spec(
+                      tensor.TensorSpec(None, dtype=dtypes.float32))),
+          'flat_values': [1.0, 2.0, 3.0, 4.0],
+          'nested_row_splits': [[0, 2, 4], [0, 2, 3, 3, 4]],
+      },
+  ])
+  def testToFromComponents(self, rt_spec, flat_values, nested_row_splits):
+    wrapped_tensor = WrappedTensor(constant_op.constant(flat_values))
+    rt = RaggedTensor.from_nested_row_splits(wrapped_tensor, nested_row_splits)
+    components = rt_spec._to_components(rt)
+    self.assertIsInstance(components[0], WrappedTensor)
+    self.assertAllEqual(components[0].value, wrapped_tensor.value)
+    self.assertAllTensorsEqual(components[1:], nested_row_splits)
+    rt_reconstructed = rt_spec._from_components(components)
+    self.assertIsInstance(rt_reconstructed.flat_values, WrappedTensor)
+    self.assertAllEqual(
+        rt_reconstructed.flat_values.value, wrapped_tensor.value
+    )
+    self.assertAllTensorsEqual(
+        rt_reconstructed.nested_row_splits, rt.nested_row_splits
+    )
+    self.assertEqual(rt_reconstructed.dtype, rt.dtype)
+
+  def testIsCompatibleWith(self):
+    spec1 = RaggedTensorSpec(
+        [32, None, None],
+        dtypes.float32,
+        2,
+        flat_values_spec=WrappedTensor.Spec(
+            tensor.TensorSpec([None, None], dtypes.float32)
+        ),
+    )
+    spec2 = RaggedTensorSpec(
+        None,
+        dtypes.float32,
+        2,
+        flat_values_spec=WrappedTensor.Spec(
+            tensor.TensorSpec(None, dtypes.float32)
+        ),
+    )
+    spec3 = RaggedTensorSpec(
+        None,
+        dtypes.int32,
+        1,
+        flat_values_spec=WrappedTensor.Spec(
+            tensor.TensorSpec(None, dtypes.int32)
+        ),
+    )
+    spec4 = RaggedTensorSpec(
+        [None],
+        dtypes.int32,
+        0,
+        flat_values_spec=WrappedTensor.Spec(
+            tensor.TensorSpec(None, dtypes.int32)
+        ),
+    )
+    spec5 = RaggedTensorSpec([None], dtypes.int32, 0)
+
+    self.assertTrue(spec1.is_compatible_with(spec2))
+    self.assertFalse(spec1.is_compatible_with(spec3))
+    self.assertFalse(spec1.is_compatible_with(spec4))
+    self.assertFalse(spec2.is_compatible_with(spec3))
+    self.assertFalse(spec2.is_compatible_with(spec4))
+    self.assertFalse(spec3.is_compatible_with(spec4))
+    self.assertFalse(spec4.is_compatible_with(spec5))
+    value = constant_op.constant([1, 2, 3])
+    self.assertFalse(spec4.is_compatible_with(value))
+    self.assertTrue(spec4.is_compatible_with(WrappedTensor(value)))
+
+  def testToList(self):
+    with context.eager_mode():
+      tensor_values = constant_op.constant(
+          ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h'])
+      row_splits = constant_op.constant([0, 2, 2, 5, 6, 8], dtypes.int64)
+      values = WrappedTensor(tensor_values)
+      rt = RaggedTensor.from_row_splits(values, row_splits)
+      expected = ragged_factory_ops.constant([['a', 'b'], [], ['c', 'd', 'e'],
+                                              ['f'], ['g', 'h']]).to_list()
+
+      with self.subTest('Raise on unsupported'):
+        with self.assertRaisesRegex(
+            ValueError,
+            'values must be convertible to a list',
+        ):
+          _ = rt.to_list()
+
+      with self.subTest('Value with numpy method'):
+
+        class WrappedTensorWithNumpy(WrappedTensor):
+
+          def numpy(self):
+            return self.value.numpy()
+
+        values = WrappedTensorWithNumpy(tensor_values)
+        rt = RaggedTensor.from_row_splits(values, row_splits)
+        self.assertEqual(rt.to_list(), expected)
+
+      with self.subTest('Value with to_list method'):
+
+        class WrappedTensorWithToList(WrappedTensor):
+
+          def to_list(self):
+            return self.value.numpy().tolist()
+
+        values = WrappedTensorWithToList(tensor_values)
+        rt = RaggedTensor.from_row_splits(values, row_splits)
+        self.assertEqual(rt.to_list(), expected)
+
+  def testFromValue(self):
+    tensor_values = constant_op.constant([[1.0, 2], [4, 5], [7, 8]])
+    values = WrappedTensor(tensor_values)
+
+    row_splits = constant_op.constant([0, 2, 3, 3, 3], dtypes.int32)
+    rt = RaggedTensor.from_row_splits(values, row_splits)
+
+    rt_spec = type_spec.type_spec_from_value(rt)
+    self.assertEqual(
+        rt_spec,
+        RaggedTensorSpec(
+            shape=[4, None, 2],
+            dtype=dtypes.float32,
+            ragged_rank=1,
+            row_splits_dtype=dtypes.int32,
+            flat_values_spec=WrappedTensor.Spec(
+                tensor.TensorSpec([None, 2], dtypes.float32))))
+    # Ensure the shape of flat_values_spec being consistent with the shape
+    # of the RaggedTensor.
+    self.assertEqual(rt_spec.shape[rt_spec.ragged_rank:],
+                     rt_spec.flat_values_spec.shape)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_tensor_test_ops.py

@@ -0,0 +1,185 @@
+# Copyright 2020 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.
+# ==============================================================================
+"""It lists ops of RaggedTensor for the interest of test."""
+
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import gen_bitwise_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops import nn_impl
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import parsing_ops
+from tensorflow.python.ops import special_math_ops
+from tensorflow.python.ops import string_ops
+
+
+# Constants listing various op types to test.  Each operation
+# should be included in at least one list below, or tested separately if
+# necessary (e.g., because it expects additional arguments).
+UNARY_FLOAT_OPS = [
+    math_ops.abs,
+    math_ops.acos,
+    math_ops.acosh,
+    math_ops.angle,
+    math_ops.asin,
+    math_ops.asinh,
+    math_ops.atan,
+    math_ops.atanh,
+    math_ops.ceil,
+    math_ops.conj,
+    math_ops.cos,
+    math_ops.cosh,
+    math_ops.digamma,
+    math_ops.erf,
+    math_ops.erfc,
+    math_ops.erfcinv,
+    math_ops.erfinv,
+    math_ops.exp,
+    math_ops.expm1,
+    math_ops.floor,
+    math_ops.imag,
+    math_ops.is_finite,
+    math_ops.is_inf,
+    math_ops.is_nan,
+    math_ops.lgamma,
+    math_ops.log,
+    math_ops.log1p,
+    math_ops.log_sigmoid,
+    math_ops.ndtri,
+    math_ops.negative,
+    math_ops.real,
+    math_ops.reciprocal,
+    math_ops.reciprocal_no_nan,
+    math_ops.rint,
+    math_ops.round,
+    math_ops.rsqrt,
+    math_ops.sign,
+    math_ops.sigmoid,
+    math_ops.sin,
+    math_ops.sinh,
+    math_ops.softplus,
+    math_ops.sqrt,
+    math_ops.square,
+    math_ops.tan,
+    math_ops.tanh,
+    nn_ops.elu,
+    nn_ops.gelu,
+    nn_ops.leaky_relu,
+    nn_ops.log_softmax,
+    nn_ops.relu,
+    nn_ops.relu6,
+    nn_ops.selu,
+    nn_ops.softsign,
+    nn_impl.swish,
+    array_ops.ones_like,
+    array_ops.ones_like_v2,
+    array_ops.zeros_like,
+    array_ops.zeros_like_v2,
+    special_math_ops.bessel_i0,
+    special_math_ops.bessel_i0e,
+    special_math_ops.bessel_i1,
+    special_math_ops.bessel_j0,
+    special_math_ops.bessel_j1,
+    special_math_ops.bessel_i1e,
+    special_math_ops.bessel_k0,
+    special_math_ops.bessel_k0e,
+    special_math_ops.bessel_k1,
+    special_math_ops.bessel_k1e,
+    special_math_ops.bessel_y0,
+    special_math_ops.bessel_y1,
+    special_math_ops.dawsn,
+    special_math_ops.expint,
+    special_math_ops.fresnel_cos,
+    special_math_ops.fresnel_sin,
+    special_math_ops.spence,
+    string_ops.as_string,
+]
+UNARY_BOOL_OPS = [
+    math_ops.logical_not,
+]
+UNARY_STRING_OPS = [
+    string_ops.decode_base64,
+    string_ops.encode_base64,
+    string_ops.string_strip,
+    string_ops.string_lower,
+    string_ops.string_upper,
+    string_ops.string_length,
+    string_ops.string_length_v2,
+    parsing_ops.decode_compressed,
+]
+BINARY_FLOAT_OPS = [
+    math_ops.add,
+    math_ops.atan2,
+    math_ops.complex,
+    math_ops.div,
+    math_ops.div_no_nan,
+    math_ops.divide,
+    math_ops.equal,
+    math_ops.floor_div,
+    math_ops.floordiv,
+    math_ops.floormod,
+    math_ops.greater,
+    math_ops.greater_equal,
+    math_ops.less,
+    math_ops.less_equal,
+    math_ops.maximum,
+    math_ops.minimum,
+    math_ops.multiply,
+    math_ops.multiply_no_nan,
+    math_ops.not_equal,
+    math_ops.pow,
+    math_ops.realdiv,
+    math_ops.squared_difference,
+    math_ops.subtract,
+    math_ops.truediv,
+    math_ops.xdivy,
+    math_ops.xlog1py,
+    math_ops.xlogy,
+    math_ops.zeta,
+]
+BINARY_BOOL_OPS = [
+    math_ops.logical_and,
+    math_ops.logical_or,
+    math_ops.logical_xor,
+]
+UNARY_INT_OPS = [
+    gen_bitwise_ops.invert,
+    string_ops.unicode_script,
+]
+BINARY_INT_OPS = [
+    gen_bitwise_ops.bitwise_and,
+    gen_bitwise_ops.bitwise_or,
+    gen_bitwise_ops.bitwise_xor,
+    gen_bitwise_ops.left_shift,
+    gen_bitwise_ops.right_shift,
+    math_ops.truncatediv,
+    math_ops.truncatemod,
+]
+BINARY_ASSERT_OPS = [
+    check_ops.assert_equal,
+    check_ops.assert_equal_v2,
+    check_ops.assert_near,
+    check_ops.assert_near_v2,
+    check_ops.assert_none_equal,
+    check_ops.assert_none_equal_v2,
+    check_ops.assert_greater,
+    check_ops.assert_greater_v2,
+    check_ops.assert_greater_equal,
+    check_ops.assert_greater_equal_v2,
+    check_ops.assert_less,
+    check_ops.assert_less_v2,
+    check_ops.assert_less_equal,
+    check_ops.assert_less_equal_v2,
+]

tensorflow/python/ops/ragged/ragged_tensor_value.py

@@ -0,0 +1,114 @@
+# 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.
+# ==============================================================================
+"""Value for RaggedTensor."""
+
+import numpy as np
+
+from tensorflow.python.ops.ragged.row_partition import RowPartition
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+@tf_export(v1=["ragged.RaggedTensorValue"])
+@dispatch.register_dispatchable_type
+class RaggedTensorValue:
+  """Represents the value of a `RaggedTensor`.
+
+  Warning: `RaggedTensorValue` should only be used in graph mode; in
+  eager mode, the `tf.RaggedTensor` class contains its value directly.
+
+  See `tf.RaggedTensor` for a description of ragged tensors.
+  """
+
+  def __init__(self, values, row_splits):
+    """Creates a `RaggedTensorValue`.
+
+    Args:
+      values: A numpy array of any type and shape; or a RaggedTensorValue.
+      row_splits: A 1-D int32 or int64 numpy array.
+    """
+    if not (isinstance(row_splits, (np.ndarray, np.generic)) and
+            row_splits.dtype in (np.int64, np.int32) and row_splits.ndim == 1):
+      raise TypeError("row_splits must be a 1D int32 or int64 numpy array")
+    if not isinstance(values, (np.ndarray, np.generic, RaggedTensorValue)):
+      raise TypeError("values must be a numpy array or a RaggedTensorValue")
+    if (isinstance(values, RaggedTensorValue) and
+        row_splits.dtype != values.row_splits.dtype):
+      raise ValueError("row_splits and values.row_splits must have "
+                       "the same dtype")
+    self._values = values
+    self._row_splits = row_splits
+
+  row_splits = property(
+      lambda self: self._row_splits,
+      doc="""The split indices for the ragged tensor value.""")
+  values = property(
+      lambda self: self._values,
+      doc="""The concatenated values for all rows in this tensor.""")
+  dtype = property(
+      lambda self: self._values.dtype,
+      doc="""The numpy dtype of values in this tensor.""")
+
+  @property
+  def flat_values(self):
+    """The innermost `values` array for this ragged tensor value."""
+    rt_values = self.values
+    while isinstance(rt_values, RaggedTensorValue):
+      rt_values = rt_values.values
+    return rt_values
+
+  @property
+  def nested_row_splits(self):
+    """The row_splits for all ragged dimensions in this ragged tensor value."""
+    rt_nested_splits = [self.row_splits]
+    rt_values = self.values
+    while isinstance(rt_values, RaggedTensorValue):
+      rt_nested_splits.append(rt_values.row_splits)
+      rt_values = rt_values.values
+    return tuple(rt_nested_splits)
+
+  @property
+  def ragged_rank(self):
+    """The number of ragged dimensions in this ragged tensor value."""
+    values_is_ragged = isinstance(self._values, RaggedTensorValue)
+    return self._values.ragged_rank + 1 if values_is_ragged else 1
+
+  @property
+  def shape(self):
+    """A tuple indicating the shape of this RaggedTensorValue."""
+    return (self._row_splits.shape[0] - 1,) + (None,) + self._values.shape[1:]
+
+  @property
+  def _nested_row_partitions(self):
+    """The row_partitions representing this shape."""
+    return [RowPartition.from_row_splits(rs) for rs in self.nested_row_splits]
+
+  def __str__(self):
+    return "<tf.RaggedTensorValue %s>" % self.to_list()
+
+  def __repr__(self):
+    return "tf.RaggedTensorValue(values=%r, row_splits=%r)" % (self._values,
+                                                               self._row_splits)
+
+  def to_list(self):
+    """Returns this ragged tensor value as a nested Python list."""
+    if isinstance(self._values, RaggedTensorValue):
+      values_as_list = self._values.to_list()
+    else:
+      values_as_list = self._values.tolist()
+    return [
+        values_as_list[self._row_splits[i]:self._row_splits[i + 1]]
+        for i in range(len(self._row_splits) - 1)
+    ]

tensorflow/python/ops/ragged/ragged_tile_op_test.py

@@ -0,0 +1,220 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.tile."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_array_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTileOpTest(test_util.TensorFlowTestCase,
+                       parameterized.TestCase):
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Docstring Example
+      #=========================================================================
+      dict(
+          descr='docstring example: ragged_rank=1, repeat axes 0 and 1',
+          rt_input=[[1, 2], [3]],
+          multiples=[3, 2],
+          expected=[
+              [1, 2, 1, 2], [3, 3], [1, 2, 1, 2], [3, 3], [1, 2, 1, 2], [3, 3]],
+      ),
+
+      #=========================================================================
+      # rank=3, ragged_rank=2
+      #=========================================================================
+      dict(
+          descr='rank=3, ragged_rank=2, repeat axis 0',
+          rt_input=[[[1, 2], [3]], [], [[4]]],
+          multiples=[2, 1, 1],
+          expected=[[[1, 2], [3]], [], [[4]],
+                    [[1, 2], [3]], [], [[4]]]),
+      dict(
+          descr='rank=3, ragged_rank=2, repeat axis 1',
+          rt_input=[[[1, 2], [3]], [], [[4]]],
+          multiples=[1, 2, 1],
+          expected=[[[1, 2], [3], [1, 2], [3]], [], [[4], [4]]]),
+      dict(
+          descr='rank=3, ragged_rank=2, repeat axis 2',
+          rt_input=[[[1, 2], [3]], [], [[4]]],
+          multiples=[1, 1, 2],
+          expected=[[[1, 2, 1, 2], [3, 3]], [], [[4, 4]]]),
+      dict(
+          descr='rank=3, ragged_rank=2, repeat axes 0 and 1',
+          rt_input=[[[1, 2], [3]], [], [[4]]],
+          multiples=[2, 2, 1],
+          expected=[[[1, 2], [3], [1, 2], [3]], [], [[4], [4]],
+                    [[1, 2], [3], [1, 2], [3]], [], [[4], [4]]]),
+      dict(
+          descr='rank=3, ragged_rank=2, repeat axes 0 and 2',
+          rt_input=[[[1, 2], [3]], [], [[4]]],
+          multiples=[2, 1, 2],
+          expected=[[[1, 2, 1, 2], [3, 3]], [], [[4, 4]],
+                    [[1, 2, 1, 2], [3, 3]], [], [[4, 4]]]),
+      dict(
+          descr='rank=3, ragged_rank=2, repeat axes 1 and 2',
+          rt_input=[[[1, 2], [3]], [], [[4]]],
+          multiples=[1, 2, 2],
+          expected=[[[1, 2, 1, 2], [3, 3], [1, 2, 1, 2], [3, 3]],
+                    [], [[4, 4], [4, 4]]]),
+      dict(
+          descr='rank=3, ragged_rank=2, repeat all axes',
+          rt_input=[[['a', 'b'], ['c']], [], [['d']]],
+          multiples=[4, 3, 2],
+          expected=[[[b'a', b'b']*2, [b'c']*2]*3, []*3, [[b'd']*2]*3]*4),
+      #=========================================================================
+      # rank=3, ragged_rank=1
+      #=========================================================================
+      dict(
+          descr='rank=3, ragged_rank=1, repeat axis 0',
+          ragged_rank=1,
+          rt_input=[[[1, 2], [3, 4]], [], [[5, 6]]],
+          multiples=[2, 1, 1],
+          expected=[[[1, 2], [3, 4]], [], [[5, 6]],
+                    [[1, 2], [3, 4]], [], [[5, 6]]]),
+      dict(
+          descr='rank=3, ragged_rank=1, repeat axis 1',
+          ragged_rank=1,
+          rt_input=[[[1, 2], [3, 4]], [], [[5, 6]]],
+          multiples=[1, 2, 1],
+          expected=[[[1, 2], [3, 4], [1, 2], [3, 4]], [], [[5, 6], [5, 6]]]),
+      dict(
+          descr='rank=3, ragged_rank=1, repeat axis 2',
+          ragged_rank=1,
+          rt_input=[[[1, 2], [3, 4]], [], [[5, 6]]],
+          multiples=[1, 1, 2],
+          expected=[[[1, 2, 1, 2], [3, 4, 3, 4]], [], [[5, 6, 5, 6]]]),
+      #=========================================================================
+      # rank=4, ragged_rank=3
+      #=========================================================================
+      dict(
+          descr='rank=4, ragged_rank=3, repeat axis 0',
+          rt_input=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]]],
+          multiples=[2, 1, 1, 1],
+          expected=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]],
+                    [[[1], [2]], [[3]]], [[]], [[[4, 5]]]]),
+      dict(
+          descr='rank=4, ragged_rank=3, repeat axis 1',
+          rt_input=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]]],
+          multiples=[1, 2, 1, 1],
+          expected=[[[[1], [2]], [[3]], [[1], [2]], [[3]]],
+                    [[], []],
+                    [[[4, 5]], [[4, 5]]]]),
+      dict(
+          descr='rank=4, ragged_rank=3, repeat axis 2',
+          rt_input=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]]],
+          multiples=[1, 1, 2, 1],
+          expected=[[[[1], [2], [1], [2]], [[3], [3]]],
+                    [[]],
+                    [[[4, 5], [4, 5]]]]),
+      dict(
+          descr='rank=4, ragged_rank=3, repeat axis 3',
+          rt_input=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]]],
+          multiples=[1, 1, 1, 2],
+          expected=[[[[1, 1], [2, 2]], [[3, 3]]], [[]], [[[4, 5, 4, 5]]]]),
+      dict(
+          descr='rank=4, ragged_rank=3, repeat all axes',
+          rt_input=[[[['a'], ['b']], [['c']]], [[]], [[['d', 'e']]]],
+          multiples=[5, 4, 3, 2],
+          expected=[[[[b'a']*2, [b'b']*2]*3, [[b'c']*2]*3]*4,
+                    [[]*3]*4,
+                    [[[b'd', b'e']*2]*3]*4]*5),
+      dict(
+          descr='rank=5, ragged_rank=4, repeat all axes',
+          rt_input=[[[[['a']]]]],
+          multiples=[6, 5, 4, 3, 2],
+          expected=[[[[[b'a']*2]*3]*4]*5]*6),
+      #=========================================================================
+      # multiple=0
+      #=========================================================================
+      dict(
+          descr='rank=4, ragged_rank=3, repeat axis 0 (multiple=0)',
+          rt_input=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]]],
+          multiples=[0, 1, 1, 1],
+          expected=[]),
+      dict(
+          descr='rank=4, ragged_rank=3, repeat axis 1 (multiple=0)',
+          rt_input=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]]],
+          multiples=[1, 0, 1, 1],
+          expected=[[], [], []]),
+      dict(
+          descr='rank=4, ragged_rank=3, repeat axis 2 (multiple=0)',
+          rt_input=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]]],
+          multiples=[1, 1, 0, 1],
+          expected=[[[], []], [[]], [[]]]),
+      dict(
+          descr='rank=4, ragged_rank=3, repeat axis 3 (multiple=0)',
+          rt_input=[[[[1], [2]], [[3]]], [[]], [[[4, 5]]]],
+          multiples=[1, 1, 1, 0],
+          expected=[[[[], []], [[]]], [[]], [[[]]]]),
+      #=========================================================================
+      # multiple=1
+      #=========================================================================
+      dict(
+          descr='rank=4, multiples=1 (no repeats)',
+          rt_input=[[[[1], [2]], [[3], [4]]], [[[5], [6]]]],
+          multiples=[1, 1, 1, 1],
+          expected=[[[[1], [2]], [[3], [4]]],
+                    [[[5], [6]]]]),
+
+  ])  # pyformat: disable
+  def testRaggedTile(self,
+                     descr,
+                     rt_input,
+                     multiples,
+                     expected,
+                     ragged_rank=None):
+    rt = ragged_factory_ops.constant(rt_input, ragged_rank)
+
+    expected_shape = [
+        None if dim is None else dim * multiple
+        for (dim, multiple) in zip(rt.shape.as_list(), multiples)
+    ]
+
+    # Test with both const & non-const multiples: ragged_tile has a few code
+    # paths that optimize the case where multiples[d] is known to be 1.
+    const_multiples = constant_op.constant(multiples, dtypes.int64)
+    non_const_multiples = array_ops.placeholder_with_default(
+        const_multiples, shape=[len(multiples)])
+
+    for multiples_tensor in (const_multiples, non_const_multiples):
+      tiled = ragged_array_ops.tile(rt, multiples_tensor)
+      self.assertEqual(tiled.ragged_rank, rt.ragged_rank)
+      self.assertEqual(tiled.shape.ndims, rt.shape.ndims)
+      if multiples_tensor is const_multiples:
+        self.assertEqual(tiled.shape.as_list(), expected_shape)
+      self.assertAllEqual(tiled, expected)
+
+  def testRaggedTileWithTensorInput(self):
+    # When the input is a `Tensor`, ragged_tile just delegates to tf.tile.
+    dt = constant_op.constant([[1, 2], [3, 4]])
+    tiled = ragged_array_ops.tile(dt, [3, 2])
+    expected = [[1, 2, 1, 2], [3, 4, 3, 4],
+                [1, 2, 1, 2], [3, 4, 3, 4],
+                [1, 2, 1, 2], [3, 4, 3, 4]]  # pyformat: disable
+    self.assertAllEqual(tiled, expected)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_to_sparse_op_test.py

@@ -0,0 +1,196 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged.to_sparse op."""
+
+from tensorflow.python.eager import context
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorToSparseOpTest(test_util.TensorFlowTestCase):
+
+  def testDocStringExample(self):
+    rt = ragged_factory_ops.constant([[1, 2, 3], [4], [], [5, 6]])
+    st = self.evaluate(rt.to_sparse())
+    self.assertAllEqual(st.indices,
+                        [[0, 0], [0, 1], [0, 2], [1, 0], [3, 0], [3, 1]])
+    self.assertAllEqual(st.values, [1, 2, 3, 4, 5, 6])
+    self.assertAllEqual(st.dense_shape, [4, 3])
+
+  def test2DRaggedTensorWithOneRaggedDimension(self):
+    rt = ragged_factory_ops.constant([['a', 'b'], ['c', 'd', 'e'], ['f'], [],
+                                      ['g']])
+    st = self.evaluate(rt.to_sparse())
+    self.assertAllEqual(
+        st.indices, [[0, 0], [0, 1], [1, 0], [1, 1], [1, 2], [2, 0], [4, 0]])
+    self.assertAllEqual(st.values, b'a b c d e f g'.split())
+    self.assertAllEqual(st.dense_shape, [5, 3])
+
+  def test3DRaggedTensorWithOneRaggedDimension(self):
+    rt = ragged_factory_ops.constant(
+        [[[1, 2], [3, 4]], [[5, 6], [7, 8], [9, 10]], [[11, 12]], [], [[13, 14]]
+        ],
+        ragged_rank=1)
+    st = self.evaluate(rt.to_sparse())
+    self.assertAllEqual(st.indices,
+                        [[0, 0, 0], [0, 0, 1], [0, 1, 0], [0, 1, 1], [1, 0, 0],
+                         [1, 0, 1], [1, 1, 0], [1, 1, 1], [1, 2, 0], [1, 2, 1],
+                         [2, 0, 0], [2, 0, 1], [4, 0, 0], [4, 0, 1]])
+    self.assertAllEqual(st.values,
+                        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
+    self.assertAllEqual(st.dense_shape, [5, 3, 2])
+
+  def test4DRaggedTensorWithOneRaggedDimension(self):
+    rt = ragged_factory_ops.constant(
+        [[[[1, 2], [3, 4]], [[5, 6], [7, 8]]], [], [[[9, 10], [11, 12]]]],
+        ragged_rank=1)
+    st = self.evaluate(rt.to_sparse())
+    self.assertAllEqual(st.values, [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
+    self.assertAllEqual(
+        st.indices,
+        [
+            [0, 0, 0, 0],  # index for value=1
+            [0, 0, 0, 1],  # index for value=2
+            [0, 0, 1, 0],  # index for value=3
+            [0, 0, 1, 1],  # index for value=4
+            [0, 1, 0, 0],  # index for value=5
+            [0, 1, 0, 1],  # index for value=6
+            [0, 1, 1, 0],  # index for value=7
+            [0, 1, 1, 1],  # index for value=8
+            [2, 0, 0, 0],  # index for value=9
+            [2, 0, 0, 1],  # index for value=10
+            [2, 0, 1, 0],  # index for value=11
+            [2, 0, 1, 1],  # index for value=12
+        ])
+    self.assertAllEqual(st.dense_shape, [3, 2, 2, 2])
+
+  def test4DRaggedTensorWithTwoRaggedDimensions(self):
+    rt = ragged_factory_ops.constant(
+        [[[[1, 2], [3, 4]], [[5, 6], [7, 8], [9, 10]]],
+         [[[11, 12]], [], [[13, 14]]], []],
+        ragged_rank=2)
+    st = self.evaluate(rt.to_sparse())
+    self.assertAllEqual(
+        st.indices,
+        [
+            [0, 0, 0, 0],  # index for value=1
+            [0, 0, 0, 1],  # index for value=2
+            [0, 0, 1, 0],  # index for value=3
+            [0, 0, 1, 1],  # index for value=4
+            [0, 1, 0, 0],  # index for value=5
+            [0, 1, 0, 1],  # index for value=6
+            [0, 1, 1, 0],  # index for value=7
+            [0, 1, 1, 1],  # index for value=8
+            [0, 1, 2, 0],  # index for value=9
+            [0, 1, 2, 1],  # index for value=10
+            [1, 0, 0, 0],  # index for value=11
+            [1, 0, 0, 1],  # index for value=12
+            [1, 2, 0, 0],  # index for value=13
+            [1, 2, 0, 1],  # index for value=14
+        ])
+    self.assertAllEqual(st.values,
+                        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
+    self.assertAllEqual(st.dense_shape, [3, 3, 3, 2])
+
+  def testShape(self):
+    rt = ragged_factory_ops.constant([[1, 2], [3, 4, 5], [6], [], [7]])
+    st = rt.to_sparse()
+    self.assertEqual(st.indices.shape.as_list(), [7, 2])
+    self.assertEqual(st.values.shape.as_list(), [7])
+    self.assertEqual(st.dense_shape.shape.as_list(), [2])
+
+    rt = ragged_factory_ops.constant([[[1, 2]], [], [[3, 4]], []],
+                                     ragged_rank=1)
+    st = rt.to_sparse()
+    self.assertEqual(st.indices.shape.as_list(), [4, 3])
+    self.assertEqual(st.values.shape.as_list(), [4])
+    self.assertEqual(st.dense_shape.shape.as_list(), [3])
+
+    rt = ragged_factory_ops.constant([[[1], [2, 3, 4, 5, 6, 7]], [[]]])
+    st = rt.to_sparse()
+    self.assertEqual(st.indices.shape.as_list(), [7, 3])
+    self.assertEqual(st.values.shape.as_list(), [7])
+    self.assertEqual(st.dense_shape.shape.as_list(), [3])
+
+  def testKernelErrors(self):
+    # An empty vector, defined using a placeholder to ensure that we can't
+    # determine that it's invalid at graph-construction time.
+    empty_vector = array_ops.placeholder_with_default(
+        array_ops.zeros([0], dtypes.int64), shape=None)
+
+    bad_rt1 = ragged_tensor.RaggedTensor.from_row_splits(
+        row_splits=[2, 3], values=[1, 2, 3], validate=False)
+    bad_split0 = r'First value of ragged splits must be 0.*'
+    with self.assertRaisesRegex(errors.InvalidArgumentError, bad_split0):
+      self.evaluate(bad_rt1.to_sparse())
+
+    bad_rt2 = ragged_tensor.RaggedTensor.from_row_splits(
+        row_splits=[0, 5], values=empty_vector, validate=False)
+    bad_rt3 = ragged_tensor.RaggedTensor.from_row_splits(
+        row_splits=[0, 1],
+        values=ragged_tensor.RaggedTensor.from_row_splits(
+            row_splits=[0, 5], values=empty_vector, validate=False),
+        validate=False)
+    split_mismatch1_error = r'Final value of ragged splits must match.*'
+    for rt in [bad_rt2, bad_rt3]:
+      with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                  split_mismatch1_error):
+        self.evaluate(rt.to_sparse())
+
+    bad_rt4 = ragged_tensor.RaggedTensor.from_row_splits(
+        row_splits=[0, 5],
+        values=ragged_tensor.RaggedTensor.from_row_splits(
+            row_splits=[0], values=empty_vector, validate=False),
+        validate=False)
+    split_mismatch2_error = r'Final value of ragged splits must match.*'
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                split_mismatch2_error):
+      self.evaluate(bad_rt4.to_sparse())
+
+    bad_rt5 = ragged_tensor.RaggedTensor.from_row_splits(
+        row_splits=empty_vector, values=[], validate=False)
+    empty_splits_error = (r'ragged splits may not be empty.*')
+    with self.assertRaisesRegex(errors.InvalidArgumentError,
+                                empty_splits_error):
+      self.evaluate(bad_rt5.to_sparse())
+
+  def testGradient(self):
+    if context.executing_eagerly():
+      return
+    # rt1.shape == rt2.shape == [2, (D2), (D3), 2].
+    rt1 = ragged_factory_ops.constant(
+        [[[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0]]]], ragged_rank=2)
+    rt2 = ragged_factory_ops.constant(
+        [[[[9.0, 8.0], [7.0, 6.0]], [[5.0, 4.0]]]], ragged_rank=2)
+    rt = ragged_functional_ops.map_flat_values(math_ops.add, rt1, rt2 * 2.0)
+    st = rt.to_sparse()
+
+    g1, g2 = gradients_impl.gradients(st.values,
+                                      [rt1.flat_values, rt2.flat_values])
+    self.assertAllEqual(g1, [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]])
+    self.assertAllEqual(g2, [[2.0, 2.0], [2.0, 2.0], [2.0, 2.0]])
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_to_tensor_op_test.py

@@ -0,0 +1,874 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged.to_tensor."""
+
+import random
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.client import session
+from tensorflow.python.eager import backprop
+from tensorflow.python.eager import context
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import tensor as tensor_lib
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import gradients_impl
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
+from tensorflow.python.platform import benchmark
+from tensorflow.python.platform import googletest
+from tensorflow.python.util import nest
+
+
+def make_placeholder(t):
+  return array_ops.placeholder_with_default(t, None)
+
+
+def rebuild_ragged_tensor_with_value_rowids(rt, feed_dict=None, sess=None):
+  """Returns a copy of `rt`, built using `from_value_rowids`.
+
+  This ensures that RaggedTensor._cached_value_rowids is populated, which
+  triggers a different code-path for converting ragged tensors to tensors.
+
+  If `feed_dict` and `sess` are specified, then build the new `RaggedTensor`
+  using placeholder tensors, and populate a feed dictionary that can be used
+  to feed the placeholders.
+
+  Args:
+    rt: The RaggedTensor to copy.
+    feed_dict: If specified, then build the new `RaggedTensor` using
+      placeholders, and populate this dict with entries to feed those
+      placeholders.
+    sess: A session used to evaluate tensors; required if feed_dict is
+      specified.
+
+  Returns:
+    A copy of `rt`, built using `from_value_rowids`.
+  """
+  if isinstance(rt, ragged_tensor.RaggedTensor):
+    values = rebuild_ragged_tensor_with_value_rowids(rt.values, feed_dict, sess)
+    rowids = rt.value_rowids()
+    nrows = rt.nrows()
+    if feed_dict is not None:
+      rowids_ph = make_placeholder(rowids)
+      nrows_ph = make_placeholder(nrows)
+      feed_dict[rowids_ph] = sess.run(rowids)
+      feed_dict[nrows_ph] = sess.run(nrows)
+      rowids, nrows = rowids_ph, nrows_ph
+    return ragged_tensor.RaggedTensor.from_value_rowids(values, rowids, nrows)
+  else:
+    if feed_dict is not None:
+      rt_ph = make_placeholder(rt)
+      feed_dict[rt_ph] = sess.run(rt)
+      rt = rt_ph
+    return rt
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedTensorToTensorOpTest(test_util.TensorFlowTestCase,
+                                 parameterized.TestCase):
+
+  def testDocStringExamples(self):
+    """Example from ragged_to_tensor.__doc__."""
+    rt = ragged_factory_ops.constant([[9, 8, 7], [], [6, 5], [4]])
+    dt = rt.to_tensor()
+    self.assertAllEqual(dt, [[9, 8, 7], [0, 0, 0], [6, 5, 0], [4, 0, 0]])
+
+  @parameterized.named_parameters(
+      # Simple 2D ragged tensors (with one ragged dimension)
+      {
+          'testcase_name': 'shape_2xN',
+          'rt_input': [[0, 1, 2], [], [3]],
+          'expected': [[0, 1, 2], [0, 0, 0], [3, 0, 0]]
+      },
+      {
+          'testcase_name': 'shape_2xN_default_0D',
+          'rt_input': [[0, 1, 2], [], [3]],
+          'default': 5,
+          'expected': [[0, 1, 2], [5, 5, 5], [3, 5, 5]]
+      },
+      {
+          'testcase_name': 'empty_first_row',
+          'rt_input': [[], [], [3, 4], []],
+          'expected': [[0, 0], [0, 0], [3, 4], [0, 0]]
+      },
+      {
+          'testcase_name': 'empty_last_row',
+          'rt_input': [[0, 1, 2], [], [3], []],
+          'expected': [[0, 1, 2], [0, 0, 0], [3, 0, 0], [0, 0, 0]]
+      },
+      {
+          'testcase_name': 'shape_4xN',
+          'rt_input': [[1, 2, 3], [], [4], [5, 6]],
+          'expected': [[1, 2, 3], [0, 0, 0], [4, 0, 0], [5, 6, 0]]
+      },
+      {
+          'testcase_name': 'shape_4xN_default_0D',
+          'rt_input': [[1, 2, 3], [], [4], [5, 6]],
+          'default': 9,
+          'expected': [[1, 2, 3], [9, 9, 9], [4, 9, 9], [5, 6, 9]]
+      },
+      {
+          'testcase_name': 'shape_2xN_already_dense',
+          'rt_input': [[6, 7, 8], [9, 10, 11]],
+          'expected': [[6, 7, 8], [9, 10, 11]],
+      },
+      {
+          'testcase_name': 'shape_2xN_string_already_dense',
+          'rt_input': [[b'a', b'b', b'c'],
+                       [b'd', b'e', b'antidisestablishmentarianism']],
+          'ragged_rank': 1,
+          'expected': [[b'a', b'b', b'c'],
+                       [b'd', b'e', b'antidisestablishmentarianism']],
+      },
+      # 3D ragged tensors with two ragged dimensions
+      {
+          'testcase_name': 'shape_4xNxM',
+          'rt_input': [[[1, 2], [], [3, 4]], [], [[5]], [[6, 7], [8]]],
+          'expected': [
+              [[1, 2], [0, 0], [3, 4]],  #
+              [[0, 0], [0, 0], [0, 0]],  #
+              [[5, 0], [0, 0], [0, 0]],  #
+              [[6, 7], [8, 0], [0, 0]],  #
+          ]
+      },
+      {
+          'testcase_name': 'shape_4xNxM_default_0D',
+          'rt_input': [[[1, 2], [], [3, 4]], [], [[5]], [[6, 7], [8]]],
+          'default': 9,
+          'expected': [
+              [[1, 2], [9, 9], [3, 4]],  #
+              [[9, 9], [9, 9], [9, 9]],  #
+              [[5, 9], [9, 9], [9, 9]],  #
+              [[6, 7], [8, 9], [9, 9]],  #
+          ]
+      },
+      {
+          'testcase_name': 'shape_1xNx1_default_0D',
+          'rt_input': [[[1], [2], [3]]],
+          'ragged_rank': 1,
+          'default': 0,
+          'expected': [[[1], [2], [3]]],
+      },
+      {
+          'testcase_name': 'shape_2xNx2_already_dense',
+          'rt_input': [[[6, 7], [8, 9], [10, 11]],
+                       [[12, 13], [14, 15], [16, 17]]],
+          'ragged_rank': 1,
+          'expected': [[[6, 7], [8, 9], [10, 11]],
+                       [[12, 13], [14, 15], [16, 17]]],
+      },
+      {
+          'testcase_name': 'shape_2xNx2_already_dense_default_1D',
+          'rt_input': [[[6, 7], [8, 9], [10, 11]],
+                       [[12, 13], [14, 15], [16, 17]]],
+          'ragged_rank': 1,
+          'default': [31, 32],
+          'expected': [[[6, 7], [8, 9], [10, 11]],
+                       [[12, 13], [14, 15], [16, 17]]],
+      },
+      {
+          'testcase_name': 'shape_2xNx2_string_already_dense',
+          'rt_input': [[[b'a', b'b'], [b'c', b'd'], [b'e', b'f']],
+                       [[b'g', b'jalapeno'], [b'kangaroo', b'llama'],
+                        [b'manzana', b'nectar']]],
+          'ragged_rank': 1,
+          'expected': [[[b'a', b'b'], [b'c', b'd'], [b'e', b'f']],
+                       [[b'g', b'jalapeno'], [b'kangaroo', b'llama'],
+                        [b'manzana', b'nectar']]],
+      },
+      # 3D ragged tensors with one ragged dimension
+      {
+          'testcase_name': 'shape_4xNx1_default_1D',
+          'rt_input': [[[1], [2], [3]], [], [[4]], [[5], [6]]],
+          'ragged_rank': 1,
+          'default': [9],
+          'expected': [[[1], [2], [3]],
+                       [[9], [9], [9]],
+                       [[4], [9], [9]],
+                       [[5], [6], [9]]]
+      },
+      {
+          'testcase_name': 'shape_2xNx2_default_0D',
+          'rt_input': [[[6, 7], [8, 9], [10, 11]],
+                       [[12, 13], [14, 15]]],
+          'ragged_rank': 1,
+          'default': 2,
+          'expected': [[[6, 7], [8, 9], [10, 11]],
+                       [[12, 13], [14, 15], [2, 2]]],
+      },
+      {
+          'testcase_name': 'shape_2xNx2_default_1D',
+          'rt_input': [[[6, 7], [8, 9], [10, 11]],
+                       [[12, 13], [14, 15]]],
+          'ragged_rank': 1,
+          'default': [2, 3],
+          'expected': [[[6, 7], [8, 9], [10, 11]],
+                       [[12, 13], [14, 15], [2, 3]]],
+      },
+      # 4D ragged tensors with 3 ragged dimensions
+      {
+          'testcase_name': 'shape_1xNxMxK_default_0D',
+          'rt_input': [[[[1], [2]], [], [[3]]]],
+          'default': 9,
+          'expected': [[[[1], [2]], [[9], [9]], [[3], [9]]]],
+      },
+      # Broadcast default
+      {
+          'testcase_name': 'shape_2xNx2x2_default_2x1',
+          'rt_input': [[[[1, 2], [3, 4]]], []],
+          'ragged_rank': 1,
+          'default': [[5], [6]],
+          'expected': [[[[1, 2], [3, 4]]],
+                       [[[5, 5], [6, 6]]]],
+      },
+      {
+          'testcase_name': 'shape_2xNx2x2_default_1x2',
+          'rt_input': [[[[1, 2], [3, 4]]], []],
+          'ragged_rank': 1,
+          'default': [[5, 6]],
+          'expected': [[[[1, 2], [3, 4]]],
+                       [[[5, 6], [5, 6]]]],
+      },
+      # Explicit shape
+      {
+          'testcase_name': 'shape_4xN_with_crop',
+          'rt_input': [[0, 1, 2, 3], [], [4], []],
+          'shape': [2, 3],
+          'expected': [[0, 1, 2], [0, 0, 0]],
+      },
+      {
+          'testcase_name': 'shape_2xN_with_pad',
+          'rt_input': [[1, 2], [3]],
+          'shape': [3, 3],
+          'expected': [[1, 2, 0], [3, 0, 0], [0, 0, 0]],
+      },
+      {
+          'testcase_name': 'shape_4xN_with_crop_and_pad',
+          'rt_input': [[0, 1, 2, 3], [], [4], []],
+          'shape': [2, 8],
+          'expected': [[0, 1, 2, 3, 0, 0, 0, 0],
+                       [0, 0, 0, 0, 0, 0, 0, 0]],
+      },
+      {
+          'testcase_name': 'shape_4xN_with_tuple_shape',
+          'rt_input': [[0, 1, 2, 3], [], [4], []],
+          'shape': (2, 3),
+          'expected': [[0, 1, 2], [0, 0, 0]],
+      },
+      {
+          'testcase_name': 'shape_4xN_with_tensorshape_shape',
+          'rt_input': [[0, 1, 2, 3], [], [4], []],
+          'shape': tensor_shape.TensorShape([2, 3]),
+          'expected': [[0, 1, 2], [0, 0, 0]],
+      },
+      {
+          'testcase_name': 'shape_4xN_with_partial_shape',
+          'rt_input': [[0, 1, 2, 3], [], [4], []],
+          'shape': tensor_shape.TensorShape([2, None]),
+          'expected': [[0, 1, 2, 3], [0, 0, 0, 0]],
+      },
+      # Empty tensors
+      {
+          'testcase_name': 'shape_0xN',
+          'rt_input': [],
+          'ragged_rank': 1,
+          'expected': [],
+          'expected_shape': [0, 0],
+      },
+      {
+          'testcase_name': 'shape_0xNxM',
+          'rt_input': [],
+          'ragged_rank': 2,
+          'expected': [],
+          'expected_shape': [0, 0, 0],
+      },
+      # {
+      #     'testcase_name': 'shape_0xNx2',
+      #     'rt_input': [],
+      #     'ragged_rank': 1,
+      #     'inner_shape': [2],
+      #     'expected': [],
+      #     'expected_shape': [0, 0, 2],
+      # },
+      {
+          'testcase_name': 'shape_2xN_empty',
+          'rt_input': [[], []],
+          'expected': [[], []],
+          'expected_shape': [2, 0],
+      },
+  )  # pyformat: disable
+  def testRaggedTensorToTensor(self,
+                               rt_input,
+                               expected,
+                               ragged_rank=None,
+                               inner_shape=None,
+                               default=None,
+                               shape=None,
+                               expected_shape=None):
+    rt1 = ragged_factory_ops.constant(
+        rt_input, ragged_rank=ragged_rank, inner_shape=inner_shape)
+    rt2 = rebuild_ragged_tensor_with_value_rowids(rt1)
+    for rt in [rt1, rt2]:
+      for use_placeholder in [False, True]:
+        if use_placeholder:
+          if default is not None:
+            default = make_placeholder(default)
+          rt = nest.map_structure(make_placeholder, rt, expand_composites=True)
+        dt = rt.to_tensor(default_value=default, shape=shape)
+        self.assertIsInstance(dt, tensor_lib.Tensor)
+        self.assertEqual(rt.dtype, dt.dtype)
+        if shape is not None:
+          self.assertTrue(dt.shape.is_compatible_with(shape))
+        else:
+          self.assertTrue(dt.shape.is_compatible_with(rt.shape))
+        if expected_shape is not None:
+          expected = np.ndarray(expected_shape, buffer=np.array(expected))
+        self.assertAllEqual(dt, expected)
+
+  @parameterized.parameters([
+      {
+          'rt_input': [[1, 2, 3]],
+          'default': 'a',
+          'error_type': TypeError,
+          'error': r'Expected int32|Cannot convert',
+      },
+      {
+          'rt_input': [[1, 2, 3]],
+          'default': [0],
+          'error': r'default_value\.shape=.* and '
+                   r'rt_input\.flat_values\.shape=.* are incompatible: '
+                   r'default_value\.rank = 1  must be less than '
+                   r'rt_input\.flat_values\.rank = 1'
+      },
+      {
+          'rt_input': [[[1, 2], [3, 4]], [[5, 6]]],
+          'ragged_rank': 1,
+          'default': [7, 8, 9],
+          'error': r'default_value\.shape.* and '
+                   r'rt_input\.flat_values\.shape.* are incompatible: '
+                   r'default_value\.shape\[-1\] = 3 but '
+                   r'rt_input\.flat_values\.shape\[-1\] = 2'
+      },
+      {
+          'rt_input': [[1, 2, 3]],
+          'shape': [3, 3, 3],
+          'error': r'rt_input\.shape and shape=\[.,.,.\] are incompatible: '
+                   r'rt_input\.rank = 2 but shape\.rank = 3'
+      },
+      {
+          'rt_input': [[[1, 2, 3]]],
+          'ragged_rank': 1,
+          'shape': [1, 1, 4],
+          'error': r'rt_input\.shape and shape=\[1,1,4\] are incompatible: '
+                   r'rt_input\.shape\[2\] = 3 but shape\[2\] = 4'
+      },
+  ])
+  def testError(self,
+                rt_input,
+                error,
+                error_type=(ValueError, errors.InvalidArgumentError),
+                default=None,
+                ragged_rank=None,
+                shape=None):
+
+    rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
+    with self.assertRaisesRegex(error_type, error):
+      self.evaluate(rt.to_tensor(default_value=default, shape=shape))
+    rt_placeholder = nest.map_structure(
+        make_placeholder, rt, expand_composites=True)
+    with self.assertRaisesRegex(error_type, error):
+      self.evaluate(
+          rt_placeholder.to_tensor(default_value=default, shape=shape))
+
+  def test_shape_limit_shape_is_tensor(self):
+    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
+    actual = input_data.to_tensor(
+        shape=constant_op.constant([2, 3], dtype=dtypes.int64))
+    self.assertAllEqual(actual, [[0, 1, 2], [0, 0, 0]])
+    self.assertEqual(actual.shape.as_list(), [2, 3])
+
+  def test_shape_limit_shape_is_tensor_unknown_rank(self):
+    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
+    actual = input_data.to_tensor(
+        shape=constant_op.constant(-1, dtype=dtypes.int64))
+    self.assertAllEqual(
+        actual, [[0, 1, 2, 3], [0, 0, 0, 0], [4, 0, 0, 0], [0, 0, 0, 0]])
+    self.assertTrue(actual.shape.is_compatible_with([4, 4]))
+
+  def test_shape_limit_shape_is_tensor_unknown_dim(self):
+    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
+    actual = input_data.to_tensor(
+        shape=constant_op.constant([2, -1], dtype=dtypes.int64))
+    self.assertAllEqual(actual, [[0, 1, 2, 3], [0, 0, 0, 0]])
+    self.assertTrue(actual.shape.is_compatible_with([2, None]))
+
+  def test_shape_limit_shape_is_tensor_int32(self):
+    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
+    actual = input_data.to_tensor(
+        shape=constant_op.constant([2, 3], dtype=dtypes.int32))
+    self.assertAllEqual(actual, [[0, 1, 2], [0, 0, 0]])
+    self.assertEqual(actual.shape.as_list(), [2, 3])
+
+  def test_shape_expand_first_dim(self):
+    input_data = ragged_factory_ops.constant([[0, 1, 2], [], [3]])
+    actual = input_data.to_tensor(shape=[4, 4])
+    self.assertAllEqual(
+        actual, [[0, 1, 2, 0], [0, 0, 0, 0], [3, 0, 0, 0], [0, 0, 0, 0]])
+    self.assertEqual(actual.shape.as_list(), [4, 4])
+
+  def test_value_transposed(self):
+    # Check that transposed data is not an issue.
+    my_value = array_ops.transpose(
+        constant_op.constant([[0, 1, 2, 3], [4, 5, 6, 7]]))
+    input_data = RaggedTensor.from_value_rowids(
+        values=my_value,
+        value_rowids=constant_op.constant([0, 1, 2, 3], dtype=dtypes.int64),
+        nrows=constant_op.constant(4, dtype=dtypes.int64),
+        validate=True)
+    self.assertAllEqual(input_data, [[[0, 4]], [[1, 5]], [[2, 6]], [[3, 7]]])
+
+  def test_broadcast_default(self):
+    # The dense dimension here is 2 x 2
+    input_data = ragged_factory_ops.constant([[[[1, 2], [3, 4]]], []],
+                                             ragged_rank=1)
+    # This placeholder has a 2 x 1 dimension.
+    default_value = make_placeholder([[5], [6]])
+    actual = input_data.to_tensor(default_value=default_value)
+    expected = [[[[1, 2], [3, 4]]], [[[5, 5], [6, 6]]]]
+    self.assertAllEqual(actual, expected)
+
+  def test_broadcast_default_no_placeholder(self):
+    input_data = ragged_factory_ops.constant([[[[1, 2], [3, 4]]], []],
+                                             ragged_rank=1)
+    # default_value has a 2 x 1 dimension.
+    default_value = constant_op.constant([[5], [6]], shape=None)
+    actual = input_data.to_tensor(default_value=default_value)
+    expected = [[[[1, 2], [3, 4]]], [[[5, 5], [6, 6]]]]
+    self.assertAllEqual(actual, expected)
+
+  def test_shape_expand_second_dim(self):
+    input_data = ragged_factory_ops.constant([[0, 1, 2], [], [3], []])
+    actual = input_data.to_tensor(shape=[3, 4])
+    self.assertAllEqual(actual, [[0, 1, 2, 0], [0, 0, 0, 0], [3, 0, 0, 0]])
+
+  @parameterized.parameters(
+      ([2, 3, 4], None, [2, 3, 4]),
+      ([2, 3, 4], [None, None, None], [2, 3, 4]),
+      ([2, 3, 4], [None, 3, None], [2, 3, 4]),
+      ([2, 3, 4], [None, 3, 4], [2, 3, 4]),
+      ([2, 3, 4], [2, 3, 4], [2, 3, 4]),
+      )
+  def test_preserve_shape_roundtrip(
+      self, input_shape, to_tensor_shape, expected_shape):
+    tensor = array_ops.zeros(input_shape)
+    ragged_from_tensor = RaggedTensor.from_tensor(tensor, ragged_rank=2)
+    recovered_tensor = ragged_from_tensor.to_tensor(shape=to_tensor_shape)
+    self.assertAllEqual(tensor.shape.as_list(), expected_shape)
+    self.assertAllEqual(ragged_from_tensor.shape.as_list(), expected_shape)
+    self.assertAllEqual(recovered_tensor.shape.as_list(), expected_shape)
+
+  def test_empty_tensor_with_shape(self):
+    input_data = RaggedTensor.from_value_rowids(
+        values=constant_op.constant([], dtype=dtypes.int64),
+        value_rowids=constant_op.constant([], dtype=dtypes.int64),
+        nrows=constant_op.constant(2, dtype=dtypes.int64),
+        validate=True)
+    actual = input_data.to_tensor(default_value=3, shape=[2, 3])
+    self.assertAllEqual(actual, [[3, 3, 3], [3, 3, 3]])
+
+  # pylint: disable=bad-whitespace
+  @parameterized.named_parameters([
+      dict(
+          testcase_name = '2d_default_shape',
+          shape         = None,
+          rt_value      = [[1, 2, 3], [4], [5, 6]],
+          rt_grad       = [[9, 8, 7], [6], [3, 2]],
+          default_value = 0,
+          default_grad  = sum([5, 4, 1]),
+          output_value  = [[1, 2, 3], [4, 0, 0], [5, 6, 0]],
+          output_grad   = [[9, 8, 7], [6, 5, 4], [3, 2, 1]]),
+      dict(
+          testcase_name = '2d_pad',
+          shape         = [4, 4],
+          rt_value      = [[1, 2, 3], [4], [5, 6]],
+          rt_grad       = [[9, 8, 7], [5], [1, 0]],
+          default_value = 0,
+          default_grad  = sum([6, 4, 3, 2, 1, 2, 3, 4, 5, 6]),
+          output_value  = [
+              [1, 2, 3, 0], [4, 0, 0, 0], [5, 6, 0, 0], [0, 0, 0, 0]],
+          output_grad   = [
+              [9, 8, 7, 6], [5, 4, 3, 2], [1, 0, 1, 2], [3, 4, 5, 6]]),
+      dict(
+          testcase_name = '2d_pad_and_crop',
+          shape         = [5, 3],
+          rt_value      = [[1, 2, 3], [4], [5, 6, 7, 8, 9], [8]],
+          rt_grad       = [[9, 8, 7], [6], [3, 2, 1, 0, 0], [2]],
+          default_value = 0,
+          default_grad  = sum([5, 4, 3, 4, 5, 6, 7]),
+          output_value  = [
+              [1, 2, 3], [4, 0, 0], [5, 6, 7], [8, 0, 0], [0, 0, 0]],
+          output_grad   = [
+              [9, 8, 7], [6, 5, 4], [3, 2, 1], [2, 3, 4], [5, 6, 7]]),
+      dict(
+          testcase_name = '3d_rrank_2',
+          shape         = [2, 2, 2],
+          rt_value      = [[[9, 8, 7], [6]], [[5, 4]]],
+          rt_grad       = [[[1, 2, 0], [3]], [[5, 6]]],
+          default_value = 3,
+          default_grad  = sum([4, 7, 8]),
+          output_value  = [[[9, 8], [6, 3]], [[5, 4], [3, 3]]],
+          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
+      dict(
+          testcase_name = '3d_rrank_1_with_0d_default',
+          ragged_rank   = 1,
+          shape         = [2, 2, 2],
+          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
+          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
+          default_value = 3,
+          default_grad  = sum([7, 8]),
+          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 3]]],
+          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
+      dict(
+          testcase_name = '3d_rrank_1_with_1d_default',
+          ragged_rank   = 1,
+          shape         = [2, 2, 2],
+          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
+          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
+          default_value = [3, 2],
+          default_grad  = [7, 8],
+          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 2]]],
+          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
+      dict(
+          testcase_name = '3d_rrank_1_with_1d_broadcast_default',
+          ragged_rank   = 1,
+          shape         = [2, 2, 2],
+          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
+          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
+          default_value = [3],
+          default_grad  = [7 + 8],
+          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 3]]],
+          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
+      dict(
+          testcase_name = '4d_rrank_1_with_2d_default',
+          ragged_rank   = 1,
+          shape         = [3, 3, 2, 1],
+          rt_value      = [[[[9], [8]], [[7], [6]]], [[[5], [4]]]],
+          rt_grad       = [[[[1], [2]], [[3], [4]]], [[[7], [8]]]],
+          default_value = [[3], [2]],
+          default_grad  = [[5 + 9 + 2 + 4 + 6 + 8], [6 + 1 + 3 + 5 + 7 + 9]],
+          output_value  = [[[[9], [8]], [[7], [6]], [[3], [2]]],
+                           [[[5], [4]], [[3], [2]], [[3], [2]]],
+                           [[[3], [2]], [[3], [2]], [[3], [2]]]],
+          output_grad   = [[[[1], [2]], [[3], [4]], [[5], [6]]],
+                           [[[7], [8]], [[9], [1]], [[2], [3]]],
+                           [[[4], [5]], [[6], [7]], [[8], [9]]]]),
+      dict(
+          testcase_name = '4d_rrank_1_with_with_0d_default',
+          ragged_rank   = 1,
+          shape         = [3, 3, 2, 1],
+          rt_value      = [[[[9], [8]], [[7], [6]]], [[[5], [4]]]],
+          rt_grad       = [[[[1], [2]], [[3], [4]]], [[[7], [8]]]],
+          default_value = 3,
+          default_grad  = 5 + 9 + 2 + 4 + 6 + 8 + 6 + 1 + 3 + 5 + 7 + 9,
+          output_value  = [[[[9], [8]], [[7], [6]], [[3], [3]]],
+                           [[[5], [4]], [[3], [3]], [[3], [3]]],
+                           [[[3], [3]], [[3], [3]], [[3], [3]]]],
+          output_grad   = [[[[1], [2]], [[3], [4]], [[5], [6]]],
+                           [[[7], [8]], [[9], [1]], [[2], [3]]],
+                           [[[4], [5]], [[6], [7]], [[8], [9]]]]),
+      dict(
+          testcase_name = 'zero_size',
+          shape         = [0, 0],
+          rt_value      = [[9, 8], [7, 6, 5], [4]],
+          rt_grad       = [[0, 0], [0, 0, 0], [0]],
+          default_value = 3,
+          default_grad  = 0,
+          output_value  = [],
+          output_grad   = [])
+  ])  # pyformat: disable
+  def test_gradient(self,
+                    shape,
+                    rt_value,
+                    rt_grad,
+                    default_value,
+                    default_grad,
+                    output_value,
+                    output_grad,
+                    ragged_rank=None):
+    """Tests that ragged_to_dense generates the right gradient.
+
+    Args:
+      shape: The `shape` arg for `ragged_to_dense`.
+      rt_value: The `rt_input` arg for `ragged_to_dense`.
+      rt_grad: The expected gradient for `rt_value`.  Corresponds 1:1 with
+        `rt_value`.
+      default_value: The `default_value` arg for `ragged_to_dense`.
+      default_grad: The expected gradient for `default_value`.  Corresponds 1:1
+        with `default_value`.
+      output_value: The expected output of `ragged_to_dense`.
+      output_grad: The gradient for the output (used to generate the gradients
+        `rt_grad` and `default_grad`).  Corresponds 1:1 with `output_value`.
+      ragged_rank: Ragged rank for `rt_value`.
+    """
+    rt_value = ragged_factory_ops.constant(
+        rt_value, dtype=dtypes.float32, ragged_rank=ragged_rank)
+    rt_grad = ragged_factory_ops.constant(
+        rt_grad, dtype=dtypes.float32, ragged_rank=ragged_rank)
+    default_value = constant_op.constant(default_value, dtype=dtypes.float32)
+    default_grad = constant_op.constant(default_grad, dtype=dtypes.float32)
+    output_value = constant_op.constant(
+        output_value, dtype=dtypes.float32, shape=shape)
+    output_grad = constant_op.constant(
+        output_grad, dtype=dtypes.float32, shape=shape)
+    shape = tensor_shape.as_shape(shape)
+
+    # There are different code paths for ragged_to_dense, depending on whether
+    # the RaggedTensor was created from row_splits or value_rowids.  Make sure
+    # that we test both.
+    for partition_type in ['row_splits', 'value_rowids']:
+      rt_val = self.rt_with_partition_type(rt_value, partition_type)
+      if context.executing_eagerly():
+        self._test_gradient_helper(rt_val, default_value, shape, output_grad,
+                                   output_value, rt_grad, default_grad)
+      else:
+        # There are different code paths when computing the gradient for
+        # default_value, depending on whether shape info is statically
+        # available; make sure that we test all code paths.
+        for shape_info in ['known', 'unknown_dims', 'unknown_rank']:
+          rt_val = self.wrap_in_placeholder(rt_val, shape_info)
+          default_val = self.wrap_in_placeholder(default_value, shape_info)
+          shape_val = self.wrap_in_placeholder(shape, shape_info)
+          self._test_gradient_helper(rt_val, default_val, shape_val,
+                                     output_grad, output_value, rt_grad,
+                                     default_grad)
+
+  def _test_gradient_helper(self, rt_val, default_val, shape_val, output_grad,
+                            expected_output_val, expected_rt_grad,
+                            expected_default_grad):
+    if context.executing_eagerly():
+      with backprop.GradientTape() as tape:
+        tape.watch([rt_val, default_val])
+        out = rt_val.to_tensor(default_val, shape=shape_val)
+        actual_rt_grad, actual_default_grad = tape.gradient(
+            out, (rt_val, default_val), output_gradients=output_grad)
+    else:
+      out = rt_val.to_tensor(default_val, shape=shape_val)
+      actual_rt_grad, actual_default_grad = gradients_impl.gradients(
+          ys=out, xs=(rt_val, default_val), grad_ys=output_grad)
+
+    self.assertAllClose(out, expected_output_val)
+    self.assertIsInstance(actual_rt_grad, RaggedTensor)
+    self.assertAllClose(actual_rt_grad, expected_rt_grad)
+    self.assertAllClose(actual_default_grad, expected_default_grad)
+
+  def rt_with_partition_type(self, rt, partition_type):
+    if isinstance(rt, tensor_lib.Tensor):
+      return rt
+    if partition_type == 'row_splits':
+      return rt
+    if partition_type == 'value_rowids':
+      return ragged_tensor.RaggedTensor.from_value_rowids(
+          self.rt_with_partition_type(rt.values, partition_type),
+          rt.value_rowids(), rt.nrows())
+    raise AssertionError('Unexpected partition_type %r' % partition_type)
+
+  def wrap_in_placeholder(self, arg, shape_info):
+    """Wraps `arg` in a placeholder to limit static shape info.
+
+    Args:
+      arg: The value to wrap.  A Tensor, RaggedTensor, or TensorShape.
+      shape_info: One of ['known', 'unknown_dims', 'unknown_rank'].
+
+    Returns:
+      * If shape_info is 'known': returns `arg`.
+      * If shape_info is 'unknown_dims': returns a placeholder wrapping `arg`
+        where the dimension sizes are unknown.  If `arg` is a TensorShape,
+        then convert it to a vector first.  If `arg` is a RaggedTensor, then
+        wrap the flat_values.
+      * If shape_info is 'unknown_rank': returns a placeholder wrapping `arg`
+        where the rank is unknown.  If `arg` is a TensorShape, then convert it
+        to a vector first.  If `arg` is a RaggedTensor, then wrap the
+        flat_values.
+    """
+    if shape_info == 'known':
+      return arg
+    if isinstance(arg, ragged_tensor.RaggedTensor):
+      return arg.with_flat_values(
+          self.wrap_in_placeholder(arg.flat_values, shape_info))
+    if isinstance(arg, tensor_shape.TensorShape):
+      if arg.ndims is None:
+        return arg
+      arg = constant_op.constant(arg.as_list())
+    if shape_info == 'unknown_rank':
+      return array_ops.placeholder_with_default(arg, None)
+    if shape_info == 'unknown_dims':
+      return array_ops.placeholder_with_default(arg, [None] * arg.shape.rank)
+    raise AssertionError('Unexpected shape_info %r' % shape_info)
+
+  def test_shape_is_list_including_tensor_element(self):
+    rt = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6]])
+    result = rt.to_tensor(shape=[2, constant_op.constant(2)])
+    self.assertAllEqual(result, [[1, 2], [4, 0]])
+
+  def test_row_splits_out_of_bounds(self):
+    from tensorflow.python.ops import gen_ragged_conversion_ops
+
+    with self.assertRaisesRegex(
+        errors.InvalidArgumentError,
+        'Row partition size.*exceeds the number of values',
+    ):
+      self.evaluate(
+          gen_ragged_conversion_ops.ragged_tensor_to_tensor(
+              shape=[1, 100],
+              values=[42.0],
+              default_value=0.0,
+              row_partition_tensors=[[0, 100]],
+              row_partition_types=['ROW_SPLITS'],
+          )
+      )
+
+
+class RaggedToDenseBenchmark(googletest.Benchmark):
+
+  # Configurations to test.  See `run_benchmark` for config param docs.
+  CONFIGS = [
+      {'shape': [10, 10]},
+      {'shape': [10, 1000]},
+      {'shape': [1000, 10]},
+      {'shape': [1000, 10], 'fill': [1, 0.95]},  # Mostly full.
+      {'shape': [1000, 10], 'fill': [1, 0.05]},  # Mostly empty.
+      {'shape': [1000, 10], 'dtype': dtypes.string},
+      {'shape': [1000, 10], 'dtype': dtypes.int64},
+      {'shape': [100, 100]},
+      {'shape': [50, 50, 32]},
+      {'shape': [100, 100, 100], 'min_iters': 100},
+      {'shape': [1000, 1000], 'min_iters': 100},
+      {'shape': [10, 10, 10, 10, 10]},
+      {'shape': [10, 10, 10, 10, 10], 'ragged_rank': 1},
+      {'shape': [10, 10, 10, 10, 10], 'ragged_rank': 2},
+      {'shape': [50, 50, 32], 'ragged_rank': 1, 'default_shape': [32]},
+      {'shape': [200, 50, 32], 'ragged_rank': 1, 'default_shape': [32]}
+  ]  # pyformat: disable
+
+  def run_benchmark(self,
+                    shape=(100, 100),
+                    ragged_rank=None,
+                    dtype=dtypes.float32,
+                    fill=None,
+                    default_shape=(),
+                    output_shape=None,
+                    min_iters=1000):
+    """Run a benchmark with the specified configuration parameters.
+
+    Args:
+      shape: Bounding box for the input ragged tensor.
+      ragged_rank: Ragged rank for the input ragged tensor.  Defaults to
+        `len(shape)-1`.
+      dtype: Data type for the input ragged tensor.
+      fill: How full each dimension should be (0-1).  Corresponds 1:1 with
+        `shape`.  Defaults to 0.8 for each dimension.
+      default_shape: Shape for the default (padding) value.
+      output_shape: Output shape -- ragged tensor will be padded or cropped to
+        this shape.
+      min_iters: Minimum iterations for benchmark.
+    """
+    if ragged_rank is None:
+      ragged_rank = len(shape) - 1
+    if fill is None:
+      fill = [0.8 for _ in shape]
+
+    # Build the inputs for the op.
+    rt_input = self._generateRaggedTensor(shape, ragged_rank, dtype, fill)
+    default_value = constant_op.constant(
+        self._generateRaggedTensor(default_shape, 0, dtype), dtype=dtype)
+
+    mbs = np.prod(shape) / (2**20)
+    with session.Session(config=benchmark.benchmark_config()) as sess:
+      extras = {
+          'shape': shape,
+          'ragged_rank': ragged_rank,
+          'dtype': dtype,
+          'fill': fill,
+          'default_shape': default_shape
+      }
+      rt = ragged_factory_ops.constant(rt_input, dtype, ragged_rank=ragged_rank)
+
+      # Inputs for with_splits:
+      splits_rt_placeholder = ragged_factory_ops.placeholder(
+          dtype, ragged_rank, shape[ragged_rank + 1:])
+      splits_feed_dict = {splits_rt_placeholder: sess.run(rt)}
+
+      # Inputs for with_rowids:
+      rowids_feed_dict = {}
+      rowids_rt_placeholder = rebuild_ragged_tensor_with_value_rowids(
+          rt, rowids_feed_dict, sess)
+
+      # Common arguments for benchmarks:
+      run_op_benchmark_kwargs = dict(
+          sess=sess,
+          store_memory_usage=True,
+          min_iters=min_iters,
+          burn_iters=max(5, min_iters // 10),
+          mbs=mbs,
+          extras=extras)
+
+      ragged_to_tensor_with_splits = splits_rt_placeholder.to_tensor(
+          default_value=default_value)
+      self.run_op_benchmark(
+          op_or_tensor=ragged_to_tensor_with_splits.op,
+          name='ragged_to_tensor_with_splits',
+          feed_dict=splits_feed_dict,
+          **run_op_benchmark_kwargs)
+
+      ragged_to_tensor_with_rowids = rowids_rt_placeholder.to_tensor(
+          default_value=default_value)
+      self.run_op_benchmark(
+          op_or_tensor=ragged_to_tensor_with_rowids.op,
+          name='ragged_to_tensor_with_rowids',
+          feed_dict=rowids_feed_dict,
+          **run_op_benchmark_kwargs)
+
+  def _generateRaggedTensor(self, shape, ragged_rank, dtype, fill=None, axis=0):
+    if axis == len(shape):
+      value = random.random()
+      if dtype == dtypes.string:
+        value = str(value)
+      if dtype.is_integer:
+        value = int(value * 1000)
+      return value
+    if axis == 0 or axis > ragged_rank:
+      slice_size = shape[axis]
+    else:
+      slice_size = (np.random.geometric(fill[axis], shape[axis]) == 1).sum()
+    return [
+        self._generateRaggedTensor(shape, ragged_rank, dtype, fill, axis + 1)
+        for _ in range(slice_size)
+    ]
+
+  def benchmark_ragged_to_dense(self):
+    random.seed(5)
+    for config in self.CONFIGS:
+      self.run_benchmark(**config)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_util.py

@@ -0,0 +1,138 @@
+# 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.
+# ==============================================================================
+"""Private convenience functions for RaggedTensors.
+
+None of these methods are exposed in the main "ragged" package.
+"""
+
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import check_ops
+from tensorflow.python.ops import control_flow_ops
+from tensorflow.python.ops import gen_ragged_math_ops
+from tensorflow.python.ops import math_ops
+
+
+def assert_splits_match(nested_splits_lists):
+  """Checks that the given splits lists are identical.
+
+  Performs static tests to ensure that the given splits lists are identical,
+  and returns a list of control dependency op tensors that check that they are
+  fully identical.
+
+  Args:
+    nested_splits_lists: A list of nested_splits_lists, where each split_list is
+      a list of `splits` tensors from a `RaggedTensor`, ordered from outermost
+      ragged dimension to innermost ragged dimension.
+
+  Returns:
+    A list of control dependency op tensors.
+  Raises:
+    ValueError: If the splits are not identical.
+  """
+  error_msg = "Inputs must have identical ragged splits"
+  for splits_list in nested_splits_lists:
+    if len(splits_list) != len(nested_splits_lists[0]):
+      raise ValueError(error_msg)
+  return [
+      check_ops.assert_equal(s1, s2, message=error_msg)
+      for splits_list in nested_splits_lists[1:]
+      for (s1, s2) in zip(nested_splits_lists[0], splits_list)
+  ]
+
+
+# Note: imported here to avoid circular dependency of array_ops.
+get_positive_axis = array_ops.get_positive_axis
+convert_to_int_tensor = array_ops.convert_to_int_tensor
+repeat = array_ops.repeat_with_axis
+
+
+def lengths_to_splits(lengths):
+  """Returns splits corresponding to the given lengths."""
+  return array_ops.concat([[0], math_ops.cumsum(lengths)], axis=-1)
+
+
+def repeat_ranges(params, splits, repeats):
+  """Repeats each range of `params` (as specified by `splits`) `repeats` times.
+
+  Let the `i`th range of `params` be defined as
+  `params[splits[i]:splits[i + 1]]`.  Then this function returns a tensor
+  containing range 0 repeated `repeats[0]` times, followed by range 1 repeated
+  `repeats[1]`, ..., followed by the last range repeated `repeats[-1]` times.
+
+  Args:
+    params: The `Tensor` whose values should be repeated.
+    splits: A splits tensor indicating the ranges of `params` that should be
+      repeated. Elements should be non-negative integers.
+    repeats: The number of times each range should be repeated. Supports
+      broadcasting from a scalar value. Elements should be non-negative
+      integers.
+
+  Returns:
+    A `Tensor` with the same rank and type as `params`.
+
+  #### Example:
+
+  >>> print(repeat_ranges(
+  ...     params=tf.constant(['a', 'b', 'c']),
+  ...     splits=tf.constant([0, 2, 3]),
+  ...     repeats=tf.constant(3)))
+  tf.Tensor([b'a' b'b' b'a' b'b' b'a' b'b' b'c' b'c' b'c'],
+      shape=(9,), dtype=string)
+  """
+  # Check if the input is valid
+  splits_checks = [
+      check_ops.assert_non_negative(
+          splits, message="Input argument 'splits' must be non-negative"
+      ),
+      check_ops.assert_integer(
+          splits,
+          message=(
+              "Input argument 'splits' must be integer, but got"
+              f" {splits.dtype} instead"
+          ),
+      ),
+  ]
+  repeats_checks = [
+      check_ops.assert_non_negative(
+          repeats, message="Input argument 'repeats' must be non-negative"
+      ),
+      check_ops.assert_integer(
+          repeats,
+          message=(
+              "Input argument 'repeats' must be integer, but got"
+              f" {repeats.dtype} instead"
+          ),
+      ),
+  ]
+  splits = control_flow_ops.with_dependencies(splits_checks, splits)
+  repeats = control_flow_ops.with_dependencies(repeats_checks, repeats)
+
+  # Divide `splits` into starts and limits, and repeat them `repeats` times.
+  if repeats.shape.ndims != 0:
+    repeated_starts = repeat(splits[:-1], repeats, axis=0)
+    repeated_limits = repeat(splits[1:], repeats, axis=0)
+  else:
+    # Optimization: we can just call repeat once, and then slice the result.
+    repeated_splits = repeat(splits, repeats, axis=0)
+    n_splits = array_ops.shape(repeated_splits, out_type=repeats.dtype)[0]
+    repeated_starts = repeated_splits[:n_splits - repeats]
+    repeated_limits = repeated_splits[repeats:]
+
+  # Get indices for each range from starts to limits, and use those to gather
+  # the values in the desired repetition pattern.
+  one = array_ops.ones((), repeated_starts.dtype)
+  offsets = gen_ragged_math_ops.ragged_range(
+      repeated_starts, repeated_limits, one)
+  return array_ops.gather(params, offsets.rt_dense_values)

tensorflow/python/ops/ragged/ragged_util_test.py

@@ -0,0 +1,241 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_util."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import test_util
+from tensorflow.python.framework.errors import InvalidArgumentError
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops.ragged import ragged_util
+from tensorflow.python.platform import googletest
+
+# Example 3d tensor for test cases.  Has shape [4, 2, 3].
+TENSOR_3D = [[[('%d%d%d' % (i, j, k)).encode('utf-8')
+               for k in range(3)]
+              for j in range(2)]
+             for i in range(4)]
+
+# Example 4d tensor for test cases.  Has shape [4, 2, 3, 5].
+TENSOR_4D = [[[[('%d%d%d%d' % (i, j, k, l)).encode('utf-8')
+                for l in range(5)]
+               for k in range(3)]
+              for j in range(2)]
+             for i in range(4)]
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedUtilTest(test_util.TensorFlowTestCase,
+                     parameterized.TestCase):
+
+  @parameterized.parameters([
+      # Docstring examples
+      dict(
+          data=['a', 'b', 'c'],
+          repeats=[3, 0, 2],
+          axis=0,
+          expected=[b'a', b'a', b'a', b'c', b'c']),
+      dict(
+          data=[[1, 2], [3, 4]],
+          repeats=[2, 3],
+          axis=0,
+          expected=[[1, 2], [1, 2], [3, 4], [3, 4], [3, 4]]),
+      dict(
+          data=[[1, 2], [3, 4]],
+          repeats=[2, 3],
+          axis=1,
+          expected=[[1, 1, 2, 2, 2], [3, 3, 4, 4, 4]]),
+
+      # Scalar repeats value
+      dict(
+          data=['a', 'b', 'c'],
+          repeats=2,
+          axis=0,
+          expected=[b'a', b'a', b'b', b'b', b'c', b'c']),
+      dict(
+          data=[[1, 2], [3, 4]],
+          repeats=2,
+          axis=0,
+          expected=[[1, 2], [1, 2], [3, 4], [3, 4]]),
+      dict(
+          data=[[1, 2], [3, 4]],
+          repeats=2,
+          axis=1,
+          expected=[[1, 1, 2, 2], [3, 3, 4, 4]]),
+
+      # data & repeats are broadcast to have at least one dimension,
+      # so these are all equivalent:
+      dict(data=3, repeats=4, axis=0, expected=[3, 3, 3, 3]),
+      dict(data=[3], repeats=4, axis=0, expected=[3, 3, 3, 3]),
+      dict(data=3, repeats=[4], axis=0, expected=[3, 3, 3, 3]),
+      dict(data=[3], repeats=[4], axis=0, expected=[3, 3, 3, 3]),
+      # Empty tensor
+      dict(data=[], repeats=[], axis=0, expected=[]),
+  ])
+  def testRepeat(self, data, repeats, expected, axis=None):
+    result = ragged_util.repeat(data, repeats, axis)
+    self.assertAllEqual(result, expected)
+
+  @parameterized.parameters([
+      dict(mode=mode, **args)
+      for mode in ['constant', 'dynamic', 'unknown_shape']
+      for args in [
+          # data & repeats are broadcast to have at least one dimension,
+          # so these are all equivalent:
+          dict(data=3, repeats=4, axis=0),
+          dict(data=[3], repeats=4, axis=0),
+          dict(data=3, repeats=[4], axis=0),
+          dict(data=[3], repeats=[4], axis=0),
+
+          # 1-dimensional data tensor.
+          dict(data=[], repeats=5, axis=0),
+          dict(data=[1, 2, 3], repeats=5, axis=0),
+          dict(data=[1, 2, 3], repeats=[3, 0, 2], axis=0),
+          dict(data=[1, 2, 3], repeats=[3, 0, 2], axis=-1),
+          dict(data=[b'a', b'b', b'c'], repeats=[3, 0, 2], axis=0),
+
+          # 2-dimensional data tensor.
+          dict(data=[[1, 2, 3], [4, 5, 6]], repeats=3, axis=0),
+          dict(data=[[1, 2, 3], [4, 5, 6]], repeats=3, axis=1),
+          dict(data=[[1, 2, 3], [4, 5, 6]], repeats=[3, 5], axis=0),
+          dict(data=[[1, 2, 3], [4, 5, 6]], repeats=[3, 5, 7], axis=1),
+
+          # 3-dimensional data tensor: shape=[4, 2, 3].
+          dict(data=TENSOR_3D, repeats=2, axis=0),
+          dict(data=TENSOR_3D, repeats=2, axis=1),
+          dict(data=TENSOR_3D, repeats=2, axis=2),
+          dict(data=TENSOR_3D, repeats=[2, 0, 4, 1], axis=0),
+          dict(data=TENSOR_3D, repeats=[3, 2], axis=1),
+          dict(data=TENSOR_3D, repeats=[1, 3, 1], axis=2),
+
+          # 4-dimensional data tensor: shape=[4, 2, 3, 5].
+          dict(data=TENSOR_4D, repeats=2, axis=0),
+          dict(data=TENSOR_4D, repeats=2, axis=1),
+          dict(data=TENSOR_4D, repeats=2, axis=2),
+          dict(data=TENSOR_4D, repeats=2, axis=3),
+          dict(data=TENSOR_4D, repeats=[2, 0, 4, 1], axis=0),
+          dict(data=TENSOR_4D, repeats=[3, 2], axis=1),
+          dict(data=TENSOR_4D, repeats=[1, 3, 1], axis=2),
+          dict(data=TENSOR_4D, repeats=[1, 3, 0, 0, 2], axis=3),
+      ]
+  ])
+  def testValuesMatchesNumpy(self, mode, data, repeats, axis):
+    # Exception: we can't handle negative axis if data.ndims is unknown.
+    if axis < 0 and mode == 'unknown_shape':
+      return
+
+    expected = np.repeat(data, repeats, axis)
+
+    if mode == 'constant':
+      data = constant_op.constant(data)
+      repeats = constant_op.constant(repeats)
+    elif mode == 'dynamic':
+      data = constant_op.constant(data)
+      repeats = constant_op.constant(repeats)
+      data = array_ops.placeholder_with_default(data, data.shape)
+      repeats = array_ops.placeholder_with_default(repeats, repeats.shape)
+    elif mode == 'unknown_shape':
+      data = array_ops.placeholder_with_default(data, None)
+      repeats = array_ops.placeholder_with_default(repeats, None)
+
+    result = ragged_util.repeat(data, repeats, axis)
+    self.assertAllEqual(result, expected)
+
+  @parameterized.parameters([
+      dict(
+          descr='axis >= rank(data)',
+          mode='dynamic',
+          data=[1, 2, 3],
+          repeats=[3, 0, 2],
+          axis=1,
+          error='axis=1 out of bounds: expected -1<=axis<1'),
+      dict(
+          descr='axis < -rank(data)',
+          mode='dynamic',
+          data=[1, 2, 3],
+          repeats=[3, 0, 2],
+          axis=-2,
+          error='axis=-2 out of bounds: expected -1<=axis<1'),
+      dict(
+          descr='len(repeats) != data.shape[axis]',
+          mode='dynamic',
+          data=[[1, 2, 3], [4, 5, 6]],
+          repeats=[2, 3],
+          axis=1,
+          error='Dimensions 3 and 2 are not compatible'),
+      dict(
+          descr='rank(repeats) > 1',
+          mode='dynamic',
+          data=[[1, 2, 3], [4, 5, 6]],
+          repeats=[[3], [5]],
+          axis=1,
+          error=r'Shape \(2, 1\) must have rank at most 1'),
+      dict(
+          descr='non-integer axis',
+          mode='constant',
+          data=[1, 2, 3],
+          repeats=2,
+          axis='foo',
+          exception=TypeError,
+          error='`axis` must be an int'),
+  ])
+  def testError(self,
+                descr,
+                mode,
+                data,
+                repeats,
+                axis,
+                exception=ValueError,
+                error=None):
+    # Make sure that this is also an error case for numpy.
+    with self.assertRaises(exception):
+      np.repeat(data, repeats, axis)
+
+    if mode == 'constant':
+      data = constant_op.constant(data)
+      repeats = constant_op.constant(repeats)
+    elif mode == 'dynamic':
+      data = constant_op.constant(data)
+      repeats = constant_op.constant(repeats)
+      data = array_ops.placeholder_with_default(data, data.shape)
+      repeats = array_ops.placeholder_with_default(repeats, repeats.shape)
+    elif mode == 'unknown_shape':
+      data = array_ops.placeholder_with_default(data, None)
+      repeats = array_ops.placeholder_with_default(repeats, None)
+
+    with self.assertRaisesRegex(exception, error):
+      ragged_util.repeat(data, repeats, axis)
+
+  @parameterized.parameters([
+      dict(
+          params=[1, 2, 3],
+          splits=[-1, -3],
+          repeats=2,
+          exception=InvalidArgumentError,
+      ),
+      dict(params=[1, 2, 3], splits=[1, 2], repeats=0.5, exception=TypeError),
+  ])
+  def testInputCheck(self, params, splits, repeats, exception):
+    params = constant_op.constant(params)
+    splits = constant_op.constant(splits)
+    repeats = constant_op.constant(repeats)
+    with self.assertRaises(exception):
+      ragged_util.repeat_ranges(params, splits, repeats)
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/ragged_where_op.py

@@ -0,0 +1,259 @@
+# 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.
+# ==============================================================================
+"""where operation for RaggedTensors."""
+
+import typing
+
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_concat_ops
+from tensorflow.python.ops.ragged import ragged_functional_ops
+from tensorflow.python.ops.ragged import ragged_gather_ops
+from tensorflow.python.ops.ragged import ragged_tensor
+from tensorflow.python.ops.ragged import ragged_tensor_shape
+from tensorflow.python.util import dispatch
+
+
+@dispatch.dispatch_for_api(array_ops.where_v2)
+def where_v2(condition: ragged_tensor.RaggedOrDense,
+             x: typing.Optional[ragged_tensor.RaggedOrDense] = None,
+             y: typing.Optional[ragged_tensor.RaggedOrDense] = None,
+             name=None):
+  """Return the elements where `condition` is `True`.
+
+  : If both `x` and `y` are None: Retrieve indices of true elements.
+
+    Returns the coordinates of true elements of `condition`. The coordinates
+    are returned in a 2-D tensor with shape
+    `[num_true_values, dim_size(condition)]`, where `result[i]` is the
+    coordinates of the `i`th true value (in row-major order).
+
+  : If both `x` and `y` are non-`None`: Multiplex between `x` and `y`.
+
+    Choose an output shape  from the shapes of `condition`, `x`, and `y` that
+    all three shapes are broadcastable to; and then use the broadcasted
+    `condition` tensor as a mask that chooses whether the corredsponding element
+    in the output should be taken from `x` (if `condition` is true) or `y` (if
+    `condition` is false).
+
+  >>> # Example: retrieve indices of true elements
+  >>> tf.where(tf.ragged.constant([[True, False], [True]]))
+  <tf.Tensor: shape=(2, 2), dtype=int64, numpy= array([[0, 0], [1, 0]])>
+
+  >>> # Example: multiplex between `x` and `y`
+  >>> tf.where(tf.ragged.constant([[True, False], [True, False, True]]),
+  ...          tf.ragged.constant([['A', 'B'], ['C', 'D', 'E']]),
+  ...          tf.ragged.constant([['a', 'b'], ['c', 'd', 'e']]))
+  <tf.RaggedTensor [[b'A', b'b'], [b'C', b'd', b'E']]>
+
+  Args:
+    condition: A potentially ragged tensor of type `bool`
+    x: A potentially ragged tensor (optional).
+    y: A potentially ragged tensor (optional).  Must be specified if `x` is
+      specified.  Must have the same rank and type as `x`.
+    name: A name of the operation (optional).
+
+  Returns:
+    : If both `x` and `y` are `None`:
+      A `Tensor` with shape `(num_true, rank(condition))`.
+    : Otherwise:
+      A potentially ragged tensor with the same type as `x` and `y`, and whose
+      shape is broadcast-compatible with `x`, `y`, and `condition`.
+
+  Raises:
+    ValueError: When exactly one of `x` or `y` is non-`None`; or when
+      `condition`, `x`, and `y` have incompatible shapes.
+  """
+  if (x is None) != (y is None):
+    raise ValueError('x and y must be either both None or both non-None')
+
+  with ops.name_scope('RaggedWhere', name, [condition, x, y]):
+    condition = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        condition, name='condition')
+    if x is None:
+      return _coordinate_where(condition)
+    else:
+      x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x, name='x')
+      y = ragged_tensor.convert_to_tensor_or_ragged_tensor(y, name='y')
+      condition, x, y = ragged_tensor.match_row_splits_dtypes(condition, x, y)
+      return _elementwise_where_v2(condition, x, y)
+
+
+@dispatch.dispatch_for_api(array_ops.where)
+def where(condition: ragged_tensor.RaggedOrDense,
+          x: typing.Optional[ragged_tensor.RaggedOrDense] = None,
+          y: typing.Optional[ragged_tensor.RaggedOrDense] = None,
+          name=None):
+  """Return the elements, either from `x` or `y`, depending on the `condition`.
+
+  : If both `x` and `y` are `None`:
+    Returns the coordinates of true elements of `condition`. The coordinates
+    are returned in a 2-D tensor with shape
+    `[num_true_values, dim_size(condition)]`, where `result[i]` is the
+    coordinates of the `i`th true value (in row-major order).
+
+  : If both `x` and `y` are non-`None`:
+    Returns a tensor formed by selecting values from `x` where condition is
+    true, and from `y` when condition is false.  In particular:
+
+    : If `condition`, `x`, and `y` all have the same shape:
+
+      * `result[i1...iN] = x[i1...iN]` if `condition[i1...iN]` is true.
+      * `result[i1...iN] = y[i1...iN]` if `condition[i1...iN]` is false.
+
+    : Otherwise:
+
+      * `condition` must be a vector.
+      * `x` and `y` must have the same number of dimensions.
+      * The outermost dimensions of `condition`, `x`, and `y` must all have the
+        same size.
+      * `result[i] = x[i]` if `condition[i]` is true.
+      * `result[i] = y[i]` if `condition[i]` is false.
+
+  Args:
+    condition: A potentially ragged tensor of type `bool`
+    x: A potentially ragged tensor (optional).
+    y: A potentially ragged tensor (optional).  Must be specified if `x` is
+      specified.  Must have the same rank and type as `x`.
+    name: A name of the operation (optional)
+
+  Returns:
+    : If both `x` and `y` are `None`:
+      A `Tensor` with shape `(num_true, dim_size(condition))`.
+    : Otherwise:
+      A potentially ragged tensor with the same type, rank, and outermost
+      dimension size as `x` and `y`.
+      `result.ragged_rank = max(x.ragged_rank, y.ragged_rank)`.
+
+  Raises:
+    ValueError: When exactly one of `x` or `y` is non-`None`; or when
+      `condition`, `x`, and `y` have incompatible shapes.
+
+  #### Examples:
+
+  >>> # Coordinates where condition is true.
+  >>> condition = tf.ragged.constant([[True, False, True], [False, True]])
+  >>> print(where(condition))
+  tf.Tensor( [[0 0] [0 2] [1 1]], shape=(3, 2), dtype=int64)
+
+  >>> # Elementwise selection between x and y, based on condition.
+  >>> condition = tf.ragged.constant([[True, False, True], [False, True]])
+  >>> x = tf.ragged.constant([['A', 'B', 'C'], ['D', 'E']])
+  >>> y = tf.ragged.constant([['a', 'b', 'c'], ['d', 'e']])
+  >>> print(where(condition, x, y))
+  <tf.RaggedTensor [[b'A', b'b', b'C'], [b'd', b'E']]>
+
+  >>> # Row selection between x and y, based on condition.
+  >>> condition = [True, False]
+  >>> x = tf.ragged.constant([['A', 'B', 'C'], ['D', 'E']])
+  >>> y = tf.ragged.constant([['a', 'b', 'c'], ['d', 'e']])
+  >>> print(where(condition, x, y))
+  <tf.RaggedTensor [[b'A', b'B', b'C'], [b'd', b'e']]>
+  """
+  if (x is None) != (y is None):
+    raise ValueError('x and y must be either both None or both non-None')
+  with ops.name_scope('RaggedWhere', name, [condition, x, y]):
+    condition = ragged_tensor.convert_to_tensor_or_ragged_tensor(
+        condition, name='condition')
+    if x is None:
+      return _coordinate_where(condition)
+    else:
+      x = ragged_tensor.convert_to_tensor_or_ragged_tensor(x, name='x')
+      y = ragged_tensor.convert_to_tensor_or_ragged_tensor(y, name='y')
+      condition, x, y = ragged_tensor.match_row_splits_dtypes(condition, x, y)
+      return _elementwise_where(condition, x, y)
+
+
+def _elementwise_where(condition, x, y):
+  """Ragged version of tf.where(condition, x, y)."""
+  condition_is_ragged = isinstance(condition, ragged_tensor.RaggedTensor)
+  x_is_ragged = isinstance(x, ragged_tensor.RaggedTensor)
+  y_is_ragged = isinstance(y, ragged_tensor.RaggedTensor)
+
+  if not (condition_is_ragged or x_is_ragged or y_is_ragged):
+    return array_ops.where(condition, x, y)
+
+  elif condition_is_ragged and x_is_ragged and y_is_ragged:
+    return ragged_functional_ops.map_flat_values(array_ops.where, condition, x,
+                                                 y)
+  elif not condition_is_ragged:
+    # Concatenate x and y, and then use `gather` to assemble the selected rows.
+    condition.shape.assert_has_rank(1)
+    x_and_y = ragged_concat_ops.concat([x, y], axis=0)
+    x_nrows = _nrows(x, out_type=x_and_y.row_splits.dtype)
+    y_nrows = _nrows(y, out_type=x_and_y.row_splits.dtype)
+    indices = array_ops.where(condition, math_ops.range(x_nrows),
+                              x_nrows + math_ops.range(y_nrows))
+    return ragged_gather_ops.gather(x_and_y, indices)
+
+  else:
+    raise ValueError('Input shapes do not match.')
+
+
+def _elementwise_where_v2(condition, x, y):
+  """Ragged version of tf.where_v2(condition, x, y)."""
+  # Broadcast x, y, and condition to have the same shape.
+  if not (condition.shape.is_fully_defined() and x.shape.is_fully_defined() and
+          y.shape.is_fully_defined() and x.shape == y.shape and
+          condition.shape == x.shape):
+    shape_c = ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(
+        condition)
+    shape_x = ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(x)
+    shape_y = ragged_tensor_shape.RaggedTensorDynamicShape.from_tensor(y)
+    shape = ragged_tensor_shape.broadcast_dynamic_shape(
+        shape_c, ragged_tensor_shape.broadcast_dynamic_shape(shape_x, shape_y))
+    condition = ragged_tensor_shape.broadcast_to(condition, shape)
+    x = ragged_tensor_shape.broadcast_to(x, shape)
+    y = ragged_tensor_shape.broadcast_to(y, shape)
+
+  condition_is_ragged = isinstance(condition, ragged_tensor.RaggedTensor)
+  x_is_ragged = isinstance(x, ragged_tensor.RaggedTensor)
+  y_is_ragged = isinstance(y, ragged_tensor.RaggedTensor)
+  if not (condition_is_ragged or x_is_ragged or y_is_ragged):
+    return array_ops.where_v2(condition, x, y)
+
+  return ragged_functional_ops.map_flat_values(array_ops.where_v2, condition, x,
+                                               y)
+
+
+def _coordinate_where(condition):
+  """Ragged version of tf.where(condition)."""
+  if not isinstance(condition, ragged_tensor.RaggedTensor):
+    return array_ops.where(condition)
+
+  # The coordinate for each `true` value in condition.values.
+  selected_coords = _coordinate_where(condition.values)
+
+  # Convert the first index in each coordinate to a row index and column index.
+  condition = condition.with_row_splits_dtype(selected_coords.dtype)
+  first_index = selected_coords[:, 0]
+  selected_rows = array_ops.gather(condition.value_rowids(), first_index)
+  selected_row_starts = array_ops.gather(condition.row_splits, selected_rows)
+  selected_cols = first_index - selected_row_starts
+
+  # Assemble the row & column index with the indices for inner dimensions.
+  return array_ops.concat([
+      array_ops.expand_dims(selected_rows, 1),
+      array_ops.expand_dims(selected_cols, 1), selected_coords[:, 1:]
+  ],
+                          axis=1)
+
+
+def _nrows(rt_input, out_type):
+  if isinstance(rt_input, ragged_tensor.RaggedTensor):
+    return rt_input.nrows(out_type=out_type)
+  else:
+    return array_ops.shape(rt_input, out_type=out_type)[0]

tensorflow/python/ops/ragged/ragged_where_op_test.py

@@ -0,0 +1,375 @@
+# 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.
+# ==============================================================================
+"""Tests for ragged_array_ops.where."""
+
+from absl.testing import parameterized
+import numpy as np
+
+from tensorflow.python.framework import errors
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_where_op
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedWhereV1OpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Docstring Examples
+      #=========================================================================
+      dict(  # shape=[D1, (D2)]
+          condition=ragged_factory_ops.constant_value(
+              [[True, False, True], [False, True]]),
+          expected=[[0, 0], [0, 2], [1, 1]]),
+      dict(  # shape=[D1, (D2)]
+          condition=ragged_factory_ops.constant_value(
+              [[True, False, True], [False, True]]),
+          x=ragged_factory_ops.constant_value(
+              [['A', 'B', 'C'], ['D', 'E']]),
+          y=ragged_factory_ops.constant_value(
+              [['a', 'b', 'c'], ['d', 'e']]),
+          expected=ragged_factory_ops.constant_value(
+              [[b'A', b'b', b'C'], [b'd', b'E']])),
+      dict(  # shape=[D1, (D2)]
+          condition=ragged_factory_ops.constant_value([True, False]),
+          x=ragged_factory_ops.constant_value([['A', 'B', 'C'], ['D', 'E']]),
+          y=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d', 'e']]),
+          expected=ragged_factory_ops.constant_value(
+              [[b'A', b'B', b'C'], [b'd', b'e']])),
+      #=========================================================================
+      # Coordinate-retrieval mode
+      #=========================================================================
+      dict(  # shape=[D1]
+          condition=[True, False, True, False, True],
+          expected=[[0], [2], [4]]),
+      dict(  # shape=[D1, D2]
+          condition=[[True, False], [False, True]],
+          expected=[[0, 0], [1, 1]]),
+      dict(  # shape=[D1, (D2)]
+          condition=ragged_factory_ops.constant_value(
+              [[True, False, True], [False, True]]),
+          expected=[[0, 0], [0, 2], [1, 1]]),
+      dict(  # shape=[D1, (D2), (D3)]
+          condition=ragged_factory_ops.constant_value([
+              [[True, False, True], [False, True]],
+              [[True], [], [False], [False, True, False]]
+          ]),
+          expected=[[0, 0, 0], [0, 0, 2], [0, 1, 1],
+                    [1, 0, 0], [1, 3, 1]]),
+      dict(  # shape=[D1, (D2), D3]
+          condition=ragged_factory_ops.constant_value([
+              [[True, False], [False, True]],
+              [[True, False], [False, False], [True, False], [False, True]]
+          ], ragged_rank=1),
+          expected=[[0, 0, 0], [0, 1, 1],
+                    [1, 0, 0], [1, 2, 0], [1, 3, 1]]),
+      dict(  # shape=[D1, (D2), (D3), (D4)]
+          condition=ragged_factory_ops.constant_value([
+              [[[], [True]]],
+              [[[True, False, True], [False, True]],
+               [[True], [], [False], [False, True, False]]]
+          ]),
+          expected=[[0, 0, 1, 0],
+                    [1, 0, 0, 0], [1, 0, 0, 2], [1, 0, 1, 1],
+                    [1, 1, 0, 0], [1, 1, 3, 1]]),
+
+      #=========================================================================
+      # Elementwise value-selection mode
+      #=========================================================================
+      dict(  # shape=[]
+          condition=True, x='A', y='a', expected=b'A'),
+      dict(  # shape=[]
+          condition=False, x='A', y='a', expected=b'a'),
+      dict(  # shape=[D1]
+          condition=[True, False, True],
+          x=['A', 'B', 'C'],
+          y=['a', 'b', 'c'],
+          expected=[b'A', b'b', b'C']),
+      dict(  # shape=[D1, D2]
+          condition=[[True, False], [False, True]],
+          x=[['A', 'B'], ['D', 'E']],
+          y=[['a', 'b'], ['d', 'e']],
+          expected=[[b'A', b'b'], [b'd', b'E']]),
+      dict(  # shape=[D1, (D2)]
+          condition=ragged_factory_ops.constant_value(
+              [[True, False, True], [False, True]]),
+          x=ragged_factory_ops.constant_value([['A', 'B', 'C'], ['D', 'E']]),
+          y=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d', 'e']]),
+          expected=ragged_factory_ops.constant_value(
+              [[b'A', b'b', b'C'], [b'd', b'E']])),
+      dict(  # shape=[D1, (D2), D3]
+          condition=ragged_factory_ops.constant_value([
+              [[True, False], [False, True]],
+              [[True, False], [False, False], [True, False], [False, True]]
+          ], ragged_rank=1),
+          x=ragged_factory_ops.constant_value([
+              [['A', 'B'], ['C', 'D']],
+              [['E', 'F'], ['G', 'H'], ['I', 'J'], ['K', 'L']]
+          ], ragged_rank=1),
+          y=ragged_factory_ops.constant_value([
+              [['a', 'b'], ['c', 'd']],
+              [['e', 'f'], ['g', 'h'], ['i', 'j'], ['k', 'l']]
+          ], ragged_rank=1),
+          expected=ragged_factory_ops.constant_value([
+              [[b'A', b'b'], [b'c', b'D']],
+              [[b'E', b'f'], [b'g', b'h'], [b'I', b'j'], [b'k', b'L']]
+          ], ragged_rank=1)),
+      dict(  # shape=[D1, (D2), (D3), (D4)]
+          condition=ragged_factory_ops.constant_value([
+              [[[], [True]]],
+              [[[True, False, True], [False, True]],
+               [[True], [], [False], [False, True, False]]]
+          ]),
+          x=ragged_factory_ops.constant_value([
+              [[[], ['A']]],
+              [[['B', 'C', 'D'], ['E', 'F']],
+               [['G'], [], ['H'], ['I', 'J', 'K']]]
+          ]),
+          y=ragged_factory_ops.constant_value([
+              [[[], ['a']]],
+              [[['b', 'c', 'd'], ['e', 'f']],
+               [['g'], [], ['h'], ['i', 'j', 'k']]]
+          ]),
+          expected=ragged_factory_ops.constant_value([
+              [[[], [b'A']]],
+              [[[b'B', b'c', b'D'], [b'e', b'F']],
+               [[b'G'], [], [b'h'], [b'i', b'J', b'k']]]
+          ])),
+
+      #=========================================================================
+      # Elementwise row-selection mode
+      #=========================================================================
+      dict(  # x.shape=[D1, D2], y.shape=[D1, D2]
+          condition=[True, False, True],
+          x=[['A', 'B'], ['C', 'D'], ['E', 'F']],
+          y=[['a', 'b'], ['c', 'd'], ['e', 'f']],
+          expected=[[b'A', b'B'], [b'c', b'd'], [b'E', b'F']]),
+      dict(  # x.shape=[D1, D2], y.shape=[D1, (D2)]
+          condition=[True, False, True],
+          x=[['A', 'B'], ['C', 'D'], ['E', 'F']],
+          y=ragged_factory_ops.constant_value(
+              [['a', 'b'], ['c'], ['d', 'e']]),
+          expected=ragged_factory_ops.constant_value(
+              [[b'A', b'B'], [b'c'], [b'E', b'F']])),
+      dict(  # x.shape=[D1, (D2)], y.shape=[D1, (D2)]
+          condition=[True, False, True],
+          x=ragged_factory_ops.constant_value(
+              [['A', 'B', 'C'], ['D', 'E'], ['F', 'G']]),
+          y=ragged_factory_ops.constant_value(
+              [['a', 'b'], ['c'], ['d', 'e']]),
+          expected=ragged_factory_ops.constant_value(
+              [[b'A', b'B', b'C'], [b'c'], [b'F', b'G']])),
+      dict(  # shape=[D1, (D2), (D3), (D4)]
+          condition=ragged_factory_ops.constant_value([True, False]),
+          x=ragged_factory_ops.constant_value([
+              [[[], ['A']]],
+              [[['B', 'C', 'D'], ['E', 'F']],
+               [['G'], [], ['H'], ['I', 'J', 'K']]]
+          ]),
+          y=ragged_factory_ops.constant_value([[[['a']]], [[['b']]]]),
+          expected=ragged_factory_ops.constant_value(
+              [[[[], [b'A']]], [[[b'b']]]])),
+  ])   # pyformat: disable
+  def testRaggedWhere(self, condition, expected, x=None, y=None):
+    result = ragged_where_op.where(condition, x, y)
+    self.assertAllEqual(result, expected)
+
+  @parameterized.parameters([
+      dict(
+          condition=[True, False],
+          x=[1, 2],
+          error=ValueError,
+          message='x and y must be either both None or both non-None'),
+      dict(
+          condition=ragged_factory_ops.constant_value([[True, False, True],
+                                                       [False, True]]),
+          x=ragged_factory_ops.constant_value([['A', 'B', 'C'], ['D', 'E']]),
+          y=[['a', 'b'], ['d', 'e']],
+          error=ValueError,
+          message='Input shapes do not match.'),
+  ])
+  def testRaggedWhereErrors(self, condition, error, message, x=None, y=None):
+    with self.assertRaisesRegex(error, message):
+      ragged_where_op.where(condition, x, y)
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class RaggedWhereV2OpTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  @parameterized.parameters([
+      #=========================================================================
+      # Coordinate-retrieval mode
+      #=========================================================================
+      dict(  # shape=[D1]
+          condition=[True, False, True, False, True],
+          expected=[[0], [2], [4]]),
+      dict(  # shape=[D1, D2]
+          condition=[[True, False], [False, True]],
+          expected=[[0, 0], [1, 1]]),
+      dict(  # shape=[D1, (D2)]
+          condition=ragged_factory_ops.constant_value(
+              [[True, False, True], [False, True]]),
+          expected=[[0, 0], [0, 2], [1, 1]]),
+      dict(  # shape=[D1, (D2), (D3)]
+          condition=ragged_factory_ops.constant_value([
+              [[True, False, True], [False, True]],
+              [[True], [], [False], [False, True, False]]
+          ]),
+          expected=[[0, 0, 0], [0, 0, 2], [0, 1, 1],
+                    [1, 0, 0], [1, 3, 1]]),
+      dict(  # shape=[D1, (D2), D3]
+          condition=ragged_factory_ops.constant_value([
+              [[True, False], [False, True]],
+              [[True, False], [False, False], [True, False], [False, True]]
+          ], ragged_rank=1),
+          expected=[[0, 0, 0], [0, 1, 1],
+                    [1, 0, 0], [1, 2, 0], [1, 3, 1]]),
+      dict(  # shape=[D1, (D2), (D3), (D4)]
+          condition=ragged_factory_ops.constant_value([
+              [[[], [True]]],
+              [[[True, False, True], [False, True]],
+               [[True], [], [False], [False, True, False]]]
+          ]),
+          expected=[[0, 0, 1, 0],
+                    [1, 0, 0, 0], [1, 0, 0, 2], [1, 0, 1, 1],
+                    [1, 1, 0, 0], [1, 1, 3, 1]]),
+
+      #=========================================================================
+      # Elementwise multiplexing
+      #=========================================================================
+      dict(  # shape=[]
+          condition=True, x='A', y='a', expected=b'A'),
+      dict(  # shape=[]
+          condition=False, x='A', y='a', expected=b'a'),
+      dict(  # shape=[D1]
+          condition=[True, False, True],
+          x=['A', 'B', 'C'],
+          y=['a', 'b', 'c'],
+          expected=[b'A', b'b', b'C']),
+      dict(  # shape=[D1, D2]
+          condition=[[True, False], [False, True]],
+          x=[['A', 'B'], ['D', 'E']],
+          y=[['a', 'b'], ['d', 'e']],
+          expected=[[b'A', b'b'], [b'd', b'E']]),
+      dict(  # shape=[D1, (D2)]
+          condition=ragged_factory_ops.constant_value(
+              [[True, False, True], [False, True]]),
+          x=ragged_factory_ops.constant_value([['A', 'B', 'C'], ['D', 'E']]),
+          y=ragged_factory_ops.constant_value([['a', 'b', 'c'], ['d', 'e']]),
+          expected=ragged_factory_ops.constant_value(
+              [[b'A', b'b', b'C'], [b'd', b'E']])),
+      dict(  # shape=[D1, (D2), D3]
+          condition=ragged_factory_ops.constant_value([
+              [[True, False], [False, True]],
+              [[True, False], [False, False], [True, False], [False, True]]
+          ], ragged_rank=1),
+          x=ragged_factory_ops.constant_value([
+              [['A', 'B'], ['C', 'D']],
+              [['E', 'F'], ['G', 'H'], ['I', 'J'], ['K', 'L']]
+          ], ragged_rank=1),
+          y=ragged_factory_ops.constant_value([
+              [['a', 'b'], ['c', 'd']],
+              [['e', 'f'], ['g', 'h'], ['i', 'j'], ['k', 'l']]
+          ], ragged_rank=1),
+          expected=ragged_factory_ops.constant_value([
+              [[b'A', b'b'], [b'c', b'D']],
+              [[b'E', b'f'], [b'g', b'h'], [b'I', b'j'], [b'k', b'L']]
+          ], ragged_rank=1)),
+      dict(  # shape=[D1, (D2), (D3), (D4)]
+          condition=ragged_factory_ops.constant_value([
+              [[[], [True]]],
+              [[[True, False, True], [False, True]],
+               [[True], [], [False], [False, True, False]]]
+          ]),
+          x=ragged_factory_ops.constant_value([
+              [[[], ['A']]],
+              [[['B', 'C', 'D'], ['E', 'F']],
+               [['G'], [], ['H'], ['I', 'J', 'K']]]
+          ]),
+          y=ragged_factory_ops.constant_value([
+              [[[], ['a']]],
+              [[['b', 'c', 'd'], ['e', 'f']],
+               [['g'], [], ['h'], ['i', 'j', 'k']]]
+          ]),
+          expected=ragged_factory_ops.constant_value([
+              [[[], [b'A']]],
+              [[[b'B', b'c', b'D'], [b'e', b'F']],
+               [[b'G'], [], [b'h'], [b'i', b'J', b'k']]]
+          ])),
+
+      #=========================================================================
+      # Broadcasting
+      #=========================================================================
+      dict(  # c.shape=[D1], x.shape=[D1, D2], y.shape=[D1, D2]
+          condition=[[True], [False], [True]],
+          x=[['A', 'B'], ['C', 'D'], ['E', 'F']],
+          y=[['a', 'b'], ['c', 'd'], ['e', 'f']],
+          expected=[[b'A', b'B'], [b'c', b'd'], [b'E', b'F']]),
+      dict(  # c.shape=[D1], x.shape=[D1, (D2)], y.shape=[D1, (D2)]
+          condition=[[True], [False], [True]],
+          x=ragged_factory_ops.constant_value(
+              [['A', 'B', 'C'], ['D', 'E'], ['F', 'G']]),
+          y=ragged_factory_ops.constant_value(
+              [['a', 'b', 'c'], ['d', 'e'], ['f', 'g']]),
+          expected=ragged_factory_ops.constant_value(
+              [[b'A', b'B', b'C'], [b'd', b'e'], [b'F', b'G']])),
+      dict(  # c.shape=[D1, None], x.shape=[], y.shape=[]
+          condition=ragged_factory_ops.constant_value(
+              [[True, False, True, True], [True, False]]),
+          x=10,
+          y=20,
+          expected=ragged_factory_ops.constant_value(
+              [[10, 20, 10, 10], [10, 20]])),
+      dict(  # c.shape=[D1, D2], x.shape=[D1, 1], y.shape=[1, D2]
+          condition=[[True, False], [True, False], [False, True]],
+          x=[[10], [20], [30]],
+          y=[[40, 50]],
+          expected=[[10, 50], [20, 50], [40, 30]]),
+      dict(  # c.shape=[D1, (D2), D3], x.shape=[D1, (D2), 1], y.shape=[D3]
+          condition=ragged_factory_ops.constant_value(
+              [[[True, False], [False, True]], [[True, True]]],
+              ragged_rank=1),
+          x=ragged_factory_ops.constant_value([[[10], [20]], [[30]]],
+                                              ragged_rank=1),
+          y=np.array([[[40, 50]]]),
+          expected=[[[10, 50], [40, 20]], [[30, 30]]]),
+  ])   # pyformat: disable
+  def testRaggedWhere(self, condition, expected, x=None, y=None):
+    result = ragged_where_op.where_v2(condition, x, y)
+    self.assertAllEqual(result, expected)
+
+  @parameterized.parameters([
+      dict(
+          condition=[True, False],
+          x=[1, 2],
+          error=ValueError,
+          message='x and y must be either both None or both non-None'),
+      dict(
+          condition=ragged_factory_ops.constant_value([[True, False, True],
+                                                       [False, True]]),
+          x=ragged_factory_ops.constant_value([['A', 'B', 'C'], ['D', 'E']]),
+          y=[['a', 'b'], ['d', 'e']],
+          error=errors.InvalidArgumentError,
+          message=r'must be broadcastable|Unable to broadcast'),
+  ])
+  def testRaggedWhereErrors(self, condition, error, message, x=None, y=None):
+    with self.assertRaisesRegex(error, message):
+      self.evaluate(ragged_where_op.where_v2(condition, x, y))
+
+
+if __name__ == '__main__':
+  googletest.main()

tensorflow/python/ops/ragged/segment_id_ops.py

@@ -0,0 +1,134 @@
+# 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.
+# ==============================================================================
+"""Ops for converting between row_splits and segment_ids."""
+
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import tensor_util
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.ops.ragged import ragged_util
+from tensorflow.python.util import dispatch
+from tensorflow.python.util.tf_export import tf_export
+
+
+# For background on "segments" and "segment ids", see:
+# https://www.tensorflow.org/api_docs/python/tf/math#Segmentation
+@tf_export("ragged.row_splits_to_segment_ids")
+@dispatch.add_dispatch_support
+def row_splits_to_segment_ids(splits, name=None, out_type=None):
+  """Generates the segmentation corresponding to a RaggedTensor `row_splits`.
+
+  Returns an integer vector `segment_ids`, where `segment_ids[i] == j` if
+  `splits[j] <= i < splits[j+1]`.  Example:
+
+  >>> print(tf.ragged.row_splits_to_segment_ids([0, 3, 3, 5, 6, 9]))
+   tf.Tensor([0 0 0 2 2 3 4 4 4], shape=(9,), dtype=int64)
+
+  Args:
+    splits: A sorted 1-D integer Tensor.  `splits[0]` must be zero.
+    name: A name prefix for the returned tensor (optional).
+    out_type: The dtype for the return value.  Defaults to `splits.dtype`,
+      or `tf.int64` if `splits` does not have a dtype.
+
+  Returns:
+    A sorted 1-D integer Tensor, with `shape=[splits[-1]]`
+
+  Raises:
+    ValueError: If `splits` is invalid.
+  """
+  with ops.name_scope(name, "RaggedSplitsToSegmentIds", [splits]) as name:
+    splits = ops.convert_to_tensor(
+        splits, name="splits",
+        preferred_dtype=dtypes.int64)
+    if splits.dtype not in (dtypes.int32, dtypes.int64):
+      raise ValueError("splits must have dtype int32 or int64")
+    splits.shape.assert_has_rank(1)
+    if tensor_shape.dimension_value(splits.shape[0]) == 0:
+      raise ValueError("Invalid row_splits: []")
+    if out_type is None:
+      out_type = splits.dtype
+    else:
+      out_type = dtypes.as_dtype(out_type)
+    row_lengths = splits[1:] - splits[:-1]
+    nrows = array_ops.shape(splits, out_type=out_type)[-1] - 1
+    indices = math_ops.range(nrows)
+    return ragged_util.repeat(indices, repeats=row_lengths, axis=0)
+
+
+# For background on "segments" and "segment ids", see:
+# https://www.tensorflow.org/api_docs/python/tf/math#Segmentation
+@tf_export("ragged.segment_ids_to_row_splits")
+@dispatch.add_dispatch_support
+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:
+
+  >>> print(tf.ragged.segment_ids_to_row_splits([0, 0, 0, 2, 2, 3, 4, 4, 4]))
+  tf.Tensor([0 3 3 5 6 9], shape=(6,), dtype=int64)
+
+  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]`.
+  """
+  # Local import bincount_ops to avoid import-cycle.
+  from tensorflow.python.ops import bincount_ops  # pylint: disable=g-import-not-at-top
+  if out_type is None:
+    if isinstance(segment_ids, tensor.Tensor):
+      out_type = segment_ids.dtype
+    elif isinstance(num_segments, tensor.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 = bincount_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

tensorflow/python/ops/ragged/string_ngrams_op_test.py

@@ -0,0 +1,346 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the b"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 b"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.
+# ==============================================================================
+"""Tests for the Tensorflow strings.ngrams op."""
+
+from absl.testing import parameterized
+
+from tensorflow.python.eager import def_function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import tensor
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_string_ops
+from tensorflow.python.platform import test
+
+
+class StringNgramsTest(test_util.TensorFlowTestCase, parameterized.TestCase):
+
+  def test_unpadded_ngrams(self):
+    data = [[b"aa", b"bb", b"cc", b"dd"], [b"ee", b"ff"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=3, separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[b"aa|bb|cc", b"bb|cc|dd"], []]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_tuple_multi_ngrams(self):
+    data = [[b"aa", b"bb", b"cc", b"dd"], [b"ee", b"ff"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=(2, 3), separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[b"aa|bb", b"bb|cc", b"cc|dd", b"aa|bb|cc", b"bb|cc|dd"],
+                       [b"ee|ff"]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_tuple_multi_ngrams_inverted_order(self):
+    data = [[b"aa", b"bb", b"cc", b"dd"], [b"ee", b"ff"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=(3, 2), separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[b"aa|bb|cc", b"bb|cc|dd", b"aa|bb", b"bb|cc", b"cc|dd"],
+                       [b"ee|ff"]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_list_multi_ngrams(self):
+    data = [[b"aa", b"bb", b"cc", b"dd"], [b"ee", b"ff"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=[2, 3], separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[b"aa|bb", b"bb|cc", b"cc|dd", b"aa|bb|cc", b"bb|cc|dd"],
+                       [b"ee|ff"]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_multi_ngram_ordering(self):
+    data = [[b"aa", b"bb", b"cc", b"dd"], [b"ee", b"ff"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=[3, 2], separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[b"aa|bb|cc", b"bb|cc|dd", b"aa|bb", b"bb|cc", b"cc|dd"],
+                       [b"ee|ff"]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_fully_padded_ngrams(self):
+    data = [[b"a"], [b"b", b"c", b"d"], [b"e", b"f"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=3, separator=b"|", pad_values=(b"LP", b"RP"))
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [
+        [b"LP|LP|a", b"LP|a|RP", b"a|RP|RP"],  # 0
+        [b"LP|LP|b", b"LP|b|c", b"b|c|d", b"c|d|RP", b"d|RP|RP"],  # 1
+        [b"LP|LP|e", b"LP|e|f", b"e|f|RP", b"f|RP|RP"]  # 2
+    ]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_ngram_padding_size_cap(self):
+    # Validate that the padding size is never greater than ngram_size - 1.
+    data = [[b"a"], [b"b", b"c", b"d"], [b"e", b"f"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor,
+        ngram_width=3,
+        separator=b"|",
+        pad_values=(b"LP", b"RP"),
+        padding_width=10)
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [
+        [b"LP|LP|a", b"LP|a|RP", b"a|RP|RP"],  # 0
+        [b"LP|LP|b", b"LP|b|c", b"b|c|d", b"c|d|RP", b"d|RP|RP"],  # 1
+        [b"LP|LP|e", b"LP|e|f", b"e|f|RP", b"f|RP|RP"]  # 2
+    ]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_singly_padded_ngrams(self):
+    data = [[b"a"], [b"b", b"c", b"d"], [b"e", b"f"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor,
+        ngram_width=5,
+        separator=b"|",
+        pad_values=(b"LP", b"RP"),
+        padding_width=1)
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[], [b"LP|b|c|d|RP"], []]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_singly_padded_ngrams_with_preserve_short(self):
+    data = [[b"a"], [b"b", b"c", b"d"], [b"e", b"f"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor,
+        ngram_width=5,
+        separator=b"|",
+        pad_values=(b"LP", b"RP"),
+        padding_width=1,
+        preserve_short_sequences=True)
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[b"LP|a|RP"], [b"LP|b|c|d|RP"], [b"LP|e|f|RP"]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_singly_padded_multiple_ngrams(self):
+    data = [[b"a"], [b"b", b"c", b"d"], [b"e", b"f"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor,
+        ngram_width=(1, 5),
+        separator=b"|",
+        pad_values=(b"LP", b"RP"),
+        padding_width=1)
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[b"a"], [b"b", b"c", b"d", b"LP|b|c|d|RP"], [b"e", b"f"]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_single_padding_string(self):
+    data = [[b"a"], [b"b", b"c", b"d"], [b"e", b"f"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor,
+        ngram_width=5,
+        separator=b"|",
+        pad_values=b"[PAD]",
+        padding_width=1)
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[], [b"[PAD]|b|c|d|[PAD]"], []]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_explicit_multiply_padded_ngrams(self):
+    data = [[b"a"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor,
+        ngram_width=5,
+        separator=b"|",
+        pad_values=(b"LP", b"RP"),
+        padding_width=2)
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[b"LP|LP|a|RP|RP"]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_ragged_inputs_with_multiple_ragged_dimensions(self):
+    data = [[[[b"aa", b"bb", b"cc", b"dd"]], [[b"ee", b"ff"]]]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=3, separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[[[b"aa|bb|cc", b"bb|cc|dd"]], [[]]]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_ragged_inputs_with_multiple_ragged_dimensions_and_preserve(self):
+    data = [[[[b"aa", b"bb", b"cc", b"dd"]], [[b"ee", b"ff"]]]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor,
+        ngram_width=3,
+        separator=b"|",
+        preserve_short_sequences=True)
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[[[b"aa|bb|cc", b"bb|cc|dd"]], [[b"ee|ff"]]]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_ragged_inputs_with_multiple_ragged_dimensions_bigrams(self):
+    data = [[[[b"aa", b"bb", b"cc", b"dd"]], [[b"ee", b"ff"]]]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=2, separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[[[b"aa|bb", b"bb|cc", b"cc|dd"]], [[b"ee|ff"]]]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_ragged_inputs_with_multiple_ragged_dimensions_and_multiple_ngrams(
+      self):
+    data = [[[[b"aa", b"bb", b"cc", b"dd"]], [[b"ee", b"ff"]]]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=(3, 4), separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[[[b"aa|bb|cc", b"bb|cc|dd", b"aa|bb|cc|dd"]], [[]]]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_dense_input_rank_3(self):
+    data = [[[b"a", b"z"], [b"b", b""]], [[b"b", b""], [b"e", b"f"]]]
+    data_tensor = constant_op.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=3, separator=b"|", pad_values=(b"LP", b"RP"))
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[[b"LP|LP|a", b"LP|a|z", b"a|z|RP", b"z|RP|RP"],
+                        [b"LP|LP|b", b"LP|b|", b"b||RP", b"|RP|RP"]],
+                       [[b"LP|LP|b", b"LP|b|", b"b||RP", b"|RP|RP"],
+                        [b"LP|LP|e", b"LP|e|f", b"e|f|RP", b"f|RP|RP"]]]
+    self.assertIsInstance(ngram_op, tensor.Tensor)
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_dense_input(self):
+    data = [[b"a", b"z"], [b"b", b""], [b"e", b"f"]]
+    data_tensor = constant_op.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=3, separator=b"|", pad_values=(b"LP", b"RP"))
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [
+        [b"LP|LP|a", b"LP|a|z", b"a|z|RP", b"z|RP|RP"],
+        [b"LP|LP|b", b"LP|b|", b"b||RP", b"|RP|RP"],
+        [b"LP|LP|e", b"LP|e|f", b"e|f|RP", b"f|RP|RP"],
+    ]
+    self.assertIsInstance(ngram_op, tensor.Tensor)
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_input_list_input(self):
+    data = [[b"a", b"z"], [b"b", b""], [b"e", b"f"]]
+    ngram_op = ragged_string_ops.ngrams(
+        data, ngram_width=3, separator=b"|", pad_values=(b"LP", b"RP"))
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [
+        [b"LP|LP|a", b"LP|a|z", b"a|z|RP", b"z|RP|RP"],
+        [b"LP|LP|b", b"LP|b|", b"b||RP", b"|RP|RP"],
+        [b"LP|LP|e", b"LP|e|f", b"e|f|RP", b"f|RP|RP"],
+    ]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_vector_input(self):
+    data = [b"a", b"z"]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=3, separator=b"|", pad_values=(b"LP", b"RP"))
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [b"LP|LP|a", b"LP|a|z", b"a|z|RP", b"z|RP|RP"]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_dense_input_with_multiple_ngrams(self):
+    data = [[b"a", b"b", b"c", b"d"], [b"e", b"f", b"g", b"h"]]
+    data_tensor = ragged_factory_ops.constant(data)
+    ngram_op = ragged_string_ops.ngrams(
+        data_tensor, ngram_width=(1, 2, 3), separator=b"|")
+    result = self.evaluate(ngram_op)
+    expected_ngrams = [[
+        b"a", b"b", b"c", b"d", b"a|b", b"b|c", b"c|d", b"a|b|c", b"b|c|d"
+    ], [b"e", b"f", b"g", b"h", b"e|f", b"f|g", b"g|h", b"e|f|g", b"f|g|h"]]
+    self.assertAllEqual(expected_ngrams, result)
+
+  def test_input_with_no_values(self):
+    data = ragged_factory_ops.constant([[], [], []], dtype=dtypes.string)
+    ngram_op = ragged_string_ops.ngrams(data, (1, 2))
+    result = self.evaluate(ngram_op)
+    self.assertAllEqual([0, 0, 0, 0], result.row_splits)
+    self.assertAllEqual(constant_op.constant([], dtype=dtypes.string),
+                        result.values)
+
+  @parameterized.parameters([
+      dict(
+          data=[b"a", b"z"],
+          ngram_width=2,
+          pad_values=5,
+          exception=TypeError,
+          error="pad_values must be a string, tuple of strings, or None."),
+      dict(
+          data=[b"a", b"z"],
+          ngram_width=2,
+          pad_values=[5, 3],
+          exception=TypeError,
+          error="pad_values must be a string, tuple of strings, or None."),
+      dict(
+          data=[b"a", b"z"],
+          ngram_width=2,
+          padding_width=0,
+          pad_values="X",
+          error="padding_width must be greater than 0."),
+      dict(
+          data=[b"a", b"z"],
+          ngram_width=2,
+          padding_width=1,
+          error="pad_values must be provided if padding_width is set."),
+      dict(
+          data=b"hello",
+          ngram_width=2,
+          padding_width=1,
+          pad_values="X",
+          error="Data must have rank>0"),
+      dict(
+          data=[b"hello", b"world"],
+          ngram_width=[1, 2, -1],
+          padding_width=1,
+          pad_values="X",
+          error="All ngram_widths must be greater than 0. Got .*"),
+  ])
+  def test_error(self,
+                 data,
+                 ngram_width,
+                 separator=" ",
+                 pad_values=None,
+                 padding_width=None,
+                 preserve_short_sequences=False,
+                 error=None,
+                 exception=ValueError):
+    with self.assertRaisesRegex(exception, error):
+      ragged_string_ops.ngrams(data, ngram_width, separator, pad_values,
+                               padding_width, preserve_short_sequences)
+
+  def test_unknown_rank_error(self):
+    # Use a tf.function that erases shape information.
+    @def_function.function(
+        input_signature=[tensor.TensorSpec(None, dtypes.string)])
+    def f(v):
+      return ragged_string_ops.ngrams(v, 2)
+
+    with self.assertRaisesRegex(ValueError, "Rank of data must be known."):
+      f([b"foo", b"bar"])
+
+
+if __name__ == "__main__":
+  test.main()

tensorflow/python/ops/ragged/strings_reduce_join_op_test.py

@@ -0,0 +1,140 @@
+# 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.
+# ==============================================================================
+"""Tests for tf.strings.reduce_join."""
+from absl.testing import parameterized
+from tensorflow.python.framework import tensor_shape
+from tensorflow.python.framework import test_util
+from tensorflow.python.ops.ragged import ragged_factory_ops
+from tensorflow.python.ops.ragged import ragged_string_ops
+from tensorflow.python.platform import googletest
+
+
+@test_util.run_all_in_graph_and_eager_modes
+class StringsReduceJoinOpTest(test_util.TensorFlowTestCase,
+                              parameterized.TestCase):
+
+  def test_rank_one(self):
+    input_array = [b'this', b'is', b'a', b'test']
+    truth = b'thisisatest'
+    truth_shape = []
+    with self.cached_session():
+      output = ragged_string_ops.reduce_join(
+          inputs=input_array, axis=-1, keepdims=False, separator='')
+      output_array = self.evaluate(output)
+    self.assertAllEqual(truth, output_array)
+    self.assertAllEqual(truth_shape, output.get_shape())
+
+  @parameterized.parameters([
+      {
+          'input_array': [[
+              b'this', b'is', b'a', b'test', b'for', b'ragged', b'tensors'
+          ], [b'please', b'do', b'not', b'panic', b'!']],
+          'axis': 0,
+          'keepdims': False,
+          'truth': [
+              b'thisplease', b'isdo', b'anot', b'testpanic', b'for!', b'ragged',
+              b'tensors'
+          ],
+          'truth_shape': [7],
+      },
+      {
+          'input_array': [[
+              b'this', b'is', b'a', b'test', b'for', b'ragged', b'tensors'
+          ], [b'please', b'do', b'not', b'panic', b'!']],
+          'axis': 1,
+          'keepdims': False,
+          'truth': [b'thisisatestforraggedtensors', b'pleasedonotpanic!'],
+          'truth_shape': [2],
+      },
+      {
+          'input_array': [[
+              b'this', b'is', b'a', b'test', b'for', b'ragged', b'tensors'
+          ], [b'please', b'do', b'not', b'panic', b'!']],
+          'axis': 1,
+          'keepdims': False,
+          'truth': [
+              b'this|is|a|test|for|ragged|tensors', b'please|do|not|panic|!'
+          ],
+          'truth_shape': [2],
+          'separator': '|',
+      },
+      {
+          'input_array': [[[b'a', b'b'], [b'b', b'c']], [[b'dd', b'ee']]],
+          'axis': -1,
+          'keepdims': False,
+          'truth': [[b'a|b', b'b|c'], [b'dd|ee']],
+          'truth_shape': [2, None],
+          'separator': '|',
+      },
+      {
+          'input_array': [[[[b'a', b'b', b'c'], [b'dd', b'ee']]],
+                          [[[b'f', b'g', b'h'], [b'ii', b'jj']]]],
+          'axis': -2,
+          'keepdims': False,
+          'truth': [[[b'a|dd', b'b|ee', b'c']], [[b'f|ii', b'g|jj', b'h']]],
+          'truth_shape': [2, None, None],
+          'separator': '|',
+      },
+      {
+          'input_array': [[[b't', b'h', b'i', b's'], [b'i', b's'], [b'a'],
+                           [b't', b'e', b's', b't']],
+                          [[b'p', b'l', b'e', b'a', b's', b'e'],
+                           [b'p', b'a', b'n', b'i', b'c']]],
+          'axis': -1,
+          'keepdims': False,
+          'truth': [[b'this', b'is', b'a', b'test'], [b'please', b'panic']],
+          'truth_shape': [2, None],
+          'separator': '',
+      },
+      {
+          'input_array': [[[[b't'], [b'h'], [b'i'], [b's']], [[b'i', b's']],
+                           [[b'a', b'n']], [[b'e'], [b'r'], [b'r']]],
+                          [[[b'p'], [b'l'], [b'e'], [b'a'], [b's'], [b'e']],
+                           [[b'p'], [b'a'], [b'n'], [b'i'], [b'c']]]],
+          'axis': -1,
+          'keepdims': False,
+          'truth': [[[b't', b'h', b'i', b's'], [b'is'], [b'an'],
+                     [b'e', b'r', b'r']],
+                    [[b'p', b'l', b'e', b'a', b's', b'e'],
+                     [b'p', b'a', b'n', b'i', b'c']]],
+          'truth_shape': [2, None, None],
+          'separator': '',
+      },
+  ])
+  def test_different_ranks(self,
+                           input_array,
+                           axis,
+                           keepdims,
+                           truth,
+                           truth_shape,
+                           separator=''):
+    with self.cached_session():
+      input_tensor = ragged_factory_ops.constant(input_array)
+      output = ragged_string_ops.reduce_join(
+          inputs=input_tensor,
+          axis=axis,
+          keepdims=keepdims,
+          separator=separator)
+      output_array = self.evaluate(output)
+    self.assertAllEqual(truth, output_array)
+    if all(isinstance(s, tensor_shape.Dimension) for s in output.shape):
+      output_shape = [dim.value for dim in output.shape]
+    else:
+      output_shape = output.shape
+    self.assertAllEqual(truth_shape, output_shape)
+
+
+if __name__ == '__main__':
+  googletest.main()