Add tests for creating ragged tensors from regular tensors.

Author: csukuangfj

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

@@ -224,6 +224,7 @@ RaggedAny::RaggedAny(torch::Tensor tensor) {
   int32_t ndim = tensor.dim();
   K2_CHECK_GE(ndim, 2) << "Expect a tensor with more than 1-D";
   ContextPtr context = GetContext(tensor);
+  DeviceGuard guard(context);
   std::vector<RaggedShape> shapes;
   shapes.reserve(ndim - 1);
   int32_t dim0 = tensor.size(0);

k2/python/tests/ragged_tensor_test.py

@@ -31,7 +31,6 @@
 
 
 class TestRaggedTensor(unittest.TestCase):
-
     @classmethod
     def setUpClass(cls):
         cls.devices = [torch.device("cpu")]
@@ -65,6 +64,46 @@ def test_create_ragged_tensor_from_string(self):
         assert b.num_axes == 3
         assert b.dim0 == 2
 
+    def test_create_ragged_tensor_from_torch_tensor(self):
+        for device in self.devices:
+            for func in [k2r.create_ragged_tensor, k2r.RaggedTensor]:
+                for dtype in self.dtypes:
+                    a = torch.arange(24, dtype=dtype, device=device).reshape(
+                        2, 3, 4
+                    )
+                    b = func(a)
+
+                    # a is contiguous, so memory is shared
+                    c = a.reshape(-1)
+                    c[0] = 10
+                    assert b.values[0] == 10
+                    b.values[1] = 100
+                    assert c[1] == 100
+
+                    assert b.dtype == dtype
+                    assert b.device == device
+
+                    assert torch.all(torch.eq(c, b.values))
+
+        for device in self.devices:
+            for func in [k2r.create_ragged_tensor, k2r.RaggedTensor]:
+                for dtype in self.dtypes:
+                    a = torch.arange(100, dtype=dtype, device=device).reshape(
+                        10, 10
+                    )[:, ::2]
+                    b = func(a)
+                    assert b.dtype == dtype
+                    assert b.device == device
+
+                    c = a.reshape(-1)
+                    assert torch.all(torch.eq(c, b.values))
+
+                    # a is not contiguous, so memory is copied
+                    c[0] = -10
+                    assert b.values[0] != -10
+                    b.values[1] = -100
+                    assert c[1] != -100
+
     def test_property_values(self):
         a = k2r.RaggedTensor([[1], [2], [], [3, 4]])
         assert torch.all(torch.eq(a.values, torch.tensor([1, 2, 3, 4])))
@@ -128,17 +167,17 @@ def test_sum_with_grad(self):
                 a = a.to(device)
                 a.requires_grad_(True)
                 b = a.sum()
-                expected_sum = torch.tensor([3, 0, 5],
-                                            dtype=dtype,
-                                            device=device)
+                expected_sum = torch.tensor(
+                    [3, 0, 5], dtype=dtype, device=device
+                )
 
                 assert torch.all(torch.eq(b, expected_sum))
 
                 c = b[0] * 10 + b[1] * 20 + b[2] * 30
                 c.backward()
-                expected_grad = torch.tensor([10, 10, 30],
-                                             device=device,
-                                             dtype=dtype)
+                expected_grad = torch.tensor(
+                    [10, 10, 30], device=device, dtype=dtype
+                )
                 assert torch.all(torch.eq(a.grad, expected_grad))
 
     def test_sum_no_grad(self):
@@ -147,26 +186,27 @@ def test_sum_no_grad(self):
                 a = k2r.RaggedTensor([[1, 2], [], [5]], dtype=dtype)
                 a = a.to(device)
                 b = a.sum()
-                expected_sum = torch.tensor([3, 0, 5],
-                                            dtype=dtype,
-                                            device=device)
+                expected_sum = torch.tensor(
+                    [3, 0, 5], dtype=dtype, device=device
+                )
 
                 assert torch.all(torch.eq(b, expected_sum))
 
     def test_getitem(self):
         for device in self.devices:
             for dtype in self.dtypes:
-                a = k2r.RaggedTensor("[ [[1 2] [] [10]] [[3] [5]] ]",
-                                     dtype=dtype)
+                a = k2r.RaggedTensor(
+                    "[ [[1 2] [] [10]] [[3] [5]] ]", dtype=dtype
+                )
                 a = a.to(device)
                 b = a[0]
-                expected = k2r.RaggedTensor("[[1 2] [] [10]]",
-                                            dtype=dtype).to(device)
+                expected = k2r.RaggedTensor("[[1 2] [] [10]]", dtype=dtype).to(
+                    device
+                )
                 assert b == expected
 
                 b = a[1]
-                expected = k2r.RaggedTensor("[[3] [5]]",
-                                            dtype=dtype).to(device)
+                expected = k2r.RaggedTensor("[[3] [5]]", dtype=dtype).to(device)
                 assert b == expected
 
     def test_getstate_2axes(self):
@@ -177,9 +217,9 @@ def test_getstate_2axes(self):
                 assert isinstance(b, tuple)
                 assert len(b) == 3
                 # b contains (row_splits, "row_ids1", values)
-                b_0 = torch.tensor([0, 2, 3, 3],
-                                   dtype=torch.int32,
-                                   device=device)
+                b_0 = torch.tensor(
+                    [0, 2, 3, 3], dtype=torch.int32, device=device
+                )
                 b_1 = "row_ids1"
                 b_2 = a.values
 
@@ -190,18 +230,19 @@ def test_getstate_2axes(self):
     def test_getstate_3axes(self):
         for device in self.devices:
             for dtype in self.dtypes:
-                a = k2r.RaggedTensor("[[[1 2] [3] []] [[4] [5 6]]]",
-                                     dtype=dtype).to(device)
+                a = k2r.RaggedTensor(
+                    "[[[1 2] [3] []] [[4] [5 6]]]", dtype=dtype
+                ).to(device)
                 b = a.__getstate__()
                 assert isinstance(b, tuple)
                 assert len(b) == 5
                 # b contains (row_splits1, "row_ids1", row_splits2,
                 # "row_ids2", values)
                 b_0 = torch.tensor([0, 3, 5], dtype=torch.int32, device=device)
                 b_1 = "row_ids1"
-                b_2 = torch.tensor([0, 2, 3, 3, 4, 6],
-                                   dtype=torch.int32,
-                                   device=device)  # noqa
+                b_2 = torch.tensor(
+                    [0, 2, 3, 3, 4, 6], dtype=torch.int32, device=device
+                )  # noqa
                 b_3 = "row_ids2"
                 b_4 = a.values
 
@@ -255,7 +296,8 @@ def test_tot_size_3axes(self):
             for dtype in self.dtypes:
                 a = k2r.RaggedTensor(
                     "[ [[1 2 3] [] [5 8]] [[] [1 5 9 10 -1] [] [] []] ]",
-                    dtype=dtype)
+                    dtype=dtype,
+                )
                 a = a.to(device)
 
                 assert a.tot_size(0) == 2