xref: /aosp_15_r20/external/tensorflow/tensorflow/python/ops/ragged/ragged_to_tensor_op_test.py (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ==============================================================================
15"""Tests for ragged.to_tensor."""
16
17import random
18
19from absl.testing import parameterized
20import numpy as np
21
22from tensorflow.python.client import session
23from tensorflow.python.eager import backprop
24from tensorflow.python.eager import context
25from tensorflow.python.framework import constant_op
26from tensorflow.python.framework import dtypes
27from tensorflow.python.framework import errors
28from tensorflow.python.framework import ops
29from tensorflow.python.framework import tensor_shape
30from tensorflow.python.framework import test_util
31from tensorflow.python.ops import array_ops
32from tensorflow.python.ops import gradients_impl
33from tensorflow.python.ops.ragged import ragged_factory_ops
34from tensorflow.python.ops.ragged import ragged_tensor
35from tensorflow.python.ops.ragged.ragged_tensor import RaggedTensor
36from tensorflow.python.platform import benchmark
37from tensorflow.python.platform import googletest
38from tensorflow.python.util import nest
39
40
41def make_placeholder(t):
42  return array_ops.placeholder_with_default(t, None)
43
44
45def rebuild_ragged_tensor_with_value_rowids(rt, feed_dict=None, sess=None):
46  """Returns a copy of `rt`, built using `from_value_rowids`.
47
48  This ensures that RaggedTensor._cached_value_rowids is populated, which
49  triggers a different code-path for converting ragged tensors to tensors.
50
51  If `feed_dict` and `sess` are specified, then build the new `RaggedTensor`
52  using placeholder tensors, and populate a feed dictionary that can be used
53  to feed the placeholders.
54
55  Args:
56    rt: The RaggedTensor to copy.
57    feed_dict: If specified, then build the new `RaggedTensor` using
58      placeholders, and populate this dict with entries to feed those
59      placeholders.
60    sess: A session used to evaluate tensors; required if feed_dict is
61      specified.
62
63  Returns:
64    A copy of `rt`, built using `from_value_rowids`.
65  """
66  if isinstance(rt, ragged_tensor.RaggedTensor):
67    values = rebuild_ragged_tensor_with_value_rowids(rt.values, feed_dict, sess)
68    rowids = rt.value_rowids()
69    nrows = rt.nrows()
70    if feed_dict is not None:
71      rowids_ph = make_placeholder(rowids)
72      nrows_ph = make_placeholder(nrows)
73      feed_dict[rowids_ph] = sess.run(rowids)
74      feed_dict[nrows_ph] = sess.run(nrows)
75      rowids, nrows = rowids_ph, nrows_ph
76    return ragged_tensor.RaggedTensor.from_value_rowids(values, rowids, nrows)
77  else:
78    if feed_dict is not None:
79      rt_ph = make_placeholder(rt)
80      feed_dict[rt_ph] = sess.run(rt)
81      rt = rt_ph
82    return rt
83
84
85@test_util.run_all_in_graph_and_eager_modes
86class RaggedTensorToTensorOpTest(test_util.TensorFlowTestCase,
87                                 parameterized.TestCase):
88
89  def testDocStringExamples(self):
90    """Example from ragged_to_tensor.__doc__."""
91    rt = ragged_factory_ops.constant([[9, 8, 7], [], [6, 5], [4]])
92    dt = rt.to_tensor()
93    self.assertAllEqual(dt, [[9, 8, 7], [0, 0, 0], [6, 5, 0], [4, 0, 0]])
94
95  @parameterized.named_parameters(
96      # Simple 2D ragged tensors (with one ragged dimension)
97      {
98          'testcase_name': 'shape_2xN',
99          'rt_input': [[0, 1, 2], [], [3]],
100          'expected': [[0, 1, 2], [0, 0, 0], [3, 0, 0]]
101      },
102      {
103          'testcase_name': 'shape_2xN_default_0D',
104          'rt_input': [[0, 1, 2], [], [3]],
105          'default': 5,
106          'expected': [[0, 1, 2], [5, 5, 5], [3, 5, 5]]
107      },
108      {
109          'testcase_name': 'empty_first_row',
110          'rt_input': [[], [], [3, 4], []],
111          'expected': [[0, 0], [0, 0], [3, 4], [0, 0]]
112      },
113      {
114          'testcase_name': 'empty_last_row',
115          'rt_input': [[0, 1, 2], [], [3], []],
116          'expected': [[0, 1, 2], [0, 0, 0], [3, 0, 0], [0, 0, 0]]
117      },
118      {
119          'testcase_name': 'shape_4xN',
120          'rt_input': [[1, 2, 3], [], [4], [5, 6]],
121          'expected': [[1, 2, 3], [0, 0, 0], [4, 0, 0], [5, 6, 0]]
122      },
123      {
124          'testcase_name': 'shape_4xN_default_0D',
125          'rt_input': [[1, 2, 3], [], [4], [5, 6]],
126          'default': 9,
127          'expected': [[1, 2, 3], [9, 9, 9], [4, 9, 9], [5, 6, 9]]
128      },
129      {
130          'testcase_name': 'shape_2xN_already_dense',
131          'rt_input': [[6, 7, 8], [9, 10, 11]],
132          'expected': [[6, 7, 8], [9, 10, 11]],
133      },
134      {
135          'testcase_name': 'shape_2xN_string_already_dense',
136          'rt_input': [[b'a', b'b', b'c'],
137                       [b'd', b'e', b'antidisestablishmentarianism']],
138          'ragged_rank': 1,
139          'expected': [[b'a', b'b', b'c'],
140                       [b'd', b'e', b'antidisestablishmentarianism']],
141      },
142      # 3D ragged tensors with two ragged dimensions
143      {
144          'testcase_name': 'shape_4xNxM',
145          'rt_input': [[[1, 2], [], [3, 4]], [], [[5]], [[6, 7], [8]]],
146          'expected': [
147              [[1, 2], [0, 0], [3, 4]],  #
148              [[0, 0], [0, 0], [0, 0]],  #
149              [[5, 0], [0, 0], [0, 0]],  #
150              [[6, 7], [8, 0], [0, 0]],  #
151          ]
152      },
153      {
154          'testcase_name': 'shape_4xNxM_default_0D',
155          'rt_input': [[[1, 2], [], [3, 4]], [], [[5]], [[6, 7], [8]]],
156          'default': 9,
157          'expected': [
158              [[1, 2], [9, 9], [3, 4]],  #
159              [[9, 9], [9, 9], [9, 9]],  #
160              [[5, 9], [9, 9], [9, 9]],  #
161              [[6, 7], [8, 9], [9, 9]],  #
162          ]
163      },
164      {
165          'testcase_name': 'shape_1xNx1_default_0D',
166          'rt_input': [[[1], [2], [3]]],
167          'ragged_rank': 1,
168          'default': 0,
169          'expected': [[[1], [2], [3]]],
170      },
171      {
172          'testcase_name': 'shape_2xNx2_already_dense',
173          'rt_input': [[[6, 7], [8, 9], [10, 11]],
174                       [[12, 13], [14, 15], [16, 17]]],
175          'ragged_rank': 1,
176          'expected': [[[6, 7], [8, 9], [10, 11]],
177                       [[12, 13], [14, 15], [16, 17]]],
178      },
179      {
180          'testcase_name': 'shape_2xNx2_already_dense_default_1D',
181          'rt_input': [[[6, 7], [8, 9], [10, 11]],
182                       [[12, 13], [14, 15], [16, 17]]],
183          'ragged_rank': 1,
184          'default': [31, 32],
185          'expected': [[[6, 7], [8, 9], [10, 11]],
186                       [[12, 13], [14, 15], [16, 17]]],
187      },
188      {
189          'testcase_name': 'shape_2xNx2_string_already_dense',
190          'rt_input': [[[b'a', b'b'], [b'c', b'd'], [b'e', b'f']],
191                       [[b'g', b'jalapeno'], [b'kangaroo', b'llama'],
192                        [b'manzana', b'nectar']]],
193          'ragged_rank': 1,
194          'expected': [[[b'a', b'b'], [b'c', b'd'], [b'e', b'f']],
195                       [[b'g', b'jalapeno'], [b'kangaroo', b'llama'],
196                        [b'manzana', b'nectar']]],
197      },
198      # 3D ragged tensors with one ragged dimension
199      {
200          'testcase_name': 'shape_4xNx1_default_1D',
201          'rt_input': [[[1], [2], [3]], [], [[4]], [[5], [6]]],
202          'ragged_rank': 1,
203          'default': [9],
204          'expected': [[[1], [2], [3]],
205                       [[9], [9], [9]],
206                       [[4], [9], [9]],
207                       [[5], [6], [9]]]
208      },
209      {
210          'testcase_name': 'shape_2xNx2_default_0D',
211          'rt_input': [[[6, 7], [8, 9], [10, 11]],
212                       [[12, 13], [14, 15]]],
213          'ragged_rank': 1,
214          'default': 2,
215          'expected': [[[6, 7], [8, 9], [10, 11]],
216                       [[12, 13], [14, 15], [2, 2]]],
217      },
218      {
219          'testcase_name': 'shape_2xNx2_default_1D',
220          'rt_input': [[[6, 7], [8, 9], [10, 11]],
221                       [[12, 13], [14, 15]]],
222          'ragged_rank': 1,
223          'default': [2, 3],
224          'expected': [[[6, 7], [8, 9], [10, 11]],
225                       [[12, 13], [14, 15], [2, 3]]],
226      },
227      # 4D ragged tensors with 3 ragged dimensions
228      {
229          'testcase_name': 'shape_1xNxMxK_default_0D',
230          'rt_input': [[[[1], [2]], [], [[3]]]],
231          'default': 9,
232          'expected': [[[[1], [2]], [[9], [9]], [[3], [9]]]],
233      },
234      # Broadcast default
235      {
236          'testcase_name': 'shape_2xNx2x2_default_2x1',
237          'rt_input': [[[[1, 2], [3, 4]]], []],
238          'ragged_rank': 1,
239          'default': [[5], [6]],
240          'expected': [[[[1, 2], [3, 4]]],
241                       [[[5, 5], [6, 6]]]],
242      },
243      {
244          'testcase_name': 'shape_2xNx2x2_default_1x2',
245          'rt_input': [[[[1, 2], [3, 4]]], []],
246          'ragged_rank': 1,
247          'default': [[5, 6]],
248          'expected': [[[[1, 2], [3, 4]]],
249                       [[[5, 6], [5, 6]]]],
250      },
251      # Explicit shape
252      {
253          'testcase_name': 'shape_4xN_with_crop',
254          'rt_input': [[0, 1, 2, 3], [], [4], []],
255          'shape': [2, 3],
256          'expected': [[0, 1, 2], [0, 0, 0]],
257      },
258      {
259          'testcase_name': 'shape_2xN_with_pad',
260          'rt_input': [[1, 2], [3]],
261          'shape': [3, 3],
262          'expected': [[1, 2, 0], [3, 0, 0], [0, 0, 0]],
263      },
264      {
265          'testcase_name': 'shape_4xN_with_crop_and_pad',
266          'rt_input': [[0, 1, 2, 3], [], [4], []],
267          'shape': [2, 8],
268          'expected': [[0, 1, 2, 3, 0, 0, 0, 0],
269                       [0, 0, 0, 0, 0, 0, 0, 0]],
270      },
271      {
272          'testcase_name': 'shape_4xN_with_tuple_shape',
273          'rt_input': [[0, 1, 2, 3], [], [4], []],
274          'shape': (2, 3),
275          'expected': [[0, 1, 2], [0, 0, 0]],
276      },
277      {
278          'testcase_name': 'shape_4xN_with_tensorshape_shape',
279          'rt_input': [[0, 1, 2, 3], [], [4], []],
280          'shape': tensor_shape.TensorShape([2, 3]),
281          'expected': [[0, 1, 2], [0, 0, 0]],
282      },
283      {
284          'testcase_name': 'shape_4xN_with_partial_shape',
285          'rt_input': [[0, 1, 2, 3], [], [4], []],
286          'shape': tensor_shape.TensorShape([2, None]),
287          'expected': [[0, 1, 2, 3], [0, 0, 0, 0]],
288      },
289      # Empty tensors
290      {
291          'testcase_name': 'shape_0xN',
292          'rt_input': [],
293          'ragged_rank': 1,
294          'expected': [],
295          'expected_shape': [0, 0],
296      },
297      {
298          'testcase_name': 'shape_0xNxM',
299          'rt_input': [],
300          'ragged_rank': 2,
301          'expected': [],
302          'expected_shape': [0, 0, 0],
303      },
304      # {
305      #     'testcase_name': 'shape_0xNx2',
306      #     'rt_input': [],
307      #     'ragged_rank': 1,
308      #     'inner_shape': [2],
309      #     'expected': [],
310      #     'expected_shape': [0, 0, 2],
311      # },
312      {
313          'testcase_name': 'shape_2xN_empty',
314          'rt_input': [[], []],
315          'expected': [[], []],
316          'expected_shape': [2, 0],
317      },
318  )  # pyformat: disable
319  def testRaggedTensorToTensor(self,
320                               rt_input,
321                               expected,
322                               ragged_rank=None,
323                               inner_shape=None,
324                               default=None,
325                               shape=None,
326                               expected_shape=None):
327    rt1 = ragged_factory_ops.constant(
328        rt_input, ragged_rank=ragged_rank, inner_shape=inner_shape)
329    rt2 = rebuild_ragged_tensor_with_value_rowids(rt1)
330    for rt in [rt1, rt2]:
331      for use_placeholder in [False, True]:
332        if use_placeholder:
333          if default is not None:
334            default = make_placeholder(default)
335          rt = nest.map_structure(make_placeholder, rt, expand_composites=True)
336        dt = rt.to_tensor(default_value=default, shape=shape)
337        self.assertIsInstance(dt, ops.Tensor)
338        self.assertEqual(rt.dtype, dt.dtype)
339        if shape is not None:
340          self.assertTrue(dt.shape.is_compatible_with(shape))
341        else:
342          self.assertTrue(dt.shape.is_compatible_with(rt.shape))
343        if expected_shape is not None:
344          expected = np.ndarray(expected_shape, buffer=np.array(expected))
345        self.assertAllEqual(dt, expected)
346
347  @parameterized.parameters([
348      {
349          'rt_input': [[1, 2, 3]],
350          'default': 'a',
351          'error_type': TypeError,
352          'error': r'Expected int32|Cannot convert',
353      },
354      {
355          'rt_input': [[1, 2, 3]],
356          'default': [0],
357          'error': r'default_value\.shape=.* and '
358                   r'rt_input\.flat_values\.shape=.* are incompatible: '
359                   r'default_value\.rank = 1  must be less than '
360                   r'rt_input\.flat_values\.rank = 1'
361      },
362      {
363          'rt_input': [[[1, 2], [3, 4]], [[5, 6]]],
364          'ragged_rank': 1,
365          'default': [7, 8, 9],
366          'error': r'default_value\.shape.* and '
367                   r'rt_input\.flat_values\.shape.* are incompatible: '
368                   r'default_value\.shape\[-1\] = 3 but '
369                   r'rt_input\.flat_values\.shape\[-1\] = 2'
370      },
371      {
372          'rt_input': [[1, 2, 3]],
373          'shape': [3, 3, 3],
374          'error': r'rt_input\.shape and shape=\[.,.,.\] are incompatible: '
375                   r'rt_input\.rank = 2 but shape\.rank = 3'
376      },
377      {
378          'rt_input': [[[1, 2, 3]]],
379          'ragged_rank': 1,
380          'shape': [1, 1, 4],
381          'error': r'rt_input\.shape and shape=\[1,1,4\] are incompatible: '
382                   r'rt_input\.shape\[2\] = 3 but shape\[2\] = 4'
383      },
384  ])
385  def testError(self,
386                rt_input,
387                error,
388                error_type=(ValueError, errors.InvalidArgumentError),
389                default=None,
390                ragged_rank=None,
391                shape=None):
392
393    rt = ragged_factory_ops.constant(rt_input, ragged_rank=ragged_rank)
394    with self.assertRaisesRegex(error_type, error):
395      self.evaluate(rt.to_tensor(default_value=default, shape=shape))
396    rt_placeholder = nest.map_structure(
397        make_placeholder, rt, expand_composites=True)
398    with self.assertRaisesRegex(error_type, error):
399      self.evaluate(
400          rt_placeholder.to_tensor(default_value=default, shape=shape))
401
402  def test_shape_limit_shape_is_tensor(self):
403    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
404    actual = input_data.to_tensor(
405        shape=constant_op.constant([2, 3], dtype=dtypes.int64))
406    self.assertAllEqual(actual, [[0, 1, 2], [0, 0, 0]])
407    self.assertEqual(actual.shape.as_list(), [2, 3])
408
409  def test_shape_limit_shape_is_tensor_unknown_rank(self):
410    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
411    actual = input_data.to_tensor(
412        shape=constant_op.constant(-1, dtype=dtypes.int64))
413    self.assertAllEqual(
414        actual, [[0, 1, 2, 3], [0, 0, 0, 0], [4, 0, 0, 0], [0, 0, 0, 0]])
415    self.assertTrue(actual.shape.is_compatible_with([4, 4]))
416
417  def test_shape_limit_shape_is_tensor_unknown_dim(self):
418    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
419    actual = input_data.to_tensor(
420        shape=constant_op.constant([2, -1], dtype=dtypes.int64))
421    self.assertAllEqual(actual, [[0, 1, 2, 3], [0, 0, 0, 0]])
422    self.assertTrue(actual.shape.is_compatible_with([2, None]))
423
424  def test_shape_limit_shape_is_tensor_int32(self):
425    input_data = ragged_factory_ops.constant([[0, 1, 2, 3], [], [4], []])
426    actual = input_data.to_tensor(
427        shape=constant_op.constant([2, 3], dtype=dtypes.int32))
428    self.assertAllEqual(actual, [[0, 1, 2], [0, 0, 0]])
429    self.assertEqual(actual.shape.as_list(), [2, 3])
430
431  def test_shape_expand_first_dim(self):
432    input_data = ragged_factory_ops.constant([[0, 1, 2], [], [3]])
433    actual = input_data.to_tensor(shape=[4, 4])
434    self.assertAllEqual(
435        actual, [[0, 1, 2, 0], [0, 0, 0, 0], [3, 0, 0, 0], [0, 0, 0, 0]])
436    self.assertEqual(actual.shape.as_list(), [4, 4])
437
438  def test_value_transposed(self):
439    # Check that transposed data is not an issue.
440    my_value = array_ops.transpose(
441        constant_op.constant([[0, 1, 2, 3], [4, 5, 6, 7]]))
442    input_data = RaggedTensor.from_value_rowids(
443        values=my_value,
444        value_rowids=constant_op.constant([0, 1, 2, 3], dtype=dtypes.int64),
445        nrows=constant_op.constant(4, dtype=dtypes.int64),
446        validate=True)
447    self.assertAllEqual(input_data, [[[0, 4]], [[1, 5]], [[2, 6]], [[3, 7]]])
448
449  def test_broadcast_default(self):
450    # The dense dimension here is 2 x 2
451    input_data = ragged_factory_ops.constant([[[[1, 2], [3, 4]]], []],
452                                             ragged_rank=1)
453    # This placeholder has a 2 x 1 dimension.
454    default_value = make_placeholder([[5], [6]])
455    actual = input_data.to_tensor(default_value=default_value)
456    expected = [[[[1, 2], [3, 4]]], [[[5, 5], [6, 6]]]]
457    self.assertAllEqual(actual, expected)
458
459  def test_broadcast_default_no_placeholder(self):
460    input_data = ragged_factory_ops.constant([[[[1, 2], [3, 4]]], []],
461                                             ragged_rank=1)
462    # default_value has a 2 x 1 dimension.
463    default_value = constant_op.constant([[5], [6]], shape=None)
464    actual = input_data.to_tensor(default_value=default_value)
465    expected = [[[[1, 2], [3, 4]]], [[[5, 5], [6, 6]]]]
466    self.assertAllEqual(actual, expected)
467
468  def test_shape_expand_second_dim(self):
469    input_data = ragged_factory_ops.constant([[0, 1, 2], [], [3], []])
470    actual = input_data.to_tensor(shape=[3, 4])
471    self.assertAllEqual(actual, [[0, 1, 2, 0], [0, 0, 0, 0], [3, 0, 0, 0]])
472
473  @parameterized.parameters(
474      ([2, 3, 4], None, [2, 3, 4]),
475      ([2, 3, 4], [None, None, None], [2, 3, 4]),
476      ([2, 3, 4], [None, 3, None], [2, 3, 4]),
477      ([2, 3, 4], [None, 3, 4], [2, 3, 4]),
478      ([2, 3, 4], [2, 3, 4], [2, 3, 4]),
479      )
480  def test_preserve_shape_roundtrip(
481      self, input_shape, to_tensor_shape, expected_shape):
482    tensor = array_ops.zeros(input_shape)
483    ragged_from_tensor = RaggedTensor.from_tensor(tensor, ragged_rank=2)
484    recovered_tensor = ragged_from_tensor.to_tensor(shape=to_tensor_shape)
485    self.assertAllEqual(tensor.shape.as_list(), expected_shape)
486    self.assertAllEqual(ragged_from_tensor.shape.as_list(), expected_shape)
487    self.assertAllEqual(recovered_tensor.shape.as_list(), expected_shape)
488
489  def test_empty_tensor_with_shape(self):
490    input_data = RaggedTensor.from_value_rowids(
491        values=constant_op.constant([], dtype=dtypes.int64),
492        value_rowids=constant_op.constant([], dtype=dtypes.int64),
493        nrows=constant_op.constant(2, dtype=dtypes.int64),
494        validate=True)
495    actual = input_data.to_tensor(default_value=3, shape=[2, 3])
496    self.assertAllEqual(actual, [[3, 3, 3], [3, 3, 3]])
497
498  # pylint: disable=bad-whitespace
499  @parameterized.named_parameters([
500      dict(
501          testcase_name = '2d_default_shape',
502          shape         = None,
503          rt_value      = [[1, 2, 3], [4], [5, 6]],
504          rt_grad       = [[9, 8, 7], [6], [3, 2]],
505          default_value = 0,
506          default_grad  = sum([5, 4, 1]),
507          output_value  = [[1, 2, 3], [4, 0, 0], [5, 6, 0]],
508          output_grad   = [[9, 8, 7], [6, 5, 4], [3, 2, 1]]),
509      dict(
510          testcase_name = '2d_pad',
511          shape         = [4, 4],
512          rt_value      = [[1, 2, 3], [4], [5, 6]],
513          rt_grad       = [[9, 8, 7], [5], [1, 0]],
514          default_value = 0,
515          default_grad  = sum([6, 4, 3, 2, 1, 2, 3, 4, 5, 6]),
516          output_value  = [
517              [1, 2, 3, 0], [4, 0, 0, 0], [5, 6, 0, 0], [0, 0, 0, 0]],
518          output_grad   = [
519              [9, 8, 7, 6], [5, 4, 3, 2], [1, 0, 1, 2], [3, 4, 5, 6]]),
520      dict(
521          testcase_name = '2d_pad_and_crop',
522          shape         = [5, 3],
523          rt_value      = [[1, 2, 3], [4], [5, 6, 7, 8, 9], [8]],
524          rt_grad       = [[9, 8, 7], [6], [3, 2, 1, 0, 0], [2]],
525          default_value = 0,
526          default_grad  = sum([5, 4, 3, 4, 5, 6, 7]),
527          output_value  = [
528              [1, 2, 3], [4, 0, 0], [5, 6, 7], [8, 0, 0], [0, 0, 0]],
529          output_grad   = [
530              [9, 8, 7], [6, 5, 4], [3, 2, 1], [2, 3, 4], [5, 6, 7]]),
531      dict(
532          testcase_name = '3d_rrank_2',
533          shape         = [2, 2, 2],
534          rt_value      = [[[9, 8, 7], [6]], [[5, 4]]],
535          rt_grad       = [[[1, 2, 0], [3]], [[5, 6]]],
536          default_value = 3,
537          default_grad  = sum([4, 7, 8]),
538          output_value  = [[[9, 8], [6, 3]], [[5, 4], [3, 3]]],
539          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
540      dict(
541          testcase_name = '3d_rrank_1_with_0d_default',
542          ragged_rank   = 1,
543          shape         = [2, 2, 2],
544          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
545          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
546          default_value = 3,
547          default_grad  = sum([7, 8]),
548          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 3]]],
549          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
550      dict(
551          testcase_name = '3d_rrank_1_with_1d_default',
552          ragged_rank   = 1,
553          shape         = [2, 2, 2],
554          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
555          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
556          default_value = [3, 2],
557          default_grad  = [7, 8],
558          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 2]]],
559          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
560      dict(
561          testcase_name = '3d_rrank_1_with_1d_broadcast_default',
562          ragged_rank   = 1,
563          shape         = [2, 2, 2],
564          rt_value      = [[[9, 8], [7, 6]], [[5, 4]]],
565          rt_grad       = [[[1, 2], [3, 4]], [[5, 6]]],
566          default_value = [3],
567          default_grad  = [7 + 8],
568          output_value  = [[[9, 8], [7, 6]], [[5, 4], [3, 3]]],
569          output_grad   = [[[1, 2], [3, 4]], [[5, 6], [7, 8]]]),
570      dict(
571          testcase_name = '4d_rrank_1_with_2d_default',
572          ragged_rank   = 1,
573          shape         = [3, 3, 2, 1],
574          rt_value      = [[[[9], [8]], [[7], [6]]], [[[5], [4]]]],
575          rt_grad       = [[[[1], [2]], [[3], [4]]], [[[7], [8]]]],
576          default_value = [[3], [2]],
577          default_grad  = [[5 + 9 + 2 + 4 + 6 + 8], [6 + 1 + 3 + 5 + 7 + 9]],
578          output_value  = [[[[9], [8]], [[7], [6]], [[3], [2]]],
579                           [[[5], [4]], [[3], [2]], [[3], [2]]],
580                           [[[3], [2]], [[3], [2]], [[3], [2]]]],
581          output_grad   = [[[[1], [2]], [[3], [4]], [[5], [6]]],
582                           [[[7], [8]], [[9], [1]], [[2], [3]]],
583                           [[[4], [5]], [[6], [7]], [[8], [9]]]]),
584      dict(
585          testcase_name = '4d_rrank_1_with_with_0d_default',
586          ragged_rank   = 1,
587          shape         = [3, 3, 2, 1],
588          rt_value      = [[[[9], [8]], [[7], [6]]], [[[5], [4]]]],
589          rt_grad       = [[[[1], [2]], [[3], [4]]], [[[7], [8]]]],
590          default_value = 3,
591          default_grad  = 5 + 9 + 2 + 4 + 6 + 8 + 6 + 1 + 3 + 5 + 7 + 9,
592          output_value  = [[[[9], [8]], [[7], [6]], [[3], [3]]],
593                           [[[5], [4]], [[3], [3]], [[3], [3]]],
594                           [[[3], [3]], [[3], [3]], [[3], [3]]]],
595          output_grad   = [[[[1], [2]], [[3], [4]], [[5], [6]]],
596                           [[[7], [8]], [[9], [1]], [[2], [3]]],
597                           [[[4], [5]], [[6], [7]], [[8], [9]]]]),
598      dict(
599          testcase_name = 'zero_size',
600          shape         = [0, 0],
601          rt_value      = [[9, 8], [7, 6, 5], [4]],
602          rt_grad       = [[0, 0], [0, 0, 0], [0]],
603          default_value = 3,
604          default_grad  = 0,
605          output_value  = [],
606          output_grad   = [])
607  ])  # pyformat: disable
608  def test_gradient(self,
609                    shape,
610                    rt_value,
611                    rt_grad,
612                    default_value,
613                    default_grad,
614                    output_value,
615                    output_grad,
616                    ragged_rank=None):
617    """Tests that ragged_to_dense generates the right gradient.
618
619    Args:
620      shape: The `shape` arg for `ragged_to_dense`.
621      rt_value: The `rt_input` arg for `ragged_to_dense`.
622      rt_grad: The expected gradient for `rt_value`.  Corresponds 1:1 with
623        `rt_value`.
624      default_value: The `default_value` arg for `ragged_to_dense`.
625      default_grad: The expected gradient for `default_value`.  Corresponds 1:1
626        with `default_value`.
627      output_value: The expected output of `ragged_to_dense`.
628      output_grad: The gradient for the output (used to generate the gradients
629        `rt_grad` and `default_grad`).  Corresponds 1:1 with `output_value`.
630      ragged_rank: Ragged rank for `rt_value`.
631    """
632    rt_value = ragged_factory_ops.constant(
633        rt_value, dtype=dtypes.float32, ragged_rank=ragged_rank)
634    rt_grad = ragged_factory_ops.constant(
635        rt_grad, dtype=dtypes.float32, ragged_rank=ragged_rank)
636    default_value = constant_op.constant(default_value, dtype=dtypes.float32)
637    default_grad = constant_op.constant(default_grad, dtype=dtypes.float32)
638    output_value = constant_op.constant(
639        output_value, dtype=dtypes.float32, shape=shape)
640    output_grad = constant_op.constant(
641        output_grad, dtype=dtypes.float32, shape=shape)
642    shape = tensor_shape.as_shape(shape)
643
644    # There are different code paths for ragged_to_dense, depending on whether
645    # the RaggedTensor was created from row_splits or value_rowids.  Make sure
646    # that we test both.
647    for partition_type in ['row_splits', 'value_rowids']:
648      rt_val = self.rt_with_partition_type(rt_value, partition_type)
649      if context.executing_eagerly():
650        self._test_gradient_helper(rt_val, default_value, shape, output_grad,
651                                   output_value, rt_grad, default_grad)
652      else:
653        # There are different code paths when computing the gradient for
654        # default_value, depending on whether shape info is statically
655        # available; make sure that we test all code paths.
656        for shape_info in ['known', 'unknown_dims', 'unknown_rank']:
657          rt_val = self.wrap_in_placeholder(rt_val, shape_info)
658          default_val = self.wrap_in_placeholder(default_value, shape_info)
659          shape_val = self.wrap_in_placeholder(shape, shape_info)
660          self._test_gradient_helper(rt_val, default_val, shape_val,
661                                     output_grad, output_value, rt_grad,
662                                     default_grad)
663
664  def _test_gradient_helper(self, rt_val, default_val, shape_val, output_grad,
665                            expected_output_val, expected_rt_grad,
666                            expected_default_grad):
667    if context.executing_eagerly():
668      with backprop.GradientTape() as tape:
669        tape.watch([rt_val, default_val])
670        out = rt_val.to_tensor(default_val, shape=shape_val)
671        actual_rt_grad, actual_default_grad = tape.gradient(
672            out, (rt_val, default_val), output_gradients=output_grad)
673    else:
674      out = rt_val.to_tensor(default_val, shape=shape_val)
675      actual_rt_grad, actual_default_grad = gradients_impl.gradients(
676          ys=out, xs=(rt_val, default_val), grad_ys=output_grad)
677
678    self.assertAllClose(out, expected_output_val)
679    self.assertIsInstance(actual_rt_grad, RaggedTensor)
680    self.assertAllClose(actual_rt_grad, expected_rt_grad)
681    self.assertAllClose(actual_default_grad, expected_default_grad)
682
683  def rt_with_partition_type(self, rt, partition_type):
684    if isinstance(rt, ops.Tensor):
685      return rt
686    if partition_type == 'row_splits':
687      return rt
688    if partition_type == 'value_rowids':
689      return ragged_tensor.RaggedTensor.from_value_rowids(
690          self.rt_with_partition_type(rt.values, partition_type),
691          rt.value_rowids(), rt.nrows())
692    raise AssertionError('Unexpected partition_type %r' % partition_type)
693
694  def wrap_in_placeholder(self, arg, shape_info):
695    """Wraps `arg` in a placeholder to limit static shape info.
696
697    Args:
698      arg: The value to wrap.  A Tensor, RaggedTensor, or TensorShape.
699      shape_info: One of ['known', 'unknown_dims', 'unknown_rank'].
700
701    Returns:
702      * If shape_info is 'known': returns `arg`.
703      * If shape_info is 'unknown_dims': returns a placeholder wrapping `arg`
704        where the dimension sizes are unknown.  If `arg` is a TensorShape,
705        then convert it to a vector first.  If `arg` is a RaggedTensor, then
706        wrap the flat_values.
707      * If shape_info is 'unknown_rank': returns a placeholder wrapping `arg`
708        where the rank is unknown.  If `arg` is a TensorShape, then convert it
709        to a vector first.  If `arg` is a RaggedTensor, then wrap the
710        flat_values.
711    """
712    if shape_info == 'known':
713      return arg
714    if isinstance(arg, ragged_tensor.RaggedTensor):
715      return arg.with_flat_values(
716          self.wrap_in_placeholder(arg.flat_values, shape_info))
717    if isinstance(arg, tensor_shape.TensorShape):
718      if arg.ndims is None:
719        return arg
720      arg = constant_op.constant(arg.as_list())
721    if shape_info == 'unknown_rank':
722      return array_ops.placeholder_with_default(arg, None)
723    if shape_info == 'unknown_dims':
724      return array_ops.placeholder_with_default(arg, [None] * arg.shape.rank)
725    raise AssertionError('Unexpected shape_info %r' % shape_info)
726
727  def test_shape_is_list_including_tensor_element(self):
728    rt = ragged_factory_ops.constant([[1, 2, 3], [4], [5, 6]])
729    result = rt.to_tensor(shape=[2, constant_op.constant(2)])
730    self.assertAllEqual(result, [[1, 2], [4, 0]])
731
732
733class RaggedToDenseBenchmark(googletest.Benchmark):
734
735  # Configurations to test.  See `run_benchmark` for config param docs.
736  CONFIGS = [
737      {'shape': [10, 10]},
738      {'shape': [10, 1000]},
739      {'shape': [1000, 10]},
740      {'shape': [1000, 10], 'fill': [1, 0.95]},  # Mostly full.
741      {'shape': [1000, 10], 'fill': [1, 0.05]},  # Mostly empty.
742      {'shape': [1000, 10], 'dtype': dtypes.string},
743      {'shape': [1000, 10], 'dtype': dtypes.int64},
744      {'shape': [100, 100]},
745      {'shape': [50, 50, 32]},
746      {'shape': [100, 100, 100], 'min_iters': 100},
747      {'shape': [1000, 1000], 'min_iters': 100},
748      {'shape': [10, 10, 10, 10, 10]},
749      {'shape': [10, 10, 10, 10, 10], 'ragged_rank': 1},
750      {'shape': [10, 10, 10, 10, 10], 'ragged_rank': 2},
751      {'shape': [50, 50, 32], 'ragged_rank': 1, 'default_shape': [32]},
752      {'shape': [200, 50, 32], 'ragged_rank': 1, 'default_shape': [32]}
753  ]  # pyformat: disable
754
755  def run_benchmark(self,
756                    shape=(100, 100),
757                    ragged_rank=None,
758                    dtype=dtypes.float32,
759                    fill=None,
760                    default_shape=(),
761                    output_shape=None,
762                    min_iters=1000):
763    """Run a benchmark with the specified configuration parameters.
764
765    Args:
766      shape: Bounding box for the input ragged tensor.
767      ragged_rank: Ragged rank for the input ragged tensor.  Defaults to
768        `len(shape)-1`.
769      dtype: Data type for the input ragged tensor.
770      fill: How full each dimension should be (0-1).  Corresponds 1:1 with
771        `shape`.  Defaults to 0.8 for each dimension.
772      default_shape: Shape for the default (padding) value.
773      output_shape: Output shape -- ragged tensor will be padded or cropped to
774        this shape.
775      min_iters: Minimum iterations for benchmark.
776    """
777    if ragged_rank is None:
778      ragged_rank = len(shape) - 1
779    if fill is None:
780      fill = [0.8 for _ in shape]
781
782    # Build the inputs for the op.
783    rt_input = self._generateRaggedTensor(shape, ragged_rank, dtype, fill)
784    default_value = constant_op.constant(
785        self._generateRaggedTensor(default_shape, 0, dtype), dtype=dtype)
786
787    mbs = np.prod(shape) / (2**20)
788    with session.Session(config=benchmark.benchmark_config()) as sess:
789      extras = {
790          'shape': shape,
791          'ragged_rank': ragged_rank,
792          'dtype': dtype,
793          'fill': fill,
794          'default_shape': default_shape
795      }
796      rt = ragged_factory_ops.constant(rt_input, dtype, ragged_rank=ragged_rank)
797
798      # Inputs for with_splits:
799      splits_rt_placeholder = ragged_factory_ops.placeholder(
800          dtype, ragged_rank, shape[ragged_rank + 1:])
801      splits_feed_dict = {splits_rt_placeholder: sess.run(rt)}
802
803      # Inputs for with_rowids:
804      rowids_feed_dict = {}
805      rowids_rt_placeholder = rebuild_ragged_tensor_with_value_rowids(
806          rt, rowids_feed_dict, sess)
807
808      # Common arguments for benchmarks:
809      run_op_benchmark_kwargs = dict(
810          sess=sess,
811          store_memory_usage=True,
812          min_iters=min_iters,
813          burn_iters=max(5, min_iters // 10),
814          mbs=mbs,
815          extras=extras)
816
817      ragged_to_tensor_with_splits = splits_rt_placeholder.to_tensor(
818          default_value=default_value)
819      self.run_op_benchmark(
820          op_or_tensor=ragged_to_tensor_with_splits.op,
821          name='ragged_to_tensor_with_splits',
822          feed_dict=splits_feed_dict,
823          **run_op_benchmark_kwargs)
824
825      ragged_to_tensor_with_rowids = rowids_rt_placeholder.to_tensor(
826          default_value=default_value)
827      self.run_op_benchmark(
828          op_or_tensor=ragged_to_tensor_with_rowids.op,
829          name='ragged_to_tensor_with_rowids',
830          feed_dict=rowids_feed_dict,
831          **run_op_benchmark_kwargs)
832
833  def _generateRaggedTensor(self, shape, ragged_rank, dtype, fill=None, axis=0):
834    if axis == len(shape):
835      value = random.random()
836      if dtype == dtypes.string:
837        value = str(value)
838      if dtype.is_integer:
839        value = int(value * 1000)
840      return value
841    if axis == 0 or axis > ragged_rank:
842      slice_size = shape[axis]
843    else:
844      slice_size = (np.random.geometric(fill[axis], shape[axis]) == 1).sum()
845    return [
846        self._generateRaggedTensor(shape, ragged_rank, dtype, fill, axis + 1)
847        for _ in range(slice_size)
848    ]
849
850  def benchmark_ragged_to_dense(self):
851    random.seed(5)
852    for config in self.CONFIGS:
853      self.run_benchmark(**config)
854
855
856if __name__ == '__main__':
857  googletest.main()
858