Remove use of deprecated NumPy list indexing

aesara-devs · Sep 29, 2022 · 8086efe · 8086efe
1 parent 1e36f60
commit 8086efe
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/aesara/tensor/sort.py b/aesara/tensor/sort.py
@@ -291,18 +291,18 @@ def _topk_py_impl(op, x, k, axis, idx_dtype):
     idx[axis] = slice(-k, None) if k > 0 else slice(-k)
 
     if not op.return_indices:
-        zv = np.partition(x, -k, axis=axis)[idx]
+        zv = np.partition(x, -k, axis=axis)[tuple(idx)]
         return zv
     elif op.return_values:
-        zi = np.argpartition(x, -k, axis=axis)[idx]
+        zi = np.argpartition(x, -k, axis=axis)[tuple(idx)]
         idx2 = tuple(
             np.arange(s).reshape((s,) + (1,) * (ndim - i - 1)) if i != axis else zi
             for i, s in enumerate(x.shape)
         )
         zv = x[idx2]
         return zv, zi.astype(idx_dtype)
     else:
-        zi = np.argpartition(x, -k, axis=axis)[idx]
+        zi = np.argpartition(x, -k, axis=axis)[tuple(idx)]
         return zi.astype(idx_dtype)