[mlir][sparse][taco] Use np.array_equal to compare integer values.

Fix MLIR-PyTACO and some tests to use np.array_equal to compare integer
values.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D120526

diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/test_simple_tensor_algebra.py b/mlir/test/Integration/Dialect/SparseTensor/taco/test_simple_tensor_algebra.py
index 68ae7ef..d02bdce 100644 (file)
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/test_simple_tensor_algebra.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/test_simple_tensor_algebra.py
@@ -28,7 +28,7 @@ C.insert([1, 2], 7)
 D[i, j] = A[i, j] + B[i, j] - C[i, j]
 
 indices, values = D.get_coordinates_and_values()
-passed = np.allclose(indices, [[0, 0], [0, 1], [1, 2]])
+passed = np.array_equal(indices, [[0, 0], [0, 1], [1, 2]])
 passed += np.allclose(values, [20.0, 5.0, 63.0])
 
 # CHECK: Number of passed: 2

diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
index 7b07821..2fdb9f4 100644 (file)
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
@@ -692,7 +692,7 @@ class Tensor:
     rank, nse, shape, values, indices = utils.sparse_tensor_to_coo_tensor(
         self._packed_sparse_value, self._dtype.value)
     assert rank == self.order
-    assert np.allclose(self.shape, shape)
+    assert np.array_equal(self.shape, shape)
     assert nse == len(values)
     self._coords = indices
     self._values = values

diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_io.py b/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_io.py
index da5c11a..a517e9f 100644 (file)
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_io.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_io.py
@@ -49,7 +49,7 @@ def test_read_mtx_matrix_general():
   a.unpack()
   passed += (a.is_unpacked())
   coords, values = a.get_coordinates_and_values()
-  passed += np.allclose(coords, [[0, 1], [2, 0], [2, 1]])
+  passed += np.array_equal(coords, [[0, 1], [2, 0], [2, 1]])
   passed += np.allclose(values, [2.0, 3.0, 4.0])
   # CHECK: 4
   print(passed)
@@ -71,8 +71,8 @@ def test_read_mtx_matrix_symmetry():
   coords, values = a.get_coordinates_and_values()
   print(coords)
   print(values)
-  passed += np.allclose(coords,
-                        [[0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1]])
+  passed += np.array_equal(coords,
+                           [[0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1]])
   passed += np.allclose(values, [2.0, 3.0, 2.0, 4.0, 3.0, 4.0])
   # CHECK: 4
   print(passed)
@@ -100,7 +100,7 @@ def test_read_tns():
   a.unpack()
   passed += (a.is_unpacked())
   coords, values = a.get_coordinates_and_values()
-  passed += np.allclose(coords, [[0, 1], [2, 0], [2, 1]])
+  passed += np.array_equal(coords, [[0, 1], [2, 0], [2, 1]])
   passed += np.allclose(values, [2.0, 3.0, 4.0])
   # CHECK: 4
   print(passed)

diff --git a/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_utils.py b/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_utils.py
index b3c0533..54354aa 100644 (file)
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_utils.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/unit_test_tensor_utils.py
@@ -80,15 +80,15 @@ def _implement_read_tns_test(
   passed = 0
 
   # Verify the output shape for the tensor.
-  if np.allclose(o_shape, t.shape):
+  if np.array_equal(o_shape, t.shape):
     passed += 1
 
   # Use the output MLIR sparse tensor pointer to retrieve the COO-flavored
   # values and verify the values.
   o_rank, o_nse, o_shape, o_values, o_indices = (
       pytaco_utils.sparse_tensor_to_coo_tensor(sparse_tensor, np.float64))
-  if o_rank == t.rank and o_nse == t.nse and np.allclose(
-      o_shape, t.shape) and np.allclose(o_values, t.values) and np.allclose(
+  if o_rank == t.rank and o_nse == t.nse and np.array_equal(
+      o_shape, t.shape) and np.allclose(o_values, t.values) and np.array_equal(
           o_indices, t.indices):
     passed += 1