Internal Change

PiperOrigin-RevId: 224642874
2018-12-08 04:40:13 -08:00 · 2018-12-08 04:40:13 -08:00 · e5bb37762e
commit e5bb37762e
parent ee0f5b839d
1 changed files with 7 additions and 0 deletions
--- a/tensorflow/python/ops/ragged/ragged_tensor_value.py
+++ b/tensorflow/python/ops/ragged/ragged_tensor_value.py
@ -98,3 +98,10 @@ class RaggedTensorValue(object):
        values_as_list[self._row_splits[i]:self._row_splits[i + 1]]
        for i in range(len(self._row_splits) - 1)
    ]
+
+  def value_rowids(self, name=None):
+    del name
+    row_lengths = self._row_splits[1:] - self._row_splits[:-1]
+    nrows = self._row_splits.shape[-1] - 1
+    indices = np.arange(nrows)
+    return np.repeat(indices, repeats=row_lengths, axis=0)