1# RUN: SUPPORTLIB=%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext %PYTHON %s | FileCheck %s
2import numpy as np
3import os
4import sys
5
6_SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
7sys.path.append(_SCRIPT_PATH)
8from tools import mlir_pytaco_api as pt
9
10compressed = pt.compressed
11
12passed = 0
13all_types = [pt.complex64, pt.complex128]
14for t in all_types:
15  i, j = pt.get_index_vars(2)
16  A = pt.tensor([2, 3], dtype=t)
17  B = pt.tensor([2, 3], dtype=t)
18  C = pt.tensor([2, 3], compressed, dtype=t)
19  A.insert([0, 1], 10 + 20j)
20  A.insert([1, 2], 40 + 0.5j)
21  B.insert([0, 0], 20)
22  B.insert([1, 2], 30 + 15j)
23  C[i, j] = A[i, j] + B[i, j]
24
25  indices, values = C.get_coordinates_and_values()
26  passed += isinstance(values[0], t.value)
27  passed += np.array_equal(indices, [[0, 0], [0, 1], [1, 2]])
28  passed += np.allclose(values, [20, 10 + 20j, 70 + 15.5j])
29
30# CHECK: Number of passed: 6
31print("Number of passed:", passed)
32