mirror of https://github.com/Nonannet/copapy.git
test for tensor type added/updated
This commit is contained in:
Nicolas Kruse
parent
0d0e675207
commit
05fa44fac7
|
|
@ -4,26 +4,24 @@ import pytest
|
|||
|
||||
def test_matrix_init():
|
||||
"""Test basic matrix initialization"""
|
||||
m1 = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
assert m1.rows == 2
|
||||
assert m1.cols == 3
|
||||
m1 = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
assert m1.shape == (2, 3)
|
||||
assert m1[0] == (1, 2, 3)
|
||||
assert m1[1] == (4, 5, 6)
|
||||
|
||||
|
||||
def test_matrix_with_variables():
|
||||
"""Test matrix initialization with variables"""
|
||||
m1 = cp.matrix([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
assert m1.rows == 2
|
||||
assert m1.cols == 2
|
||||
assert isinstance(m1[0][0], cp.value)
|
||||
assert isinstance(m1[1][1], cp.value)
|
||||
m1 = cp.tensor([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
assert m1.shape == (2, 2)
|
||||
assert isinstance(m1[0][0], cp.tensor)
|
||||
assert isinstance(m1[1][1], cp.tensor)
|
||||
|
||||
|
||||
def test_matrix_addition():
|
||||
"""Test matrix addition"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[5, 6], [7, 8]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
m3 = m1 + m2
|
||||
|
||||
assert m3[0] == (6, 8)
|
||||
|
|
@ -32,7 +30,7 @@ def test_matrix_addition():
|
|||
|
||||
def test_matrix_scalar_addition():
|
||||
"""Test matrix addition with scalar"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = m1 + 5
|
||||
|
||||
assert m2[0] == (6, 7)
|
||||
|
|
@ -41,8 +39,8 @@ def test_matrix_scalar_addition():
|
|||
|
||||
def test_matrix_subtraction():
|
||||
"""Test matrix subtraction"""
|
||||
m1 = cp.matrix([[5, 6], [7, 8]])
|
||||
m2 = cp.matrix([[1, 2], [3, 4]])
|
||||
m1 = cp.tensor([[5, 6], [7, 8]])
|
||||
m2 = cp.tensor([[1, 2], [3, 4]])
|
||||
m3 = m1 - m2
|
||||
|
||||
assert m3[0] == (4, 4)
|
||||
|
|
@ -51,7 +49,7 @@ def test_matrix_subtraction():
|
|||
|
||||
def test_matrix_scalar_subtraction():
|
||||
"""Test matrix subtraction with scalar"""
|
||||
m1 = cp.matrix([[5, 6], [7, 8]])
|
||||
m1 = cp.tensor([[5, 6], [7, 8]])
|
||||
m2 = m1 - 2
|
||||
|
||||
assert m2[0] == (3, 4)
|
||||
|
|
@ -60,17 +58,19 @@ def test_matrix_scalar_subtraction():
|
|||
|
||||
def test_matrix_negation():
|
||||
"""Test matrix negation"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = -m1
|
||||
|
||||
assert m1[0] == (1, 2)
|
||||
assert m1[1] == (3, 4)
|
||||
assert m2[0] == (-1, -2)
|
||||
assert m2[1] == (-3, -4)
|
||||
|
||||
|
||||
def test_matrix_element_wise_multiplication():
|
||||
"""Test element-wise matrix multiplication"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[5, 6], [7, 8]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
m3 = m1 * m2
|
||||
|
||||
assert m3[0] == (5, 12)
|
||||
|
|
@ -79,7 +79,7 @@ def test_matrix_element_wise_multiplication():
|
|||
|
||||
def test_matrix_scalar_multiplication():
|
||||
"""Test matrix multiplication with scalar"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = m1 * 3
|
||||
|
||||
assert m2[0] == (3, 6)
|
||||
|
|
@ -88,8 +88,8 @@ def test_matrix_scalar_multiplication():
|
|||
|
||||
def test_matrix_element_wise_division():
|
||||
"""Test element-wise matrix division"""
|
||||
m1 = cp.matrix([[6.0, 8.0], [12.0, 16.0]])
|
||||
m2 = cp.matrix([[2.0, 2.0], [3.0, 4.0]])
|
||||
m1 = cp.tensor([[6.0, 8.0], [12.0, 16.0]])
|
||||
m2 = cp.tensor([[2.0, 2.0], [3.0, 4.0]])
|
||||
m3 = m1 / m2
|
||||
|
||||
assert m3[0][0] == pytest.approx(3.0) # pyright: ignore[reportUnknownMemberType]
|
||||
|
|
@ -100,7 +100,7 @@ def test_matrix_element_wise_division():
|
|||
|
||||
def test_matrix_scalar_division():
|
||||
"""Test matrix division by scalar"""
|
||||
m1 = cp.matrix([[6.0, 8.0], [12.0, 16.0]])
|
||||
m1 = cp.tensor([[6.0, 8.0], [12.0, 16.0]])
|
||||
m2 = m1 / 2.0
|
||||
|
||||
assert list(m2[0]) == pytest.approx((3.0, 4.0)) # pyright: ignore[reportUnknownMemberType]
|
||||
|
|
@ -109,25 +109,24 @@ def test_matrix_scalar_division():
|
|||
|
||||
def test_matrix_vector_multiplication():
|
||||
"""Test matrix-vector multiplication using @ operator"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
v = cp.vector([7, 8, 9])
|
||||
result = m @ v
|
||||
|
||||
assert isinstance(result, cp.vector)
|
||||
assert isinstance(result, cp.tensor)
|
||||
assert len(result.values) == 2
|
||||
assert result.values[0] == 1*7 + 2*8 + 3*9
|
||||
assert result.values[1] == 4*7 + 5*8 + 6*9
|
||||
assert result[0] == 1*7 + 2*8 + 3*9
|
||||
assert result[1] == 4*7 + 5*8 + 6*9
|
||||
|
||||
|
||||
def test_matrix_matrix_multiplication():
|
||||
"""Test matrix-matrix multiplication using @ operator"""
|
||||
m1 = cp.matrix([[1, 2], [3, 4]])
|
||||
m2 = cp.matrix([[5, 6], [7, 8]])
|
||||
m1 = cp.tensor([[1, 2], [3, 4]])
|
||||
m2 = cp.tensor([[5, 6], [7, 8]])
|
||||
result = m1 @ m2
|
||||
|
||||
assert isinstance(result, cp.matrix)
|
||||
assert result.rows == 2
|
||||
assert result.cols == 2
|
||||
assert isinstance(result, cp.tensor)
|
||||
assert result.shape == (2, 2)
|
||||
assert result[0][0] == 1*5 + 2*7
|
||||
assert result[0][1] == 1*6 + 2*8
|
||||
assert result[1][0] == 3*5 + 4*7
|
||||
|
|
@ -136,11 +135,10 @@ def test_matrix_matrix_multiplication():
|
|||
|
||||
def test_matrix_transpose():
|
||||
"""Test matrix transpose"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
mt = m.transpose()
|
||||
|
||||
assert mt.rows == 3
|
||||
assert mt.cols == 2
|
||||
assert mt.shape == (3, 2)
|
||||
assert mt[0] == (1, 4)
|
||||
assert mt[1] == (2, 5)
|
||||
assert mt[2] == (3, 6)
|
||||
|
|
@ -148,35 +146,34 @@ def test_matrix_transpose():
|
|||
|
||||
def test_matrix_transpose_property():
|
||||
"""Test matrix transpose using .T property"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
mt = m.T
|
||||
|
||||
assert mt.rows == 3
|
||||
assert mt.cols == 2
|
||||
assert mt.shape == (3, 2)
|
||||
assert mt[0] == (1, 4)
|
||||
|
||||
|
||||
def test_matrix_row_access():
|
||||
"""Test getting a row as a vector"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
row0 = m.row(0)
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
row0 = m[0]
|
||||
|
||||
assert isinstance(row0, cp.vector)
|
||||
assert isinstance(row0, cp.tensor)
|
||||
assert row0.values == (1, 2, 3)
|
||||
|
||||
|
||||
def test_matrix_col_access():
|
||||
"""Test getting a column as a vector"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
col1 = m.col(1)
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
col1 = m[:, 1]
|
||||
|
||||
assert isinstance(col1, cp.vector)
|
||||
assert col1.values == (2, 5)
|
||||
assert isinstance(col1, cp.tensor)
|
||||
assert col1 == (2, 5)
|
||||
|
||||
|
||||
def test_matrix_trace():
|
||||
"""Test matrix trace (sum of diagonal elements)"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
|
||||
trace = m.trace()
|
||||
|
||||
assert trace == 1 + 5 + 9
|
||||
|
|
@ -184,7 +181,7 @@ def test_matrix_trace():
|
|||
|
||||
def test_matrix_sum():
|
||||
"""Test sum of all matrix elements"""
|
||||
m = cp.matrix([[1, 2, 3], [4, 5, 6]])
|
||||
m = cp.tensor([[1, 2, 3], [4, 5, 6]])
|
||||
total = m.sum()
|
||||
|
||||
assert total == 1 + 2 + 3 + 4 + 5 + 6
|
||||
|
|
@ -192,7 +189,7 @@ def test_matrix_sum():
|
|||
|
||||
def test_matrix_map():
|
||||
"""Test mapping a function over matrix elements"""
|
||||
m = cp.matrix([[1, 2], [3, 4]])
|
||||
m = cp.tensor([[1, 2], [3, 4]])
|
||||
m_doubled = m.map(lambda x: x * 2)
|
||||
|
||||
assert m_doubled[0] == (2, 4)
|
||||
|
|
@ -201,20 +198,19 @@ def test_matrix_map():
|
|||
|
||||
def test_matrix_homogenize():
|
||||
"""Test homogenizing matrix (converting to all variables)"""
|
||||
m = cp.matrix([[1, cp.value(2)], [3, 4]])
|
||||
m = cp.tensor([[1, cp.value(2)], [3, 4]])
|
||||
m_homo = m.homogenize()
|
||||
|
||||
for row in m_homo:
|
||||
for elem in row:
|
||||
assert isinstance(elem, cp.value)
|
||||
assert isinstance(elem, cp.tensor) and elem.ndim == 0
|
||||
|
||||
|
||||
def test_identity_matrix():
|
||||
"""Test identity matrix creation"""
|
||||
m = cp.identity(3)
|
||||
|
||||
assert m.rows == 3
|
||||
assert m.cols == 3
|
||||
assert m.shape == (3, 3)
|
||||
assert m[0] == (1, 0, 0)
|
||||
assert m[1] == (0, 1, 0)
|
||||
assert m[2] == (0, 0, 1)
|
||||
|
|
@ -222,20 +218,18 @@ def test_identity_matrix():
|
|||
|
||||
def test_zeros_matrix():
|
||||
"""Test zeros matrix creation"""
|
||||
m = cp.zeros(2, 3)
|
||||
m = cp.zeros([2, 3])
|
||||
|
||||
assert m.rows == 2
|
||||
assert m.cols == 3
|
||||
assert m.shape == (2, 3)
|
||||
assert m[0] == (0, 0, 0)
|
||||
assert m[1] == (0, 0, 0)
|
||||
|
||||
|
||||
def test_ones_matrix():
|
||||
"""Test ones matrix creation"""
|
||||
m = cp.ones(2, 3)
|
||||
m = cp.ones([2, 3])
|
||||
|
||||
assert m.rows == 2
|
||||
assert m.cols == 3
|
||||
assert m.shape == (2, 3)
|
||||
assert m[0] == (1, 1, 1)
|
||||
assert m[1] == (1, 1, 1)
|
||||
|
||||
|
|
@ -245,8 +239,7 @@ def test_diagonal_matrix():
|
|||
v = cp.vector([1, 2, 3])
|
||||
m = cp.diagonal(v)
|
||||
|
||||
assert m.rows == 3
|
||||
assert m.cols == 3
|
||||
assert m.shape == (3, 3)
|
||||
assert m[0] == (1, 0, 0)
|
||||
assert m[1] == (0, 2, 0)
|
||||
assert m[2] == (0, 0, 3)
|
||||
|
|
@ -254,7 +247,7 @@ def test_diagonal_matrix():
|
|||
|
||||
def test_matrix_with_variables_compiled():
|
||||
"""Test matrix operations with variables in compilation"""
|
||||
m = cp.matrix([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
m = cp.tensor([[cp.value(1), 2], [3, cp.value(4)]])
|
||||
v = cp.vector([cp.value(5), 6])
|
||||
result = m @ v
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,257 @@
|
|||
#!/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!")
|
||||
|
||||
Loading…
Reference in New Issue