def test_expand():
data = np.array([1, 2, 3])
x_gpu = Tensor(data, device='cuda')
x_cpu = Tensor(data, device='cpu')
print(x_gpu.expand(dim=0, copies=5))
print(x_cpu.expand(dim=0, copies=5))
print('norm', x_gpu.norm())
def test_basic_ops(x: Tensor, y: Tensor):
print('x = ', x)
print('y = ', y)
print('x + y = ', x + y)
print('x.sigmoid = ', x.sigmoid())
print('x.tanh = ', x.tanh())
# print('x.softmax = ', x.softmax())
print('2 * x = ', Tensor(2) * x)
print('x * y = ', x * y)
print('x - y = ', x - y)
print('x - 2 = ', x - Tensor(2, device='cuda'))
print('x + 2 = ', x + Tensor(2, device='cuda'))
print('-x = ', -x)
print('x.sum = ', x.sum(0))
print('x.expand(0, 3) = ', x.expand(0, 3))
print('x.expand(1, 3) = ', x.expand(1, 3))