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 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 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 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 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 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 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 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 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 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 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 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 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)
def testBinaryOpSparseAndRagged(self): x = ragged_factory_ops.constant([[1, 2, 3], [4, 5]]) y = sparse_tensor.SparseTensor([[0, 0], [0, 1], [2, 0]], [1, 2, 3], [3, 2]) with self.assertRaises((TypeError, ValueError)): self.evaluate(math_ops.add(x, y)) with self.assertRaises((TypeError, ValueError)): self.evaluate(math_ops.add_n([x, y]))
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) ]
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 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 testRagged(self): def _ragged(i): return ragged_tensor.RaggedTensor.from_tensor(i * [[1]]) dataset = dataset_ops.Dataset.range(5).map(_ragged) self.assertDatasetProduces( dataset, expected_output=[ragged_factory_ops.constant([[i]]) for i in range(5)])
def testElementwiseOpUnknownRankError(self): if context.executing_eagerly(): return x = ragged_factory_ops.constant([[1, 2], [3]]) y = ragged_tensor.RaggedTensor.from_row_splits( array_ops.placeholder_with_default([1, 2, 3], shape=None), x.row_splits) with self.assertRaisesRegexp(ValueError, r'Unable to broadcast: unknown rank'): math_ops.add(x, y)
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.assertRaggedAlmostEqual(segmented, expected, places=5)
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.assertRaggedEqual(v1, [[1, 1, 1], [], [1, 1], [1]]) self.assertRaggedEqual(v2, [[1, 4, 9], [], [16, 25], [36]]) self.assertRaggedEqual(v3, [[6, 7, 8], [], [9, 10], [11]])
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.assertRaggedEqual(segmented, expected)
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.assertRaggedEqual(concatenated, [[1, 2], [3, 4]])
def test_merge_with_ragged_input(self, layer): ragged_data = ragged_factory_ops.constant( [[1., 1., 1.], [1., 1.], [1., 1., 1., 1.]], ragged_rank=1) dense_data = ragged_data.to_tensor() input1 = keras.Input(shape=(None, ), ragged=True) input2 = keras.Input(shape=(None, ), ragged=True) out = keras.layers.Add()([input1, input2]) model = keras.models.Model(inputs=[input1, input2], outputs=out) out_ragged = model.predict([ragged_data, ragged_data], steps=1) out_ragged = ragged_tensor.convert_to_tensor_or_ragged_tensor( out_ragged).to_tensor() input1 = keras.Input(shape=(None, )) input2 = keras.Input(shape=(None, )) out = keras.layers.Add()([input1, input2]) model = keras.models.Model(inputs=[input1, input2], outputs=out) out_dense = model.predict([dense_data, dense_data], steps=1) self.assertAllEqual(out_dense, out_ragged)
def preserveStaticShape(self): rt = ragged_factory_ops.constant([[1, 2], [], [3]]) rt_s = structure.type_spec_from_value(rt) rt_after = structure.from_tensor_list( rt_s, structure.to_tensor_list(rt_s, rt)) self.assertEqual(rt_after.row_splits.shape.as_list(), rt.row_splits.shape.as_list()) self.assertEqual(rt_after.values.shape.as_list(), [None]) st = sparse_tensor.SparseTensor(indices=[[3, 4]], values=[-1], dense_shape=[4, 5]) st_s = structure.type_spec_from_value(st) st_after = structure.from_tensor_list( st_s, structure.to_tensor_list(st_s, st)) self.assertEqual(st_after.indices.shape.as_list(), [None, 2]) self.assertEqual(st_after.values.shape.as_list(), [None]) self.assertEqual(st_after.dense_shape.shape.as_list(), st.dense_shape.shape.as_list())
def test_ragged_tensor_model_predict(self): # Create a model that accepts a sparse input and runs a "Dense" layer on it. model_input = input_layer.Input(shape=(None, ), ragged=True) self.assertEqual([None, None], model_input.shape.as_list()) layers = [Embedding(input_dim=7, output_dim=5)] model = get_model_from_layers_with_input(layers, model_input=model_input) ragged_input = ragged_factory_ops.constant([ [1, 2, 3, 4, 5], [2, 4], ]) shape = model(ragged_input).shape self.assertEqual((2, None, 5), self._normalize_shape(shape)) shape = model.predict(ragged_input, steps=1).shape self.assertEqual((2, None, 5), self._normalize_shape(shape))
def testNestedFields(self): PossiblyRaggedTensor = typing.Union[ops.Tensor, ragged_tensor.RaggedTensor] ToyFeatures = typing.Mapping[str, PossiblyRaggedTensor] class ToyInfo(extension_type.ExtensionType): version: str toys: typing.Tuple[typing.Tuple[str, ops.Tensor, ToyFeatures], ...] boxes: typing.Mapping[str, ops.Tensor] authors = [[b'A', b'Aardvark'], [b'Z', b'Zhook']] toys = [('car', 1.0, { 'size': [8, 3, 2], 'color': [0.3, 0.2, 0.8] }), ('book', 3.7, { 'authors': ragged_factory_ops.constant(authors) })] boxes = {'green': ['car'], 'blue': ['car', 'book', 'book']} toy_info = ToyInfo(version='1.0 alpha', toys=toys, boxes=boxes) self.assertEqual(toy_info.version, '1.0 alpha') self.assertEqual(toy_info.toys[0][0], 'car') self.assertIsInstance(toy_info.toys[0][1], ops.Tensor) self.assertAllEqual(toy_info.toys[0][1], 1.0) self.assertEqual(set(toy_info.toys[0][2].keys()), {'size', 'color'}) self.assertIsInstance(toy_info.toys[0][2]['size'], ops.Tensor) self.assertAllEqual(toy_info.toys[0][2]['size'], [8, 3, 2]) self.assertIsInstance(toy_info.toys[1][2]['authors'], ragged_tensor.RaggedTensor) self.assertAllEqual(toy_info.toys[1][2]['authors'], authors) self.assertAllEqual(toy_info.boxes['green'], [b'car']) self.assertAllEqual(toy_info.boxes['blue'], ['car', 'book', 'book']) expected_repr = ( r"ToyInfo\(version='1.0 alpha', toys=\(" r"\('car', <tf.Tensor[^>]*>, ImmutableDict\(" r"{'size': <tf.Tensor[^>]*>, 'color': <tf.Tensor[^>]*>}\)\), " r"\('book', <tf.Tensor[^>]*>, ImmutableDict\(" r"{'authors': (<tf.RaggedTensor[^>]*>|tf.RaggedTensor\(.*\))}\)\)\), " r'boxes=ImmutableDict\(' r"{'green': <tf.Tensor[^>]*>, 'blue': <tf.Tensor[^>]*>}\)\)") self.assertRegex(repr(toy_info), expected_repr)
def test_Bidirectional_ragged_input(self, merge_mode): if test.is_built_with_rocm(): # ragged tenors are not supported in ROCM RNN implementation self.skipTest('Test not supported on the ROCm platform') np.random.seed(100) rnn = keras.layers.LSTM units = 3 x = ragged_factory_ops.constant( [[[1, 1, 1], [1, 1, 1]], [[1, 1, 1]], [[1, 1, 1], [1, 1, 1], [1, 1, 1], [1, 1, 1]], [[1, 1, 1], [1, 1, 1], [1, 1, 1]]], ragged_rank=1) x = math_ops.cast(x, 'float32') # pylint: disable=g-long-lambda with self.cached_session(): if merge_mode == 'ave': merge_func = lambda y, y_rev: (y + y_rev) / 2 elif merge_mode == 'concat': merge_func = lambda y, y_rev: ragged_concat_ops.concat( (y, y_rev), axis=-1) elif merge_mode == 'mul': merge_func = lambda y, y_rev: (y * y_rev) # pylint: enable=g-long-lambda inputs = keras.Input( shape=(None, 3), batch_size=4, dtype='float32', ragged=True) layer = keras.layers.Bidirectional( rnn(units, return_sequences=True), merge_mode=merge_mode) f_merged = keras.backend.function([inputs], layer(inputs)) f_forward = keras.backend.function([inputs], layer.forward_layer(inputs)) f_backward = keras.backend.function( [inputs], array_ops.reverse(layer.backward_layer(inputs), axis=[1])) y_merged = f_merged(x) y_expected = merge_func( ragged_tensor.convert_to_tensor_or_ragged_tensor(f_forward(x)), ragged_tensor.convert_to_tensor_or_ragged_tensor(f_backward(x))) y_merged = ragged_tensor.convert_to_tensor_or_ragged_tensor(y_merged) self.assertAllClose(y_merged.flat_values, y_expected.flat_values)
def testReduceVar(self): x = np.array([[0, 0, 0], [0, 0, 0]], "float32") self.assertAllClose(self.evaluate(math_ops.reduce_variance(x)), 0) self.assertAllClose( self.evaluate(math_ops.reduce_variance(x, axis=0)), [0, 0, 0]) x = [[1, 2, 1, 1], [1, 1, 0, 1]] with self.assertRaisesRegex(TypeError, "must be either real or complex"): math_ops.reduce_variance(x) x = [[1., 2., 1., 1.], [1., 1., 0., 1.]] self.assertEqual(self.evaluate(math_ops.reduce_variance(x)), 0.25) x_np = np.array(x) self.assertEqual(np.var(x_np), 0.25) self.assertEqual(self.evaluate(math_ops.reduce_variance(x_np)), 0.25) x = ragged_factory_ops.constant([[5., 1., 4., 1.], [], [5., 9., 2.], [5.], []]) self.assertAllClose(math_ops.reduce_variance(x, axis=0), [0., 16., 1., 0.])
def test_ragged_int_input(self): vocab_data = np.array([10, 11, 12, 13], dtype=np.int64) input_array = ragged_factory_ops.constant( [[10, 11, 13], [13, 12, 10, 42]], dtype=np.int64) expected_output = [[2, 3, 5], [5, 4, 2, 1]] input_data = keras.Input(shape=(None, ), dtype=dtypes.int64, ragged=True) layer = get_layer_class()(max_tokens=None, dtype=dtypes.int64, num_oov_indices=1, mask_token=0, oov_token=-1) layer.set_vocabulary(vocab_data) int_data = layer(input_data) model = keras.Model(inputs=input_data, outputs=int_data) output_dataset = model.predict(input_array) self.assertAllEqual(expected_output, output_dataset)
def test_ragged_string_input(self): vocab_data = ["earth", "wind", "and", "fire"] input_array = ragged_factory_ops.constant( [["earth", "wind", "fire"], ["fire", "and", "earth", "michigan"]]) expected_output = [[2, 3, 5], [5, 4, 2, 1]] input_data = keras.Input(shape=(None, ), dtype=dtypes.string, ragged=True) layer = get_layer_class()(max_tokens=None, num_oov_indices=1, mask_token="", oov_token="[OOV]", dtype=dtypes.string) layer.set_vocabulary(vocab_data) int_data = layer(input_data) model = keras.Model(inputs=input_data, outputs=int_data) output_dataset = model.predict(input_array) self.assertAllEqual(expected_output, output_dataset)
def testDetokenizeFailsForSparseVocab(self): vocab = ["a", "##b", "##c"] ids = [0, 10, 20] init = lookup_ops.KeyValueTensorInitializer(vocab, ids, key_dtype=dtypes.string, value_dtype=dtypes.int64) table = lookup_ops.StaticVocabularyTableV1( init, num_oov_buckets=1, lookup_key_dtype=dtypes.string) self.evaluate(table.initializer) tokenizer = WordpieceTokenizer(table) words = ragged_factory_ops.constant([["abb", "abc"], ["abcbc"]]) subwords_ids = tokenizer.tokenize(words) with self.assertRaisesRegex(errors_impl.InvalidArgumentError, "detokenize.*?dense on the interval"): result = tokenizer.detokenize(subwords_ids) self.evaluate(result)
def testEmptyDimensions(self): test_value = ragged_factory_ops.constant( [[[b'I love Flume!', b'Good day. . .'], []], [], [[b'I love Zhu!', b'Good night'], [b'A scrub is', b'a guy']]]) expected_tokens = [[[[b'I', b'love', b'Flume!'], [b'Good', b'day.', b'.', b'.']], []], [], [[[b'I', b'love', b'Zhu!'], [b'Good', b'night']], [[b'A', b'scrub', b'is'], [b'a', b'guy']]]] expected_offset_starts = [[[[0, 2, 7], [0, 5, 10, 12]], []], [], [[[0, 2, 7], [0, 5]], [[0, 2, 8], [0, 2]]]] expected_offset_limits = [[[[1, 6, 13], [4, 9, 11, 13]], []], [], [[[1, 6, 11], [4, 10]], [[1, 7, 10], [1, 5]]]] tokens = self.whitespace_tokenizer.tokenize(test_value) self.assertRaggedEqual(tokens, expected_tokens) (tokens, starts, limits) = ( self.whitespace_tokenizer.tokenize_with_offsets(test_value)) self.assertRaggedEqual(tokens, expected_tokens) self.assertRaggedEqual(starts, expected_offset_starts) self.assertRaggedEqual(limits, expected_offset_limits)
def test_globalpooling_1d_with_ragged(self): ragged_data = ragged_factory_ops.constant( [[[1.0, 1.0], [2.0, 2.0], [3.0, 3.0]], [[1.0, 1.0], [2.0, 2.0]]], ragged_rank=1) dense_data = ragged_data.to_tensor() inputs = keras.Input(shape=(None, 2), dtype='float32', ragged=True) out = keras.layers.GlobalAveragePooling1D()(inputs) model = keras.models.Model(inputs=inputs, outputs=out) output_ragged = model.predict(ragged_data, steps=1) inputs = keras.Input(shape=(None, 2), dtype='float32') masking = keras.layers.Masking(mask_value=0., input_shape=(3, 2))(inputs) out = keras.layers.GlobalAveragePooling1D()(masking) model = keras.models.Model(inputs=inputs, outputs=out) output_dense = model.predict(dense_data, steps=1) self.assertAllEqual(output_ragged, output_dense)
def test_bert_tokenizer(self, text_inputs, expected, vocab=None, expected_extracted=None, lower_case=True, num_oov=1): text_inputs = constant_op.constant(text_inputs) if not vocab: vocab = _VOCAB table = _create_table(vocab, num_oov) self.evaluate(table.initializer) tokenizer = bert_tokenizer.BertTokenizer( table, token_out_type=dtypes.string, lower_case=lower_case) results = tokenizer.tokenize(text_inputs) self.assertAllEqual(results, expected) # Verify that the int ids are the same. expected_rt = ragged_factory_ops.constant(expected) expected_int = table.lookup(expected_rt.flat_values) expected_int_rt = ragged_tensor.RaggedTensor.from_nested_row_splits( expected_int, expected_rt.nested_row_splits) int_tokenizer = bert_tokenizer.BertTokenizer( vocab_lookup_table=table, token_out_type=dtypes.int64, lower_case=lower_case) results_int = int_tokenizer.tokenize(text_inputs) self.assertAllEqual(results_int, expected_int_rt) # Verify that the offsets can extract the expected tokens _, begin, end = tokenizer.tokenize_with_offsets(text_inputs) extracted_wordpieces = _ragged_substr(text_inputs, begin, end) if expected_extracted: self.assertAllEqual(extracted_wordpieces, expected_extracted) else: # The extracted won't have any wordpieces with '##' prefix. Strip them # out. stripped_prefix_flat = string_ops.regex_replace(expected_rt.flat_values, '##', '') stripped_prefix = expected_rt.with_flat_values(stripped_prefix_flat) self.assertAllEqual(extracted_wordpieces, stripped_prefix)
def testWordPieceOpAndVerifyOffsets(self, tokens, expected_subwords, vocab, expected_start=None, expected_limit=None, use_unknown_token=True, unknown_token="[UNK]", token_out_type=dtypes.string, max_bytes_per_word=100, split_unknown_characters=False): for horizon in self._FORWARD_COMPATIBILITY_HORIZONS: with compat.forward_compatibility_horizon(*horizon): tokens_t = ragged_factory_ops.constant(tokens) vocab_table = _CreateTable(vocab) self.evaluate(vocab_table.initializer) tokenizer = WordpieceTokenizer( vocab_table, unknown_token=unknown_token, token_out_type=token_out_type, max_bytes_per_word=max_bytes_per_word, split_unknown_characters=split_unknown_characters, ) subwords_t, begin_t, end_t = tokenizer.tokenize_with_offsets( tokens_t) self.assertAllEqual(subwords_t, expected_subwords) # Verify the indices by performing the following: # - Extract subwords and join them together to form the original tokens. # - Then compare the extracted tokens and original tokens. begin, end = (self.evaluate((begin_t, end_t))) # If expected start/limit offsets were provided, check them explicitly. # Otherwise test the offsets by extracting subwords using token offsets # from the original 'tokens' input. if expected_start is None or expected_limit is None: extracted_tokens = _GetTokensFromWordpieceOffsets( tokens, begin, end) self.assertAllEqual(extracted_tokens, tokens) else: self.assertAllEqual(begin, expected_start) self.assertAllEqual(end, expected_limit)
def test_ragged_input(self, x, expected_indices, expected_values, expected_shape, maxlength=None, minlength=None, binary_count=False, weights=None, axis=-1): x_ragged = ragged_factory_ops.constant(x) y = bincount.sparse_bincount(x_ragged, weights=weights, minlength=minlength, maxlength=maxlength, binary_count=binary_count, axis=axis) self.assertAllEqual(expected_indices, y.indices) self.assertAllEqual(expected_values, y.values) self.assertAllEqual(expected_shape, y.dense_shape)
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 testSimpleSignature(self): int_checker = dispatch.MakeInstanceChecker(int) rt_checker = dispatch.MakeInstanceChecker(ragged_tensor.RaggedTensor) checker = dispatch.PySignatureChecker([(0, int_checker), (2, rt_checker)]) rt = ragged_factory_ops.constant([[1, 2], [3]]) self.check_signatures(checker, [ ((1, 2, rt), True), ((1, 2, 3), False), ((1, 2), False), ((), False), ((5, 'x', rt, None), True), (([5], 'x', rt, None), False), ((5, 'x', [rt], None), False), ]) # pyformat: disable self.assertEqual( repr(checker), '<PySignatureChecker args[0]:int, args[2]:RaggedTensor>')
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([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 test_tf_ragged_tensor(self): def body(i): nonlocal s s = s * 10 + i[0] def set_state(loop_vars): nonlocal s s, = loop_vars s = 0 control_flow.for_stmt(ragged_factory_ops.constant([[1], [2, 4], [3]]), extra_test=None, body=body, get_state=lambda: (s, ), set_state=set_state, symbol_names=('s', ), opts={}) self.assertEqual(s, (123, )) self.assertOpCreated('StatelessWhile')
def test3DimMatrixRagged(self): test_value = ragged_factory_ops.constant( [[['I love Flume!'], ['I don\'t want', 'no scrubs']], [['I love Zhu!', 'Good night']]]) expected_tokens = [[[['I', 'love', 'Flume', '!']], [['I', 'don', '\'', 't', 'want'], ['no', 'scrubs']]], [[['I', 'love', 'Zhu', '!'], ['Good', 'night']]]] expected_offset_starts = [[[[0, 2, 7, 12]], [[0, 2, 5, 6, 8], [0, 3]]], [[[0, 2, 7, 10], [0, 5]]]] expected_offset_limits = [[[[1, 6, 12, 13]], [[1, 5, 6, 7, 12], [2, 9]]], [[[1, 6, 10, 11], [4, 10]]]] tokens = self.tokenizer.tokenize(test_value) self.assertRaggedEqual(tokens, expected_tokens) (tokens, starts, limits) = (self.tokenizer.tokenize_with_offsets(test_value)) self.assertRaggedEqual(tokens, expected_tokens) self.assertRaggedEqual(starts, expected_offset_starts) self.assertRaggedEqual(limits, expected_offset_limits)
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)
def test_hash_ragged_string_input(self): layer = hashing.Hashing(num_bins=2) inp_data = ragged_factory_ops.constant( [['omar', 'stringer', 'marlo', 'wire'], ['marlo', 'skywalker', 'wire']], dtype=dtypes.string) out_data = layer(inp_data) self.assertEqual(out_data.values.numpy().max(), 1) self.assertEqual(out_data.values.numpy().min(), 0) # hash of 'marlo' should be same. self.assertAllClose(out_data[0][2], out_data[1][0]) # hash of 'wire' should be same. self.assertAllClose(out_data[0][3], out_data[1][2]) inp_t = input_layer.Input(shape=(None, ), ragged=True, dtype=dtypes.string) out_t = layer(inp_t) model = training.Model(inputs=inp_t, outputs=out_t) self.assertAllClose(out_data, model.predict(inp_data))
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)
def testNextSentencePredictionExtractor(self, sentences, expected_segment_a, expected_segment_b, expected_labels, random_next_sentence_threshold=0.5, test_description=""): sentences = ragged_factory_ops.constant(sentences) # Set seed and rig the shuffle function to a deterministic reverse function # instead. This is so that we have consistent and deterministic results. random_seed.set_seed(1234) nsp = segment_extractor_ops.NextSentencePredictionExtractor( shuffle_fn=functools.partial(array_ops.reverse, axis=[-1]), random_next_sentence_threshold=random_next_sentence_threshold, ) results = nsp.get_segments(sentences) actual_segment_a, actual_segment_b, actual_labels = results self.assertAllEqual(expected_segment_a, actual_segment_a) self.assertAllEqual(expected_segment_b, actual_segment_b) self.assertAllEqual(expected_labels, actual_labels)
def test3DimMatrixWithRagged2ndDim(self, encoding): test_value = ragged_factory_ops.constant( [[[72, 101, 108, 108, 111], [87, 111, 114, 108, 100]], [[102, 105, 120, 101, 100]], [[72, 121, 112, 101, 114], [119, 111, 114, 100, 115], [99, 117, 98, 101, 46]]], np.int32) expected_value = [ [u"Hello".encode(encoding), u"World".encode(encoding)], [u"fixed".encode(encoding)], [ u"Hyper".encode(encoding), u"words".encode(encoding), u"cube.".encode(encoding) ] ] unicode_encode_op = ragged_string_ops.unicode_encode( test_value, encoding) with self.cached_session(): result = unicode_encode_op.eval() self.assertEqual(unicode_encode_op.ragged_rank, 1) self.assertAllEqual(result.tolist(), expected_value)
def testGetSelectionMask(self, masking_inputs, expected_selected_items, unselectable_ids=None, axis=1, shuffle_fn="", description=""): shuffle_fn = (functools.partial(array_ops.reverse, axis=[-1]) if shuffle_fn == "reverse" else array_ops.identity) masking_inputs = ragged_factory_ops.constant(masking_inputs) item_selector = item_selector_ops.RandomItemSelector( max_selections_per_batch=2, selection_rate=1, shuffle_fn=shuffle_fn, unselectable_ids=unselectable_ids, ) selection_mask = item_selector.get_selection_mask(masking_inputs, axis) selected_items = ragged_array_ops.boolean_mask(masking_inputs, selection_mask) self.assertAllEqual(selected_items, expected_selected_items)
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)
def testDocStringExamples(self): # Sliding window (width=3) across a sequence of tokens data = constant_op.constant( ['one', 'two', 'three', 'four', 'five', 'six']) output = sliding_window_op.sliding_window(data=data, width=3, axis=0) self.assertRaggedEqual( output, [[b'one', b'two', b'three'], [b'two', b'three', b'four'], [b'three', b'four', b'five'], [b'four', b'five', b'six']]) self.assertEqual('Shape: %s -> %s' % (data.shape, output.shape), 'Shape: (6,) -> (4, 3)') # Sliding window (width=2) across the inner dimension of a ragged matrix # containing a batch of token sequences data = ragged_factory_ops.constant([['Up', 'high', 'in', 'the', 'air'], ['Down', 'under', 'water'], ['Away', 'to', 'outer', 'space']]) output = sliding_window_op.sliding_window(data, width=2, axis=-1) self.assertRaggedEqual( output, [[[b'Up', b'high'], [b'high', b'in'], [b'in', b'the'], [b'the', b'air']], [[b'Down', b'under'], [b'under', b'water']], [[b'Away', b'to'], [b'to', b'outer'], [b'outer', b'space']] ]) # pyformat: disable self.assertEqual( 'Shape: %s -> %s' % (data.shape.as_list(), output.shape.as_list()), 'Shape: [3, None] -> [3, None, 2]') # Sliding window across the second dimension of a 3-D tensor containing # batches of sequences of embedding vectors: data = constant_op.constant([[[1, 1, 1], [2, 2, 1], [3, 3, 1], [4, 4, 1], [5, 5, 1]], [[1, 1, 2], [2, 2, 2], [3, 3, 2], [4, 4, 2], [5, 5, 2]]]) output = sliding_window_op.sliding_window(data=data, width=2, axis=1) self.assertRaggedEqual( output, [[[[1, 1, 1], [2, 2, 1]], [[2, 2, 1], [3, 3, 1]], [[3, 3, 1], [4, 4, 1]], [[4, 4, 1], [5, 5, 1]]], [[[1, 1, 2], [2, 2, 2]], [[2, 2, 2], [3, 3, 2]], [[3, 3, 2], [4, 4, 2]], [[4, 4, 2], [5, 5, 2]]]]) self.assertEqual('Shape: %s -> %s' % (data.shape, output.shape), 'Shape: (2, 5, 3) -> (2, 4, 2, 3)')
def SKIP_test_ragged_input_with_padding(self): input_data = get_input_dataset( ragged_factory_ops.constant([[[1, 2, 3, 4, 5]], [[2], [3]]])) expected_output = np.array([[[1., 2., 3., 4., 5.], [-1., -1., -1., -1., -1.]], [[2., -1., -1., -1., -1.], [3., -1., -1., -1., -1.]]]) layers = [ToDense(pad_value=-1), Final()] model = testing_utils.get_model_from_layers(layers, input_shape=(None, None), input_ragged=True, input_dtype=dtypes.int32) model.compile(optimizer="sgd", loss="mse", metrics=["accuracy"], run_eagerly=testing_utils.should_run_eagerly(), experimental_run_tf_function=testing_utils. should_run_tf_function()) output = model.predict(input_data) self.assertAllEqual(output, expected_output)
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)
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)