Remove KJT GPU test (#2869)

Summary:
Pull Request resolved: #2869

KJT's GPU test has been moved to its own test file.

Reviewed By: TroyGarden

Differential Revision: D72401855

fbshipit-source-id: 2ac75c73d8e89d0627cdd2f7f93f90ec01757be9

Author: jd7-tr

torchrec/sparse/tests/test_jagged_tensor_gpu.py

@@ -11,11 +11,7 @@
 import unittest
 
 import torch
-from torchrec.sparse.jagged_tensor import (
-    _regroup_keyed_tensors,
-    KeyedJaggedTensor,
-    KeyedTensor,
-)
+from torchrec.sparse.jagged_tensor import _regroup_keyed_tensors, KeyedTensor
 from torchrec.sparse.tests.utils import build_groups, build_kts
 from torchrec.test_utils import skip_if_asan_class
 
@@ -115,195 +111,3 @@ def test_regroup_backward(self) -> None:
 
         torch.allclose(actual_kt_0_grad, expected_kt_0_grad)
         torch.allclose(actual_kt_1_grad, expected_kt_1_grad)
-
-
-@skip_if_asan_class
-class TestKeyedJaggedTensorGPU(unittest.TestCase):
-    def setUp(self) -> None:
-        super().setUp()
-        self.device = torch.cuda.current_device()
-
-    # pyre-ignore
-    @unittest.skipIf(
-        torch.cuda.device_count() <= 0,
-        "Not enough GPUs, this test requires at least one GPUs",
-    )
-    def test_permute(self) -> None:
-        values = torch.tensor(
-            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], device=self.device
-        )
-        lengths = torch.tensor([0, 2, 0, 1, 1, 1, 0, 3, 0], device=self.device)
-        keys = ["index_0", "index_1", "index_2"]
-
-        jag_tensor = KeyedJaggedTensor.from_lengths_sync(
-            values=values,
-            keys=keys,
-            lengths=lengths,
-        )
-        indices = [1, 0, 2]
-        permuted_jag_tensor = jag_tensor.permute(indices)
-
-        self.assertEqual(permuted_jag_tensor.keys(), ["index_1", "index_0", "index_2"])
-        self.assertEqual(
-            permuted_jag_tensor.offset_per_key(),
-            [0, 3, 5, 8],
-        )
-        self.assertEqual(
-            permuted_jag_tensor.values().tolist(),
-            [3.0, 4.0, 5.0, 1.0, 2.0, 6.0, 7.0, 8.0],
-        )
-        self.assertEqual(
-            permuted_jag_tensor.lengths().tolist(), [1, 1, 1, 0, 2, 0, 0, 3, 0]
-        )
-        self.assertEqual(permuted_jag_tensor.weights_or_none(), None)
-
-    # pyre-ignore
-    @unittest.skipIf(
-        torch.cuda.device_count() <= 0,
-        "Not enough GPUs, this test requires at least one GPUs",
-    )
-    def test_permute_vb(self) -> None:
-        values = torch.tensor(
-            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], device=self.device
-        )
-        lengths = torch.tensor([1, 0, 1, 3, 0, 1, 0, 2, 0], device=self.device)
-        keys = ["index_0", "index_1", "index_2"]
-        stride_per_key_per_rank = [[2], [4], [3]]
-
-        jag_tensor = KeyedJaggedTensor.from_lengths_sync(
-            values=values,
-            keys=keys,
-            lengths=lengths,
-            stride_per_key_per_rank=stride_per_key_per_rank,
-        )
-
-        indices = [1, 0, 2]
-        permuted_jag_tensor = jag_tensor.permute(indices)
-
-        self.assertEqual(permuted_jag_tensor.keys(), ["index_1", "index_0", "index_2"])
-        self.assertEqual(
-            permuted_jag_tensor.offset_per_key(),
-            [0, 5, 6, 8],
-        )
-        self.assertEqual(
-            permuted_jag_tensor.values().tolist(),
-            [2.0, 3.0, 4.0, 5.0, 6.0, 1.0, 7.0, 8.0],
-        )
-        self.assertEqual(
-            permuted_jag_tensor.lengths().tolist(), [1, 3, 0, 1, 1, 0, 0, 2, 0]
-        )
-        self.assertEqual(permuted_jag_tensor.weights_or_none(), None)
-
-    # pyre-ignore
-    @unittest.skipIf(
-        torch.cuda.device_count() <= 0,
-        "Not enough GPUs, this test requires at least one GPUs",
-    )
-    def test_permute_vb_duplicate(self) -> None:
-        values = torch.tensor(
-            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], device=self.device
-        )
-        lengths = torch.tensor([1, 0, 1, 3, 0, 1, 0, 2, 0], device=self.device)
-        keys = ["index_0", "index_1", "index_2"]
-        stride_per_key_per_rank = [[2], [4], [3]]
-
-        jag_tensor = KeyedJaggedTensor.from_lengths_sync(
-            values=values,
-            keys=keys,
-            lengths=lengths,
-            stride_per_key_per_rank=stride_per_key_per_rank,
-        )
-
-        indices = [1, 1, 0, 0, 2, 2]
-        permuted_jag_tensor = jag_tensor.permute(indices)
-
-        self.assertEqual(
-            permuted_jag_tensor.keys(),
-            ["index_1", "index_1", "index_0", "index_0", "index_2", "index_2"],
-        )
-        self.assertTrue(
-            torch.equal(
-                permuted_jag_tensor.values().cpu(),
-                torch.Tensor(
-                    [
-                        2.0,
-                        3.0,
-                        4.0,
-                        5.0,
-                        6.0,
-                        2.0,
-                        3.0,
-                        4.0,
-                        5.0,
-                        6.0,
-                        1.0,
-                        1.0,
-                        7.0,
-                        8.0,
-                        7.0,
-                        8.0,
-                    ]
-                ),
-            )
-        )
-        self.assertTrue(
-            torch.equal(
-                permuted_jag_tensor.lengths().cpu(),
-                torch.IntTensor([1, 3, 0, 1, 1, 3, 0, 1, 1, 0, 1, 0, 0, 2, 0, 0, 2, 0]),
-            )
-        )
-        self.assertEqual(permuted_jag_tensor.weights_or_none(), None)
-
-    # pyre-ignore
-    @unittest.skipIf(
-        torch.cuda.device_count() <= 0,
-        "Not enough GPUs, this test requires at least one GPUs",
-    )
-    def test_permute_duplicates(self) -> None:
-        values = torch.tensor(
-            [1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], device=self.device
-        )
-        lengths = torch.tensor([0, 2, 0, 1, 1, 1, 0, 3, 0], device=self.device)
-        keys = ["index_0", "index_1", "index_2"]
-
-        jag_tensor = KeyedJaggedTensor.from_lengths_sync(
-            values=values,
-            keys=keys,
-            lengths=lengths,
-        )
-
-        indices = [1, 0, 2, 1, 1]
-        permuted_jag_tensor = jag_tensor.permute(indices)
-
-        self.assertEqual(
-            permuted_jag_tensor.keys(),
-            ["index_1", "index_0", "index_2", "index_1", "index_1"],
-        )
-        self.assertEqual(
-            permuted_jag_tensor.offset_per_key(),
-            [0, 3, 5, 8, 11, 14],
-        )
-        self.assertEqual(
-            permuted_jag_tensor.values().tolist(),
-            [
-                3.0,
-                4.0,
-                5.0,
-                1.0,
-                2.0,
-                6.0,
-                7.0,
-                8.0,
-                3.0,
-                4.0,
-                5.0,
-                3.0,
-                4.0,
-                5.0,
-            ],
-        )
-        self.assertEqual(
-            permuted_jag_tensor.lengths().tolist(),
-            [1, 1, 1, 0, 2, 0, 0, 3, 0, 1, 1, 1, 1, 1, 1],
-        )
-        self.assertEqual(permuted_jag_tensor.weights_or_none(), None)