mirror of https://github.com/Nonannet/copapy.git
258 lines
8.7 KiB
Python
258 lines
8.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Basic tests for the tensor class."""
|
|
|
|
import copapy as cp
|
|
|
|
def test_tensor_basic():
|
|
# Test 1: Create a scalar tensor
|
|
print("Test 1: Scalar tensor")
|
|
t0 = cp.tensor(42)
|
|
print(f"Scalar tensor: {t0}")
|
|
print(f"Shape: {t0.shape}, ndim: {t0.ndim}")
|
|
assert t0.shape == ()
|
|
assert t0 == 42
|
|
print()
|
|
|
|
# Test 2: Create a 1D tensor from list
|
|
print("Test 2: 1D tensor")
|
|
t1 = cp.tensor([1, 2, 3, 4, 5])
|
|
print(f"1D tensor: shape={t1.shape}, ndim={t1.ndim}")
|
|
print(f"Elements: {[t1[i] for i in range(len(t1))]}")
|
|
assert t1.shape == (5,)
|
|
assert t1.ndim == 1
|
|
assert t1[0] == 1
|
|
print()
|
|
|
|
# Test 3: Create a 2D tensor (matrix)
|
|
print("Test 3: 2D tensor")
|
|
t2 = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
|
print(f"2D tensor: shape={t2.shape}, ndim={t2.ndim}")
|
|
print(f"Element [0,1]: {t2[0, 1]}")
|
|
print(f"Row 1: {t2[1]}")
|
|
assert t2.shape == (2, 3)
|
|
assert t2.ndim == 2
|
|
assert t2[0, 1] == 2
|
|
|
|
assert t2[1][2] == 6
|
|
print()
|
|
|
|
# Test 4: Create a 3D tensor
|
|
print("Test 4: 3D tensor")
|
|
t3 = cp.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
|
|
print(f"3D tensor: shape={t3.shape}, ndim={t3.ndim}")
|
|
print(f"Element [0,1,0]: {t3[0, 1, 0]}")
|
|
assert t3.shape == (2, 2, 2)
|
|
assert t3.ndim == 3
|
|
assert t3[0, 1, 0] == 3
|
|
print()
|
|
|
|
# Test 6: Broadcasting with scalar
|
|
print("Test 6: Broadcasting with scalar")
|
|
t = cp.tensor([1.0, 2.0, 3.0])
|
|
result = t * 2.0
|
|
print(f"tensor * 2.0: shape={result.shape}")
|
|
print(f"Elements: {[result[i] for i in range(len(result))]}")
|
|
assert result.shape == (3,)
|
|
assert result[0] == 2.0
|
|
assert result[1] == 4.0
|
|
print()
|
|
|
|
# Test 6b: Broadcasting with different dimensions
|
|
print("Test 6b: Broadcasting with different dimensions")
|
|
# 2D tensor + 1D tensor
|
|
t2d = cp.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
|
|
t1d = cp.tensor([10.0, 20.0, 30.0])
|
|
result_2d_1d = t2d + t1d
|
|
print(f"2D tensor {t2d.shape} + 1D tensor {t1d.shape} = shape {result_2d_1d.shape}")
|
|
print(f"Elements: {[[result_2d_1d[i, j] for j in range(3)] for i in range(2)]}")
|
|
assert result_2d_1d.shape == (2, 3)
|
|
assert result_2d_1d[0, 0] == 11.0
|
|
assert result_2d_1d[1, 2] == 36.0
|
|
|
|
# 3D tensor + 2D tensor
|
|
t3d = cp.tensor([[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]])
|
|
t2d_broadcast = cp.tensor([[100.0, 200.0], [300.0, 400.0]])
|
|
result_3d_2d = t3d + t2d_broadcast
|
|
print(f"3D tensor {t3d.shape} + 2D tensor {t2d_broadcast.shape} = shape {result_3d_2d.shape}")
|
|
assert result_3d_2d.shape == (2, 2, 2)
|
|
assert result_3d_2d[0, 0, 0] == 101.0
|
|
assert result_3d_2d[1, 1, 1] == 408.0
|
|
|
|
# 3D tensor + 1D tensor
|
|
t1d_broadcast = cp.tensor([1.0, 2.0])
|
|
result_3d_1d = t3d + t1d_broadcast
|
|
print(f"3D tensor {t3d.shape} + 1D tensor {t1d_broadcast.shape} = shape {result_3d_1d.shape}")
|
|
assert result_3d_1d.shape == (2, 2, 2)
|
|
assert result_3d_1d[0, 0, 0] == 2.0
|
|
assert result_3d_1d[1, 1, 1] == 10.0
|
|
print()
|
|
|
|
# 3D tensor + vector
|
|
t1d_broadcast = cp.vector([1.0, 2.0])
|
|
result_3d_1d = t3d + t1d_broadcast
|
|
print(f"3D tensor {t3d.shape} + 1D tensor {t1d_broadcast.shape} = shape {result_3d_1d.shape}")
|
|
assert result_3d_1d.shape == (2, 2, 2)
|
|
assert result_3d_1d[0, 0, 0] == 2.0
|
|
assert result_3d_1d[1, 1, 1] == 10.0
|
|
print()
|
|
|
|
# Test 6c: Element-wise operations with different dimensions
|
|
print("Test 6c: Element-wise operations with different dimensions")
|
|
a2d = cp.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
b2d = cp.tensor([[2.0, 3.0], [4.0, 5.0]])
|
|
c2d = a2d * b2d
|
|
print(f"2D * 2D: shape={c2d.shape}")
|
|
print(f"Elements: {[[c2d[i, j] for j in range(2)] for i in range(2)]}")
|
|
assert c2d.shape == (2, 2)
|
|
assert c2d[0, 0] == 2.0
|
|
assert c2d[1, 1] == 20.0
|
|
|
|
# 3D - 2D
|
|
t3d_sub = cp.tensor([[[10.0, 20.0], [30.0, 40.0]], [[50.0, 60.0], [70.0, 80.0]]])
|
|
t2d_sub = cp.tensor([[1.0, 2.0], [3.0, 4.0]])
|
|
result_sub = t3d_sub - t2d_sub
|
|
print(f"3D - 2D: shape={result_sub.shape}")
|
|
assert result_sub.shape == (2, 2, 2)
|
|
assert result_sub[0, 0, 0] == 9.0
|
|
assert result_sub[1, 1, 1] == 76.0
|
|
print()
|
|
|
|
# Test 7: Reshape
|
|
print("Test 7: Reshape")
|
|
t = cp.tensor([1, 2, 3, 4, 5, 6])
|
|
print(f"Original: shape={t.shape}")
|
|
t_reshaped = t.reshape(2, 3)
|
|
print(f"Reshaped to (2, 3): shape={t_reshaped.shape}")
|
|
print(f"Element [1,2]: {t_reshaped[1, 2]}")
|
|
assert t_reshaped.shape == (2, 3)
|
|
assert t_reshaped[1, 2] == 6
|
|
assert t_reshaped[0, 0] == 1
|
|
print()
|
|
|
|
# Test 8: Flatten
|
|
print("Test 8: Flatten")
|
|
t = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
|
flat = t.flatten()
|
|
print(f"Original: shape={t.shape}")
|
|
print(f"Flattened: shape={flat.shape}")
|
|
print(f"Elements: {[flat[i] for i in range(len(flat))]}")
|
|
assert flat.shape == (6,)
|
|
assert flat[0] == 1
|
|
assert flat[5] == 6
|
|
print()
|
|
|
|
# Test 9: Transpose
|
|
print("Test 9: Transpose")
|
|
t = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
|
print(f"Original: shape={t.shape}")
|
|
t_t = t.transpose()
|
|
print(f"Transposed: shape={t_t.shape}")
|
|
print(f"Element [2,1]: {t_t[2, 1]}")
|
|
assert t_t.shape == (3, 2)
|
|
assert t_t[2, 1] == 6
|
|
print()
|
|
|
|
# Test 10: Sum operations
|
|
print("Test 10: Sum operations")
|
|
t = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
|
print(f"Original: shape={t.shape}")
|
|
total = t.sum()
|
|
print(f"Sum all: {total}")
|
|
sum_axis0 = t.sum(axis=0)
|
|
print(f"Sum along axis 0: shape={sum_axis0.shape}")
|
|
sum_axis1 = t.sum(axis=1)
|
|
print(f"Sum along axis 1: shape={sum_axis1.shape}")
|
|
assert total == 21
|
|
assert sum_axis0.shape == (3,)
|
|
assert sum_axis0[0] == 5
|
|
assert sum_axis1.shape == (2,)
|
|
assert sum_axis1[1] == 15
|
|
print()
|
|
|
|
# Test 10b: Sum with multiple axes and keepdims
|
|
print("Test 10b: Sum with multiple axes and keepdims")
|
|
t3d = cp.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
|
|
print(f"Original 3D tensor: shape={t3d.shape}")
|
|
|
|
# Sum along multiple axes
|
|
sum_axes_0_2 = t3d.sum(axis=(0, 2))
|
|
print(f"Sum along axes (0, 2): shape={sum_axes_0_2.shape}")
|
|
assert sum_axes_0_2.shape == (2,), f"Expected (2,), got {sum_axes_0_2.shape}"
|
|
assert sum_axes_0_2[0] == 1 + 2 + 5 + 6 # Elements from [0,:,*] and [1,:,*]
|
|
assert sum_axes_0_2[1] == 3 + 4 + 7 + 8
|
|
print(f"Values: {[sum_axes_0_2[i] for i in range(len(sum_axes_0_2))]}")
|
|
|
|
# Sum with keepdims
|
|
sum_keepdims = t3d.sum(axis=1, keepdims=True)
|
|
print(f"Sum along axis 1 with keepdims: shape={sum_keepdims.shape}")
|
|
assert sum_keepdims.shape == (2, 1, 2), f"Expected (2, 1, 2), got {sum_keepdims.shape}"
|
|
|
|
# Sum multiple axes with keepdims
|
|
sum_multi_keepdims = t3d.sum(axis=(0, 2), keepdims=True)
|
|
print(f"Sum along axes (0, 2) with keepdims: shape={sum_multi_keepdims.shape}")
|
|
assert sum_multi_keepdims.shape == (1, 2, 1), f"Expected (1, 2, 1), got {sum_multi_keepdims.shape}"
|
|
|
|
# Sum all axes with keepdims
|
|
sum_all_keepdims = t3d.sum(keepdims=True)
|
|
print(f"Sum all with keepdims: shape={sum_all_keepdims.shape}")
|
|
assert sum_all_keepdims.shape == (1, 1, 1), f"Expected (1, 1, 1), got {sum_all_keepdims.shape}"
|
|
assert sum_all_keepdims[0, 0, 0] == 36 # Sum of all elements
|
|
print()
|
|
|
|
# Test 11: Factory functions
|
|
print("Test 11: Factory functions")
|
|
z = cp.zeros((2, 3))
|
|
print(f"zeros((2, 3)): shape={z.shape}")
|
|
o = cp.ones((3, 2))
|
|
print(f"ones((3, 2)): shape={o.shape}")
|
|
e = cp.eye(3)
|
|
print(f"eye(3): shape={e.shape}")
|
|
ar = cp.arange(0, 10, 2)
|
|
print(f"arange(0, 10, 2): shape={ar.shape}")
|
|
assert z.shape == (2, 3)
|
|
assert z[0, 0] == 0
|
|
assert o.shape == (3, 2)
|
|
assert o[1, 1] == 1
|
|
print()
|
|
|
|
# Test 12: Size and properties
|
|
print("Test 12: Size and properties")
|
|
t = cp.tensor([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
|
|
print(f"Shape: {t.shape}")
|
|
print(f"ndim: {t.ndim}")
|
|
print(f"size: {t.size()}")
|
|
assert t.shape == (2, 2, 2)
|
|
assert t.ndim == 3
|
|
print()
|
|
|
|
def test_tensor_slicing():
|
|
print("Test Numpy-style slicing")
|
|
t = cp.tensor([[10, 20, 30], [40, 50, 60], [70, 80, 90]])
|
|
print(f"Original tensor: shape={t.shape}")
|
|
|
|
slice1 = t[1]
|
|
print(f"t[1]: {slice1}, shape={slice1.shape}")
|
|
assert slice1.shape == (3,)
|
|
assert slice1[0] == 40
|
|
|
|
slice2 = t[:, 2]
|
|
print(f"t[:, 2]: {slice2}, shape={slice2.shape}")
|
|
assert slice2.shape == (3,)
|
|
assert slice2[1] == 60
|
|
|
|
slice3 = t[0:2, 1:3]
|
|
print(f"t[0:2, 1:3]: {slice3}, shape={slice3.shape}")
|
|
assert slice3.shape == (2, 2)
|
|
assert slice3[0, 0] == 20
|
|
|
|
slice4 = t[-1, :]
|
|
print(f"t[-1, :]: {slice4}, shape={slice4.shape}")
|
|
assert slice4.shape == (3,)
|
|
assert slice4[2] == 90
|
|
print()
|
|
|
|
if __name__ == "__main__":
|
|
test_tensor_basic()
|
|
print("All tests completed!")
|
|
|