def testRaggedMap( self, fn, elems, expected_output, expected_ragged_rank=None, result_ragged_rank=None, elems_ragged_rank=None, dtype=dtypes.int64, result_dtype=None, infer_shape=False, ): elems = ragged.constant(elems, dtype, elems_ragged_rank) output = ragged.map_fn(fn=fn, elems=elems, dtype=result_dtype, infer_shape=infer_shape) expected_rt = ragged.constant(expected_output, ragged_rank=expected_ragged_rank) with self.test_session(): if ragged.is_ragged(expected_output): self.assertEqual(output.ragged_rank, expected_rt.ragged_rank) output_values = self.evaluate(output) self.assertAllEqual(expected_output, output_values.tolist())
def testRaggedMapOnStructure_RaggedOutputs(self): batman = ragged.constant([[1, 2, 3], [4], [5, 6, 7]]) # [[10, 20, 30], [40], [50, 60, 70]] robin = ragged.map_inner_values(mo.multiply, batman, 10) features = {'batman': batman, 'robin': robin} def _increment(f): return { 'batman': ragged.add(f['batman'], 1), 'robin': ragged.add(f['robin'], 1), } output = ragged.map_fn( fn=_increment, elems=features, infer_shape=False, dtype={ 'batman': ragged.RaggedTensorType(dtype=dtypes.int32, ragged_rank=1), 'robin': ragged.RaggedTensorType(dtype=dtypes.int32, ragged_rank=1) }, ) with self.test_session(): self.assertAllEqual(output['batman'].eval().tolist(), [[2, 3, 4], [5], [6, 7, 8]]) self.assertAllEqual(output['robin'].eval().tolist(), [[11, 21, 31], [41], [51, 61, 71]])
def testMismatchRaggedRank2(self): elems = ragged.constant([[1, 2, 3], [4, 5], [6, 7]]) fn = lambda x: ragged.from_row_starts(x, [0]) with self.assertRaisesWithLiteralMatch( ValueError, r'The declared ragged rank (10) mismatches the result (1)'): _ = ragged.map_fn( fn, elems, dtype=ragged.RaggedTensorType(dtype=dtypes.int64, ragged_rank=10))
def testMismatchRaggedRank(self): elems = ragged.constant([[[1, 2, 3]], [[4, 5], [6, 7]]]) fn = lambda x: ragged.reduce_sum(x, axis=0) with self.assertRaisesWithLiteralMatch( ValueError, r'The declared ragged rank (23) mismatches the result (1)'): _ = ragged.map_fn( fn, elems, dtype=ragged.RaggedTensorType(dtype=dtypes.int64, ragged_rank=23))
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.RaggedTensor.from_sparse(s) id_t2 = ragged.map_fn( lambda x: x, t2, ) self.assertRaggedEqual(id_t2, [[0, 5], [0, 4]])
def testRaggedMapOnStructure(self): batman = ragged.constant([[1, 2, 3], [4], [5, 6, 7]]) # [[10, 20, 30], [40], [50, 60, 70]] robin = ragged.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_fn( fn=_reduce_sum_from_all, elems=features, dtype=dtypes.int32, ) self.assertRaggedEqual(output, [66, 44, 198])
def testRaggedMapOnStructure(self): batman = ragged.constant([[1, 2, 3], [4], [5, 6, 7]]) # [[10, 20, 30], [40], [50, 60, 70]] robin = ragged.map_inner_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_fn( fn=_reduce_sum_from_all, elems=features, dtype=dtypes.int32, ) with self.test_session(): self.assertAllEqual(output.eval().tolist(), [66, 44, 198])
def testBatchGather(self): tokens = ragged.constant([['hello', '.', 'there'], ['merhaba'], ['bonjour', '.', 'ca va', '?']]) indices = ragged.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_fn(gather, data, dtype=ragged.RaggedTensorType(dtype=dtypes.string, ragged_rank=1), infer_shape=False) self.assertRaggedEqual( out, [[b'hello', b'there'], [b'merhaba'], [b'bonjour', b'ca va']])
def testBatchGather(self): tokens = ragged.constant([['hello', '.', 'there'], ['merhaba'], ['bonjour', '.', 'ca va', '?']]) indices = ragged.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_fn( gather, data, dtype=ragged.RaggedTensorType(dtype=dtypes.string, ragged_rank=1), infer_shape=False) with self.test_session(): self.assertAllEqual( self.evaluate(out).tolist(), [[b'hello', b'there'], [b'merhaba'], [b'bonjour', b'ca va']])
def testZip(self): x = ragged.constant( [[10, 20], [30, 40], [50, 60], [70], [80, 90, 100]], dtypes.int64) y = array_ops.expand_dims(mo.range(x.nrows(), dtype=dtypes.int64), axis=1) def _zip(foo): y_val, x_val = foo bar = backend.tile(y_val, array_ops.shape(x_val)) return array_ops.stack([bar, x_val], axis=1) output = ragged.map_fn(_zip, (y, x), dtype=ragged.RaggedTensorType( dtype=dtypes.int64, ragged_rank=1), infer_shape=False) self.assertRaggedEqual( output, [[[0, 10], [0, 20]], [[1, 30], [1, 40]], [[2, 50], [2, 60]], [[3, 70]], [[4, 80], [4, 90], [4, 100]]])
def testZip(self): x = ragged.constant([[10, 20], [30, 40], [50, 60], [70], [80, 90, 100]], dtypes.int64) y = array_ops.expand_dims( mo.range(ragged.nrows(x), dtype=dtypes.int64), axis=1) def _zip(foo): y_val, x_val = foo bar = backend.tile(y_val, array_ops.shape(x_val)) return array_ops.stack([bar, x_val], axis=1) output = ragged.map_fn( _zip, (y, x), dtype=ragged.RaggedTensorType(dtype=dtypes.int64, ragged_rank=1), infer_shape=False) with self.test_session(): result = self.evaluate(output).tolist() self.assertAllEqual( result, [[[0, 10], [0, 20]], [[1, 30], [1, 40]], [[2, 50], [2, 60]], [[3, 70]], [[4, 80], [4, 90], [4, 100]]])
def testRaggedMap( self, fn, elems, expected_output, expected_ragged_rank=None, result_ragged_rank=None, elems_ragged_rank=None, dtype=dtypes.int64, result_dtype=None, infer_shape=False, ): elems = ragged.constant(elems, dtype, elems_ragged_rank) output = ragged.map_fn(fn=fn, elems=elems, dtype=result_dtype, infer_shape=infer_shape) expected_rt = ragged.constant(expected_output, ragged_rank=expected_ragged_rank) self.assertRaggedEqual(expected_rt, output)
def testRaggedMap( self, fn, elems, expected_output, expected_ragged_rank=None, result_ragged_rank=None, elems_ragged_rank=None, dtype=dtypes.int64, result_dtype=None, infer_shape=False, ): elems = ragged.constant(elems, dtype, elems_ragged_rank) output = ragged.map_fn( fn=fn, elems=elems, dtype=result_dtype, infer_shape=infer_shape) expected_rt = ragged.constant( expected_output, ragged_rank=expected_ragged_rank) with self.test_session(): if ragged.is_ragged(expected_output): self.assertEqual(output.ragged_rank, expected_rt.ragged_rank) output_values = self.evaluate(output) self.assertAllEqual(expected_output, output_values.tolist())