Support indexing 2-axes RaggedTensor, Support slicing for RaggedTensor (

#825)

* Support index 2-axes RaggedTensor, Support slicing for RaggedTensor

* Fix compiling errors

* Fix unit test

* Change RaggedTensor.data to RaggedTensor.values

* Fix style

* Add docs

* Run nightly-cpu when pushing code to nightly-cpu branch

.github/workflows/nightly-cpu.yml

@@ -17,6 +17,9 @@
 name: nightly-cpu
 
 on:
+  push:
+    branches:
+      - nightly-cpu
   schedule:
     # minute (0-59)
     # hour (0-23)

k2/python/csrc/torch/v2/any.cu

@@ -70,15 +70,48 @@ void PybindRaggedAny(py::module &m) {
 
   any.def(
       "__getitem__",
-      [](RaggedAny &self, int32_t i) -> RaggedAny {
-        return self.Index(/*axis*/ 0, i);
+      [](RaggedAny &self, int32_t i) -> py::object {
+        if (self.any.NumAxes() > 2) {
+          RaggedAny ragged = self.Index(/*axis*/ 0, i);
+          return py::cast(ragged);
+        } else {
+          K2_CHECK_EQ(self.any.NumAxes(), 2);
+          Array1<int32_t> row_split = self.any.RowSplits(1).To(GetCpuContext());
+          const int32_t *row_split_data = row_split.Data();
+          int32_t begin = row_split_data[i],
+                  end = row_split_data[i + 1];
+          Dtype t = self.any.GetDtype();
+          FOR_REAL_AND_INT32_TYPES(t, T, {
+            Array1<T> array =
+                self.any.Specialize<T>().values.Arange(begin, end);
+            torch::Tensor tensor = ToTorch(array);
+            return py::cast(tensor);
+          });
+        }
+        // Unreachable code
+        return py::none();
       },
       py::arg("i"), kRaggedAnyGetItemDoc);
 
+  any.def(
+      "__getitem__",
+      [](RaggedAny &self, const py::slice &slice) -> RaggedAny {
+        py::ssize_t start = 0, stop = 0, step = 0, slicelength = 0;
+        if (!slice.compute(self.any.Dim0(), &start, &stop, &step, &slicelength))
+          throw py::error_already_set();
+        int32_t istart = static_cast<int32_t>(start);
+        int32_t istop  = static_cast<int32_t>(stop);
+        int32_t istep  = static_cast<int32_t>(step);
+        K2_CHECK_EQ(istep, 1) << "Only support slicing with step 1, given : "
+                              << istep;
+
+        return self.Arange(/*axis*/ 0, istart, istop);
+      }, py::arg("key"), kRaggedAnyGetItemSliceDoc);
+
   any.def("index",
-          static_cast<RaggedAny (RaggedAny::*)(RaggedAny &, bool)>(
+          static_cast<RaggedAny (RaggedAny::*)(RaggedAny &)>(
               &RaggedAny::Index),
-          py::arg("indexes"), py::arg("remove_axis") = true,
+          py::arg("indexes"),
           kRaggedAnyRaggedIndexDoc);
 
   any.def("index",
@@ -325,8 +358,8 @@ void PybindRaggedAny(py::module &m) {
   // Return the underlying memory of this tensor.
   // No data is copied. Memory is shared.
   any.def_property_readonly(
-      "data", [](RaggedAny &self) -> torch::Tensor { return self.Data(); },
-      kRaggedAnyDataDoc);
+      "values", [](RaggedAny &self) -> torch::Tensor { return self.Data(); },
+      kRaggedAnyValuesDoc);
 
   any.def_property_readonly(
       "shape", [](RaggedAny &self) -> RaggedShape { return self.any.shape; },

k2/python/csrc/torch/v2/doc/any.h

@@ -350,9 +350,6 @@ Select the i-th sublist along axis 0.
 Caution:
   Support for autograd is to be implemented.
 
-Note:
-  It requires that this tensor has at least 3 axes.
-
 >>> import torch
 >>> import k2.ragged as k2r
 >>> a = k2r.RaggedTensor('[ [[1 3] [] [9]]  [[8]] ]')
@@ -363,11 +360,45 @@ Select the i-th sublist along axis 0.
 >>> a[1]
 [ [ 8 ] ]
 
+>>> a = k2r.RaggedTensor('[ [1 3] [9] [8] ]')
+>>> a
+[ [ 1 3 ] [ 9 ] [ 8 ] ]
+>>> a[0]
+tensor([1, 3], dtype=torch.int32)
+>>> a[1]
+tensor([9], dtype=torch.int32)
+
 Args:
   i:
     The i-th sublist along axis 0.
 Returns:
-  Return a new ragged tensor with one fewer axis.
+  Return a new ragged tensor with one fewer axis. If `num_axes == 2`, the
+  return value will be a 1D tensor.
+)doc";
+
+static constexpr const char *kRaggedAnyGetItemSliceDoc = R"doc(
+Slices sublists along axis 0 with the given range. Only support slicing step
+equals to 1.
+
+Caution:
+  Support for autograd is to be implemented.
+
+>>> import torch
+>>> import k2.ragged as k2r
+>>> a = k2r.RaggedTensor('[ [[1 3] [] [9]]  [[8]] [[10 11]] ]')
+>>> a
+[ [ [ 1 3 ] [ ] [ 9 ] ] [ [ 8 ] ] [ [ 10 11 ] ] ]
+>>> a[0:2]
+[ [ [ 1 3 ] [ ] [ 9 ] [ [ 8 ] ] ] ]
+>>> a[1:2]
+[ [ [ 8 ] ] [ [ 10 11 ] ] ]
+
+Args:
+  key:
+    Slice containing integer constants.
+Returns:
+  Return a new ragged tensor with the same axes as original ragged tensor, but
+  only contains the sublists within the range.
 )doc";
 
 static constexpr const char *kRaggedAnyCloneDoc = R"doc(
@@ -644,23 +675,23 @@ device(type='cuda', index=0)
 >>> b.device == torch.device('cuda:0')
 )doc";
 
-static constexpr const char *kRaggedAnyDataDoc = R"doc(
+static constexpr const char *kRaggedAnyValuesDoc = R"doc(
 Return the underlying memory as a 1-D tensor.
 
 >>> import torch
 >>> import k2.ragged as k2r
 >>> a = k2r.RaggedTensor([[1, 2], [], [5], [], [8, 9, 10]])
->>> a.data
+>>> a.values
 tensor([ 1,  2,  5,  8,  9, 10], dtype=torch.int32)
->>> isinstance(a.data, torch.Tensor)
+>>> isinstance(a.values, torch.Tensor)
 True
->>> a.data[0] = -1
+>>> a.values[-2] = -1
 >>> a
 [ [ -1 2 ] [ ] [ 5 ] [ ] [ 8 9 10 ] ]
->>> a.data[3] = -3
+>>> a.values[3] = -3
 >>> a
 [ [ -1 2 ] [ ] [ 5 ] [ ] [ -3 9 10 ] ]
->>> a.data[2] = -2
+>>> a.values[2] = -2
 >>> a
 [ [ -1 2 ] [ ] [ -2 ] [ ] [ -3 9 10 ] ]
 )doc";
@@ -1301,24 +1332,18 @@ Index a ragged tensor with a ragged tensor.
   >>> import k2.ragged as k2r
   >>> src = k2r.RaggedTensor([[10, 11], [12, 13.5]])
   >>> indexes = k2r.RaggedTensor([[0, 1]])
-  >>> src.index(indexes, remove_axis=True)
-  [ [ 10 11 12 13.5 ] ]
-  >>> src.index(indexes, remove_axis=False)
+  >>> src.index(indexes)
   [ [ [ 10 11 ] [ 12 13.5 ] ] ]
   >>> i = k2r.RaggedTensor([[0], [1], [0, 0]])
-  >>> src.index(i, remove_axis=True)
-  [ [ 10 11 ] [ 12 13.5 ] [ 10 11 10 11 ] ]
-  >>> src.index(i, remove_axis=False)
+  >>> src.index(i)
   [ [ [ 10 11 ] ] [ [ 12 13.5 ] ] [ [ 10 11 ] [ 10 11 ] ] ]
 
 **Example 2**:
 
   >>> import k2.ragged as k2r
   >>> src = k2r.RaggedTensor([ [[1, 0], [], [2]], [[], [3], [0, 0, 1]], [[1, 2], [-1]]])
   >>> i = k2r.RaggedTensor([[[0, 2], [1]], [[0]]])
-  >>> src.index(i, remove_axis=True)
-  [ [ [ [ 1 0 2 ] [ 1 2 -1 ] ] [ [ 3 0 0 1 ] ] ] [ [ [ 1 0 2 ] ] ] ]
-  >>> src.index(i, remove_axis=False)
+  >>> src.index(i)
   [ [ [ [ [ 1 0 ] [ ] [ 2 ] ] [ [ 1 2 ] [ -1 ] ] ] [ [ [ ] [ 3 ] [ 0 0 1 ] ] ] ] [ [ [ [ 1 0 ] [ ] [ 2 ] ] ] ] ]
 
 Args:
@@ -1328,13 +1353,6 @@ Index a ragged tensor with a ragged tensor.
     Caution:
       Its dtype has to be ``torch.int32``.
 
-  remove_axis:
-    If ``True``, then we remove the last-but-one axis,
-    which has the effect of appending lists, e.g.
-    if ``self`` is ``[[ 10 11 ] [ 12 13 ]]`` and ``indexes``
-    is ``[[0 1]]`, this function will give us ``[[ 10 11 12 13 ]]``.
-    If ``False`` the answer will have at least 3 axes, e.g., ``[[[10 11]] [12 13]]]`` ,
-    in this case.
 Returns:
   Return indexed tensor.
 )doc";

k2/python/csrc/torch/v2/ragged_any.cu

@@ -560,13 +560,13 @@ torch::optional<torch::Tensor> RaggedAny::Sort(
   return ans;
 }
 
-RaggedAny RaggedAny::Index(RaggedAny &indexes,
-                           bool remove_axis /* = true*/) /*const*/ {
+RaggedAny RaggedAny::Index(RaggedAny &indexes) /*const*/ {
   K2_CHECK_EQ(indexes.any.GetDtype(), kInt32Dtype)
       << "Unsupported dtype: " << TraitsOf(indexes.any.GetDtype()).Name();
 
   DeviceGuard guard(any.Context());
 
+  bool remove_axis = false;
   Dtype t = any.GetDtype();
   FOR_REAL_AND_INT32_TYPES(t, T, {
     return RaggedAny(k2::Index<T>(any.Specialize<T>(),

k2/python/csrc/torch/v2/ragged_any.h

@@ -226,7 +226,7 @@ struct RaggedAny {
                                       bool need_new2old_indexes = false);
 
   /// Wrapper for k2::Index
-  RaggedAny Index(RaggedAny &indexes, bool remove_axis = true) /*const*/;
+  RaggedAny Index(RaggedAny &indexes) /*const*/;
 
   /// Wrapper for k2::Index
   std::pair<RaggedAny, torch::optional<torch::Tensor>> Index(

k2/python/k2/autograd_utils.py

@@ -105,7 +105,7 @@ def backward(ctx, out_fsa_scores_grad: torch.Tensor
                           dtype=torch.float32,
                           device=unused_in_fsa_scores.device,
                           requires_grad=False)
-        _k2.index_add(arc_map.data, expanded, ans)
+        _k2.index_add(arc_map.values, expanded, ans)
 
         return (
             None,  # out_fsa

k2/python/k2/fsa.py

@@ -1394,7 +1394,7 @@ def set_scores_stochastic_(self, scores) -> None:
 
         # Note we use `to` here since `scores` and `self.scores` may not
         # be on the same device.
-        self.scores = ragged_scores.data.to(self.scores.device)
+        self.scores = ragged_scores.values.to(self.scores.device)
 
     def convert_attr_to_ragged_(self, name: str,
                                 remove_eps: bool = True) -> 'Fsa':

k2/python/k2/fsa_algo.py

@@ -466,7 +466,7 @@ def shortest_path(fsa: Fsa, use_double_scores: bool) -> Fsa:
     '''
     entering_arcs = fsa._get_entering_arcs(use_double_scores)
     ragged_arc, ragged_int = _k2.shortest_path(fsa.arcs, entering_arcs)
-    arc_map = ragged_int.data
+    arc_map = ragged_int.values
 
     out_fsa = k2.utils.fsa_from_unary_function_tensor(fsa, ragged_arc, arc_map)
     return out_fsa
@@ -1016,7 +1016,7 @@ def ctc_graph(symbols: Union[List[List[int]], k2.RaggedTensor],
     if isinstance(symbols, k2.RaggedTensor):
         assert device is None
         assert symbols.num_axes == 2
-        symbol_values = symbols.data
+        symbol_values = symbols.values
     else:
         symbol_values = torch.tensor(
             [it for symbol in symbols for it in symbol],

k2/python/k2/utils.py

@@ -169,7 +169,7 @@ def convert_aux_label_to_symbol(
         if end == begin:
             return ':<eps>'
 
-        labels = aux_labels.data[begin:end]
+        labels = aux_labels.values[begin:end]
         ans = []
         for label in labels.tolist():
             if label == -1:
@@ -538,6 +538,7 @@ def fsa_from_unary_function_ragged(src: Fsa,
                 # We currently don't support float ragged attributes
                 assert value.dtype == torch.int32
                 new_value = value.index(arc_map)
+                new_value = new_value.remove_axis(new_value.num_axes - 2)
             setattr(dest, name, new_value)
 
     for name, value in src.named_non_tensor_attr():

k2/python/tests/index_test.py

@@ -145,6 +145,7 @@ def test(self):
                                         device=device)
             ragged_index = k2.RaggedTensor(index_shape, index_values)
             ans = src.index(ragged_index)
+            ans = ans.remove_axis(1)
             expected_row_splits = torch.tensor([0, 5, 5, 5, 9],
                                                dtype=torch.int32,
                                                device=device)
@@ -153,7 +154,7 @@ def test(self):
             expected_values = torch.tensor([1, 2, 4, 5, 6, 3, 3, 1, 2],
                                            dtype=torch.int32,
                                            device=device)
-            self.assertTrue(torch.allclose(ans.data, expected_values))
+            self.assertTrue(torch.allclose(ans.values, expected_values))
 
             # index with tensor
             tensor_index = torch.tensor([0, 3, 2, 1, 2, 1],
@@ -168,7 +169,7 @@ def test(self):
             expected_values = torch.tensor([1, 2, 4, 5, 6, 3, 3],
                                            dtype=torch.int32,
                                            device=device)
-            self.assertTrue(torch.allclose(ans.data, expected_values))
+            self.assertTrue(torch.allclose(ans.values, expected_values))
 
 
 class TestIndexTensorWithRaggedInt(unittest.TestCase):
@@ -203,7 +204,7 @@ def test(self):
             expected_values = torch.tensor([1, 4, 3, 4, 6, 2, 4],
                                            dtype=torch.int32,
                                            device=device)
-            self.assertTrue(torch.allclose(ans.data, expected_values))
+            self.assertTrue(torch.allclose(ans.values, expected_values))
 
 
 if __name__ == '__main__':