def testErrors(self):
    if not context.executing_eagerly():
      self.assertRaisesRegexp(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.assertRaisesRegexp(
        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.assertRaisesRegexp(
        ValueError, r'Shapes \(1, 2\) and \(1, 3\) are incompatible',
        ragged_array_ops.boolean_mask, [[1, 2]], [[True, False, True]])

    self.assertRaisesRegexp(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.assertRaisesRegexp(ValueError, 'mask cannot be scalar',
                            ragged_array_ops.boolean_mask, [[1, 2]], True)

    self.assertRaisesRegexp(ValueError, 'mask cannot be scalar',
                            ragged_array_ops.boolean_mask,
                            ragged_factory_ops.constant([[1, 2]]), True)

 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)

 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 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 testRaggedTensorSplitsValueMismatchError(self):
   x = ragged_factory_ops.constant([[3, 1, 4], [], [1, 5]])
   y = ragged_factory_ops.constant([[1], [2, 3], [4, 5]])
   self.assertRaisesRegexp(errors.InvalidArgumentError,
                           r'Inputs must have identical ragged splits.*',
                           ragged_functional_ops.map_flat_values, math_ops.add,
                           x, y)

 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 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 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,
       args=(condition, x, y),
       expected=[[b'a', b'b', b'C'], [], [b'd', b'E']])

 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.assertRaggedEqual(result, expected)

  def testBatchRagged(self):

    def _ragged(i):
      return ragged_tensor.RaggedTensor.from_tensor(i * [[1]])

    dataset = dataset_ops.Dataset.range(10).map(_ragged).batch(5)
    expected_output = [
        ragged_factory_ops.constant([[[0]], [[1]], [[2]], [[3]], [[4]]]),
        ragged_factory_ops.constant([[[5]], [[6]], [[7]], [[8]], [[9]]])
    ]
    self.assertDatasetProduces(dataset, expected_output=expected_output)

 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.assertRaggedEqual(segmented, expected)

  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.assertRaggedEqual(g1, [[1.0, 1.0], [1.0, 1.0], [1.0, 1.0]])
    self.assertRaggedEqual(g2, [[2.0, 2.0], [2.0, 2.0], [2.0, 2.0]])

  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)
      with self.test_session():
        self.assertEqual(tiled.eval().tolist(), expected)

 def testShapeMismatchError1(self):
   dt = constant_op.constant([1, 2, 3, 4, 5, 6])
   segment_ids = ragged_factory_ops.constant([[1, 2], []])
   self.assertRaisesRegexp(
       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 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.assertRaggedEqual(output['batman'], [[2, 3, 4], [5], [6, 7, 8]])
    self.assertRaggedEqual(output['robin'], [[11, 21, 31], [41], [51, 61, 71]])

 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.assertRaggedEqual(rt_s, ref)

 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 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 testSplitV2(self,
                  input,
                  expected,
                  input_is_ragged=False,
                  **kwargs):  # pylint: disable=redefined-builtin
    # Check that we are matching the behavior of Python's str.split:
    self.assertEqual(expected, self._py_split(input, **kwargs))

    # Prepare the input tensor.
    if input_is_ragged:
      input = ragged_factory_ops.constant(input, dtype=dtypes.string)
    else:
      input = constant_op.constant(input, dtype=dtypes.string)

    # Check that the public version (which returns a RaggedTensor) works
    # correctly.
    expected_ragged = ragged_factory_ops.constant(
        expected, ragged_rank=input.shape.ndims)
    actual_ragged_v1 = ragged_string_ops.strings_split_v1(
        input, result_type="RaggedTensor", **kwargs)
    actual_ragged_v1_input_kwarg = ragged_string_ops.strings_split_v1(
        input=input, result_type="RaggedTensor", **kwargs)
    actual_ragged_v1_source_kwarg = ragged_string_ops.strings_split_v1(
        source=input, result_type="RaggedTensor", **kwargs)
    actual_ragged_v2 = ragged_string_ops.string_split_v2(input, **kwargs)
    actual_ragged_v2_input_kwarg = ragged_string_ops.string_split_v2(
        input=input, **kwargs)
    self.assertRaggedEqual(expected_ragged, actual_ragged_v1)
    self.assertRaggedEqual(expected_ragged, actual_ragged_v1_input_kwarg)
    self.assertRaggedEqual(expected_ragged, actual_ragged_v1_source_kwarg)
    self.assertRaggedEqual(expected_ragged, actual_ragged_v2)
    self.assertRaggedEqual(expected_ragged, actual_ragged_v2_input_kwarg)

    # Check that the internal version (which returns a SparseTensor) works
    # correctly.  Note: the internal version oly supports vector inputs.
    if input.shape.ndims == 1:
      expected_sparse = self.evaluate(expected_ragged.to_sparse())
      actual_sparse_v1 = ragged_string_ops.strings_split_v1(
          input, result_type="SparseTensor", **kwargs)
      actual_sparse_v2 = string_ops.string_split_v2(input, **kwargs)
      for actual_sparse in [actual_sparse_v1, actual_sparse_v2]:
        self.assertEqual(expected_sparse.indices.tolist(),
                         self.evaluate(actual_sparse.indices).tolist())
        self.assertEqual(expected_sparse.values.tolist(),
                         self.evaluate(actual_sparse.values).tolist())
        self.assertEqual(expected_sparse.dense_shape.tolist(),
                         self.evaluate(actual_sparse.dense_shape).tolist())

  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)

    # 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.inner_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)

    with self.test_session():
      result = self.evaluate(rt)
      if rt.shape.ndims > 0:
        self.assertEqual(result.tolist(), pylist)
        if expected_shape is not None:
          self.assertEqual(result.shape, expected_shape)
      else:
        self.assertEqual(result, pylist)
        if expected_shape is not None:
          self.assertEqual((), expected_shape)