Updates for pytorch 1.6.0

CHANGELOG.md

@@ -41,6 +41,7 @@
 - [#635](https://github.com/helmholtz-analytics/heat/pull/635) `DNDarray.__getitem__` balances and resplits the given key to None if the key is a DNDarray
 - [#638](https://github.com/helmholtz-analytics/heat/pull/638) Fix: arange returns float32 with single input of type float & update skipped device tests
 - [#639](https://github.com/helmholtz-analytics/heat/pull/639) Bufix: balanced array in demo_knn, changed behaviour of knn
+- [#653](https://github.com/helmholtz-analytics/heat/pull/653) Update unittests argmax & argmin + force index order in mpi_argmax & mpi_argmin.
 
 # v0.4.0
 

heat/core/statistics.py

@@ -1184,8 +1184,14 @@ def mpi_argmax(a, b, _):
     rhs = torch.from_numpy(np.frombuffer(b, dtype=np.float64))
 
     # extract the values and minimal indices from the buffers (first half are values, second are indices)
-    values = torch.stack((lhs.chunk(2)[0], rhs.chunk(2)[0]), dim=1)
-    indices = torch.stack((lhs.chunk(2)[1], rhs.chunk(2)[1]), dim=1)
+    idx_l, idx_r = lhs.chunk(2)[1], rhs.chunk(2)[1]
+
+    if idx_l[0] < idx_r[0]:
+        values = torch.stack((lhs.chunk(2)[0], rhs.chunk(2)[0]), dim=1)
+        indices = torch.stack((idx_l, idx_r), dim=1)
+    else:
+        values = torch.stack((rhs.chunk(2)[0], lhs.chunk(2)[0]), dim=1)
+        indices = torch.stack((idx_r, idx_l), dim=1)
 
     # determine the minimum value and select the indices accordingly
     max, max_indices = torch.max(values, dim=1)
@@ -1201,8 +1207,14 @@ def mpi_argmin(a, b, _):
     lhs = torch.from_numpy(np.frombuffer(a, dtype=np.float64))
     rhs = torch.from_numpy(np.frombuffer(b, dtype=np.float64))
     # extract the values and minimal indices from the buffers (first half are values, second are indices)
-    values = torch.stack((lhs.chunk(2)[0], rhs.chunk(2)[0]), dim=1)
-    indices = torch.stack((lhs.chunk(2)[1], rhs.chunk(2)[1]), dim=1)
+    idx_l, idx_r = lhs.chunk(2)[1], rhs.chunk(2)[1]
+
+    if idx_l[0] < idx_r[0]:
+        values = torch.stack((lhs.chunk(2)[0], rhs.chunk(2)[0]), dim=1)
+        indices = torch.stack((idx_l, idx_r), dim=1)
+    else:
+        values = torch.stack((rhs.chunk(2)[0], lhs.chunk(2)[0]), dim=1)
+        indices = torch.stack((idx_r, idx_l), dim=1)
 
     # determine the minimum value and select the indices accordingly
     min, min_indices = torch.min(values, dim=1)

heat/core/tests/test_statistics.py

@@ -65,6 +65,7 @@ def test_argmax(self):
         data = ht.tril(ht.ones((size, size), split=0), k=-1)
 
         result = ht.argmax(data, axis=0)
+        expected = torch.tensor(np.argmax(data.numpy(), axis=0))
         self.assertIsInstance(result, ht.DNDarray)
         self.assertEqual(result.dtype, ht.int64)
         self.assertEqual(result._DNDarray__array.dtype, torch.int64)
@@ -73,24 +74,24 @@ def test_argmax(self):
         self.assertEqual(result.split, None)
         # skip test on gpu; argmax works different
         if not (torch.cuda.is_available() and result.device == ht.gpu):
-            self.assertTrue((result._DNDarray__array != 0).all())
+            self.assertTrue((result._DNDarray__array == expected).all())
 
         # 2D split tensor, across the axis, output tensor
         size = ht.MPI_WORLD.size * 2
         data = ht.tril(ht.ones((size, size), split=0), k=-1)
 
         output = ht.empty((size,))
         result = ht.argmax(data, axis=0, out=output)
-
+        expected = torch.tensor(np.argmax(data.numpy(), axis=0))
         self.assertIsInstance(result, ht.DNDarray)
         self.assertEqual(output.dtype, ht.int64)
         self.assertEqual(output._DNDarray__array.dtype, torch.int64)
         self.assertEqual(output.shape, (size,))
         self.assertEqual(output.lshape, (size,))
         self.assertEqual(output.split, None)
         # skip test on gpu; argmax works different
-        if not (torch.cuda.is_available() and output.device == ht.gpu):
-            self.assertTrue((output._DNDarray__array != 0).all())
+        if not (torch.cuda.is_available() and result.device == ht.gpu):
+            self.assertTrue((output._DNDarray__array == expected).all())
 
         # check exceptions
         with self.assertRaises(TypeError):
@@ -155,6 +156,7 @@ def test_argmin(self):
         data = ht.triu(ht.ones((size, size), split=0), k=1)
 
         result = ht.argmin(data, axis=0)
+        expected = torch.tensor(np.argmin(data.numpy(), axis=0))
         self.assertIsInstance(result, ht.DNDarray)
         self.assertEqual(result.dtype, ht.int64)
         self.assertEqual(result._DNDarray__array.dtype, torch.int64)
@@ -163,24 +165,24 @@ def test_argmin(self):
         self.assertEqual(result.split, None)
         # skip test on gpu; argmin works different
         if not (torch.cuda.is_available() and result.device == ht.gpu):
-            self.assertTrue((result._DNDarray__array != 0).all())
+            self.assertTrue((result._DNDarray__array == expected).all())
 
         # 2D split tensor, across the axis, output tensor
         size = ht.MPI_WORLD.size * 2
         data = ht.triu(ht.ones((size, size), split=0), k=1)
 
         output = ht.empty((size,))
         result = ht.argmin(data, axis=0, out=output)
-
+        expected = torch.tensor(np.argmin(data.numpy(), axis=0))
         self.assertIsInstance(result, ht.DNDarray)
         self.assertEqual(output.dtype, ht.int64)
         self.assertEqual(output._DNDarray__array.dtype, torch.int64)
         self.assertEqual(output.shape, (size,))
         self.assertEqual(output.lshape, (size,))
         self.assertEqual(output.split, None)
         # skip test on gpu; argmin works different
-        if not (torch.cuda.is_available() and output.device == ht.gpu):
-            self.assertTrue((output._DNDarray__array != 0).all())
+        if not (torch.cuda.is_available() and result.device == ht.gpu):
+            self.assertTrue((output._DNDarray__array == expected).all())
 
         # check exceptions
         with self.assertRaises(TypeError):

setup.py

@@ -30,7 +30,7 @@
         "Intended Audience :: Science/Research",
         "Topic :: Scientific/Engineering",
     ],
-    install_requires=["mpi4py>=3.0.0", "numpy>=1.13.0", "torch>=1.5.0", "scipy>=0.14.0"],
+    install_requires=["mpi4py>=3.0.0", "numpy>=1.13.0", "torch>=1.6.0", "scipy>=0.14.0"],
     extras_require={
         "hdf5": ["h5py>=2.8.0"],
         "netcdf": ["netCDF4>=1.4.0,<=1.5.2"],