sparse_test.py 文件源码

def test_sparse_getitem_two_dim_int_slice():
    actual = dense[:, 1]
    res = sparse_getitem(sparse, (slice(None, None, None), 1))
    assert torch.equal(actual, res.to_dense())

    actual = dense[1, :]
    res = sparse_getitem(sparse, (1, slice(None, None, None)))
    assert torch.equal(actual, res.to_dense())