mirror of https://github.com/hpcaitech/ColossalAI
46 lines
1.2 KiB
Python
46 lines
1.2 KiB
Python
import torch
|
|
from colossalai.tensor import ColoTensor
|
|
from numpy import allclose
|
|
|
|
|
|
def test_tensor_indexing():
|
|
torch_t = torch.randn(2, 3)
|
|
colo_t = ColoTensor.init_from_torch_tensor(torch_t)
|
|
assert allclose(torch_t[:, 1], colo_t[:, 1].torch_tensor())
|
|
|
|
|
|
def test_lazy_init_tensor():
|
|
lazy_t = ColoTensor(2, 3, dtype=torch.float32, requires_grad=True)
|
|
assert lazy_t._torch_tensor.numel() == 0
|
|
assert lazy_t.numel() == 6 == lazy_t.torch_tensor().numel()
|
|
|
|
|
|
def test_wrapped_tensor_func():
|
|
t_ref = torch.randn(4, 5)
|
|
t = ColoTensor.init_from_torch_tensor(t_ref.clone())
|
|
|
|
# non-func attr
|
|
assert t.is_cuda == t_ref.is_cuda
|
|
|
|
# TODO I don't find out a tensor function which returns None.
|
|
|
|
# return 1 torch.Tensor
|
|
t_abs = t.abs()
|
|
assert isinstance(t_abs, ColoTensor) and torch.equal(t_abs.torch_tensor(), t_ref.abs())
|
|
|
|
# return 1 non-torch.Tensor
|
|
assert t.dim() == t_ref.dim()
|
|
|
|
# return >1 torch.Tensor
|
|
t_split1, t_split2 = t.split(2)
|
|
assert isinstance(t_split1, ColoTensor) and isinstance(t_split2, ColoTensor)
|
|
|
|
|
|
def test_operand():
|
|
t_ref = torch.randn(4, 5)
|
|
t = ColoTensor.init_from_torch_tensor(t_ref.clone())
|
|
|
|
t_ref_res = t_ref + t_ref
|
|
t_res = t + t
|
|
assert torch.allclose(t_ref_res, t_res)
|