def test_sparse_getitem_two_dim_int_slice(self):
actual = self.dense[:, 1]
res = sparse_getitem(self.sparse, (slice(None, None, None), 1))
self.assertTrue(torch.equal(actual, res.to_dense()))
actual = self.dense[1, :]
res = sparse_getitem(self.sparse, (1, slice(None, None, None)))
self.assertTrue(torch.equal(actual, res.to_dense()))
def test_sparse_repeat_1d(self):
sparse_1d = sparse_getitem(self.sparse, 1)
actual = sparse_1d.to_dense().repeat(3, 1)
res = sparse_repeat(sparse_1d, 3, 1)
self.assertTrue(torch.equal(actual, res.to_dense()))
actual = sparse_1d.to_dense().repeat(2, 3)
res = sparse_repeat(sparse_1d, 2, 3)
self.assertTrue(torch.equal(actual, res.to_dense()))
def test_sparse_getitem_two_dim_slice(self):
actual = self.dense[2:4, 1:3]
res = sparse_getitem(self.sparse, (slice(2, 4), slice(1, 3)))
self.assertTrue(torch.equal(actual, res.to_dense()))
def test_sparse_getitem_two_dim_int(self):
actual = self.dense[2, 1]
res = sparse_getitem(self.sparse, (2, 1))
self.assertEqual(actual, res)
def test_sparse_getitem_one_dim_slice(self):
actual = self.dense[2:4]
res = sparse_getitem(self.sparse, slice(2, 4))
self.assertTrue(torch.equal(actual, res.to_dense()))
def test_sparse_getitem_one_dim_int(self):
actual = self.dense[3]
res = sparse_getitem(self.sparse, 3)
self.assertTrue(torch.equal(actual, res.to_dense()))
