1# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
2
3import numpy as np
4import os
5import sys
6
7_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
8sys.path.append(_SCRIPT_PATH)
9from tools import mlir_pytaco_api as pt
10
11compressed = pt.compressed
12dense = pt.dense
13
14# Ensure that we can run an unmodified PyTACO program with a simple tensor
15# algebra expression using tensor index notation, and produce the expected
16# result.
17i, j = pt.get_index_vars(2)
18A = pt.tensor([2, 3])
19B = pt.tensor([2, 3])
20C = pt.tensor([2, 3])
21D = pt.tensor([2, 3], compressed)
22A.insert([0, 1], 10)
23A.insert([1, 2], 40)
24B.insert([0, 0], 20)
25B.insert([1, 2], 30)
26C.insert([0, 1], 5)
27C.insert([1, 2], 7)
28D[i, j] = A[i, j] + B[i, j] - C[i, j]
29
30indices, values = D.get_coordinates_and_values()
31passed = np.array_equal(indices, [[0, 0], [0, 1], [1, 2]])
32passed += np.allclose(values, [20.0, 5.0, 63.0])
33
34# PyTACO doesn't allow the use of index values, but MLIR-PyTACO removes this
35# restriction.
36E = pt.tensor([3])
37E[i] = i
38indices, values = E.get_coordinates_and_values()
39passed += np.array_equal(indices, [[0], [1], [2]])
40passed += np.allclose(values, [0.0, 1.0, 2.0])
41
42F = pt.tensor([3])
43G = pt.tensor([3])
44F.insert([0], 10)
45F.insert([2], 40)
46G[i] = F[i] + i
47indices, values = G.get_coordinates_and_values()
48passed += np.array_equal(indices, [[0], [1], [2]])
49passed += np.allclose(values, [10.0, 1.0, 42.0])
50
51H = pt.tensor([3])
52I = pt.tensor([3])
53H.insert([0], 10)
54H.insert([2], 40)
55I[i] = H[i] * i
56indices, values = I.get_coordinates_and_values()
57passed += np.array_equal(indices, [[0], [2]])
58passed += np.allclose(values, [0.0, 80.0])
59
60# CHECK: Number of passed: 8
61print("Number of passed:", passed)
62