def test_binary_add_reduce_3d_2():
i, j, k, l = mlir_pytaco.get_index_vars(4)
I = 2
J = 3
K = 4
L = 5
A = mlir_pytaco.Tensor([I, J, K], [_COMPRESSED, _COMPRESSED, _DENSE])
B = mlir_pytaco.Tensor([I, L, K], [_DENSE, _COMPRESSED, _COMPRESSED])
C = mlir_pytaco.Tensor([J, K], [_COMPRESSED, _COMPRESSED])
D = mlir_pytaco.Tensor([L])
E = mlir_pytaco.Tensor([I], _DENSE)
_init_3d(A, I, J, K)
_init_3d(B, I, L, K)
_init_2d(C, J, K)
_init_1d_with_value(D, L, 2)
E[i] = A[i, j, k] + C[j, k] + B[i, l, k] * D[l]
indices, values = E.get_coordinates_and_values()
passed = np.array_equal(indices, [[0], [1]])
passed += np.array_equal(values, [264., 316.])
# The expression is implemented as:
# temp0[i, k] = A[i, j, k] + C[j, k]
# temp1[i, k] = B[i, l, k] * D[l]
# E[i] = temp0[i, k] + temp1[i, k]
# Check the two temporary tensors introduced by the implementation.
stats = E._stats
passed += (stats.get_total() == 2)
passed += (stats.get_formats(0) == (_COMPRESSED, _DENSE))
passed += (stats.get_dimensions(0) == (I, K))
passed += (stats.get_formats(1) == (_DENSE, _COMPRESSED))
passed += (stats.get_dimensions(1) == (I, K))
# CHECK: Number of passed: 7
print("Number of passed:", passed)
def test_tensor_copy():
i, j = mlir_pytaco.get_index_vars(2)
I = 2
J = 3
A = mlir_pytaco.Tensor([I, J])
A.insert([0, 1], 5.0)
A.insert([1, 2], 6.0)
B = mlir_pytaco.Tensor([I, J])
B[i, j] = A[i, j]
passed = (B._assignment is not None)
passed += (B._engine is None)
try:
B.compute()
except ValueError as e:
passed += (str(e).startswith("Need to invoke compile"))
B.compile()
passed += (B._engine is not None)
B.compute()
passed += (B._assignment is None)
passed += (B._engine is None)
indices, values = B.get_coordinates_and_values()
passed += np.array_equal(indices, [[0, 1], [1, 2]])
passed += np.allclose(values, [5.0, 6.0])
# No temporary tensor is used.
passed += (B._stats.get_total() == 0)
# CHECK: Number of passed: 9
print("Number of passed:", passed)
def test_index_var():
i = mlir_pytaco.IndexVar()
j = mlir_pytaco.IndexVar()
passed = (i.name != j.name)
vars = mlir_pytaco.get_index_vars(10)
passed += (len(vars) == 10)
passed += (all([isinstance(e, mlir_pytaco.IndexVar) for e in vars]))
# CHECK: Number of passed: 3
print("Number of passed:", passed)
def test_tensor_copy():
i, j = mlir_pytaco.get_index_vars(2)
I = 2
J = 3
A = mlir_pytaco.Tensor([I, J])
A.insert([0, 1], 5.0)
A.insert([1, 2], 6.0)
B = mlir_pytaco.Tensor([I, J])
B[i, j] = A[i, j]
indices, values = B.get_coordinates_and_values()
passed = np.array_equal(indices, [[0, 1], [1, 2]])
passed += np.allclose(values, [5.0, 6.0])
# CHECK: Number of passed: 2
print("Number of passed:", passed)
def test_tensor_trivial_reduction():
i, j = mlir_pytaco.get_index_vars(2)
I = 2
J = 3
A = mlir_pytaco.Tensor([I, J])
A.insert([0, 1], 5.0)
A.insert([0, 2], 3.0)
A.insert([1, 2], 6.0)
B = mlir_pytaco.Tensor([I])
B[i] = A[i, j]
indices, values = B.get_coordinates_and_values()
passed = np.array_equal(indices, [[0], [1]])
passed += np.allclose(values, [8.0, 6.0])
# No temporary tensor is used.
passed += (B._stats.get_total() == 0)
# CHECK: Number of passed: 3
print("Number of passed:", passed)
def test_write_packed_tns():
a = mlir_pytaco.Tensor([2, 3])
a.insert([0, 1], 10)
a.insert([1, 2], 40)
a.insert([0, 0], 20)
b = mlir_pytaco.Tensor([2, 3])
i, j = mlir_pytaco.get_index_vars(2)
b[i, j] = a[i, j] + a[i, j]
with tempfile.TemporaryDirectory() as test_dir:
file_name = os.path.join(test_dir, "data.tns")
mlir_pytaco_io.write(file_name, b)
with open(file_name, "r") as file:
lines = file.readlines()
passed = 0
# Skip the comment line in the output.
if lines[1:] == ["2 3\n", "2 3\n", "1 1 40\n", "1 2 20\n", "2 3 80\n"]:
passed = 1
# CHECK: 1
print(passed)