From b4c5217c27c6ad59f7fcef977c59a1d31f5944ef Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 4 Feb 2023 12:11:14 +0300
Subject: [PATCH 1/8] ENH: add flip and rot90

---
 autogen/numpy_api_dump.py                     | 20 --------
 torch_np/_detail/_flips.py                    | 32 ++++++++++++
 torch_np/_helpers.py                          |  6 +++
 torch_np/_ndarray.py                          |  5 +-
 torch_np/_wrapper.py                          | 31 +++++++++---
 .../numpy_tests/lib/test_function_base.py     | 50 ++++++++++---------
 6 files changed, 92 insertions(+), 52 deletions(-)
 create mode 100644 torch_np/_detail/_flips.py

diff --git a/autogen/numpy_api_dump.py b/autogen/numpy_api_dump.py
index 47c04728..2b7a2473 100644
--- a/autogen/numpy_api_dump.py
+++ b/autogen/numpy_api_dump.py
@@ -146,10 +146,6 @@ def asmatrix(data, dtype=None):
     raise NotImplementedError
 
 
-def average(a, axis=None, weights=None, returned=False, *, keepdims=NoValue):
-    raise NotImplementedError
-
-
 def bartlett(M):
     raise NotImplementedError
 
@@ -242,10 +238,6 @@ def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None):
     raise NotImplementedError
 
 
-def cumsum(a, axis=None, dtype=None, out=None):
-    raise NotImplementedError
-
-
 def datetime_as_string(arr, unit=None, timezone="naive", casting="same_kind"):
     raise NotImplementedError
 
@@ -330,10 +322,6 @@ def fix(x, out=None):
     raise NotImplementedError
 
 
-def flip(m, axis=None):
-    raise NotImplementedError
-
-
 def fliplr(m):
     raise NotImplementedError
 
@@ -770,10 +758,6 @@ def printoptions(*args, **kwargs):
     raise NotImplementedError
 
 
-def product(*args, **kwargs):
-    raise NotImplementedError
-
-
 def put(a, ind, v, mode="raise"):
     raise NotImplementedError
 
@@ -818,10 +802,6 @@ def roots(p):
     raise NotImplementedError
 
 
-def rot90(m, k=1, axes=(0, 1)):
-    raise NotImplementedError
-
-
 def safe_eval(source):
     raise NotImplementedError
 
diff --git a/torch_np/_detail/_flips.py b/torch_np/_detail/_flips.py
new file mode 100644
index 00000000..13cb29fa
--- /dev/null
+++ b/torch_np/_detail/_flips.py
@@ -0,0 +1,32 @@
+"""Implementations of flip-based routines and related animals. 
+"""
+
+import torch
+
+from . import _scalar_types, _util
+
+def flip(m_tensor, axis=None):
+    # XXX: semantic difference: np.flip returns a view, torch.flip copies
+    if axis is None:
+        axis = tuple(range(m_tensor.ndim))
+    else:
+        axis = _util.normalize_axis_tuple(axis, m_tensor.ndim)
+    return torch.flip(m_tensor, axis)
+
+
+def flipud(m_tensor):
+    return torch.flipud(m_tensor)
+
+
+def fliplr(m_tensor):
+    return torch.fliplr(m_tensor)
+
+
+def rot90(m_tensor, k=1, axes=(0, 1)):
+    axes = _util.normalize_axis_tuple(axes, m_tensor.ndim)
+    return torch.rot90(m_tensor, k, axes)
+
+
+def swapaxes(tensor, axis1, axis2):
+    return torch.swapaxes(tensor, axis1, axis2)
+
diff --git a/torch_np/_helpers.py b/torch_np/_helpers.py
index cead8ce6..5785ae67 100644
--- a/torch_np/_helpers.py
+++ b/torch_np/_helpers.py
@@ -74,3 +74,9 @@ def ndarrays_to_tensors(*inputs):
 def to_tensors(*inputs):
     """Convert all array_likes from `inputs` to tensors."""
     return tuple(asarray(value).get() for value in inputs)
+
+
+def _outer(x, y):
+    x_tensor, y_tensor = to_tensors(x, y)
+    result = torch.outer(x_tensor, y_tensor)
+    return asarray(result)
diff --git a/torch_np/_ndarray.py b/torch_np/_ndarray.py
index ad011767..45b79af7 100644
--- a/torch_np/_ndarray.py
+++ b/torch_np/_ndarray.py
@@ -10,7 +10,7 @@
     dtype_to_torch,
     emulate_out_arg,
 )
-from ._detail import _reductions, _util
+from ._detail import _reductions, _util, _flips
 
 newaxis = None
 
@@ -264,6 +264,9 @@ def transpose(self, *axes):
             raise ValueError("axes don't match array")
         return ndarray._from_tensor_and_base(tensor, self)
 
+    def swapaxes(self, axis1, axis2):
+        return _flips.swapaxes(self._tensor, axis1, axis2)
+
     def ravel(self, order="C"):
         if order != "C":
             raise NotImplementedError
diff --git a/torch_np/_wrapper.py b/torch_np/_wrapper.py
index e8a9845c..5e9f5282 100644
--- a/torch_np/_wrapper.py
+++ b/torch_np/_wrapper.py
@@ -8,7 +8,7 @@
 import torch
 
 from . import _dtypes, _helpers
-from ._detail import _reductions, _util
+from ._detail import _reductions, _util, _flips
 from ._ndarray import (
     array,
     asarray,
@@ -440,12 +440,22 @@ def expand_dims(a, axis):
 
 @asarray_replacer()
 def flip(m, axis=None):
-    # XXX: semantic difference: np.flip returns a view, torch.flip copies
-    if axis is None:
-        axis = tuple(range(m.ndim))
-    else:
-        axis = _util.normalize_axis_tuple(axis, m.ndim)
-    return torch.flip(m, axis)
+    return _flips.flip(m, axis)
+
+
+@asarray_replacer()
+def flipud(m):
+    return _flips.flipud(m)
+
+
+@asarray_replacer()
+def fliplr(m):
+    return _flips.fliplr(m)
+
+
+@asarray_replacer()
+def rot90(m, k=1, axes=(0, 1)):
+    return _flips.rot90(m, k, axes)
 
 
 @asarray_replacer()
@@ -469,6 +479,12 @@ def moveaxis(a, source, destination):
     return asarray(torch.moveaxis(a, source, destination))
 
 
+def swapaxis(a, axis1, axis2):
+    arr = asarray(a)
+    return arr.swapaxes(axis1, axis2)
+
+
+
 def unravel_index(indices, shape, order="C"):
     # cf https://github.com/pytorch/pytorch/pull/66687
     # this version is from
@@ -644,6 +660,7 @@ def prod(
         axis=axis, dtype=dtype, out=out, keepdims=keepdims, initial=initial, where=where
     )
 
+product = prod
 
 def cumprod(a, axis=None, dtype=None, out=None):
     arr = asarray(a)
diff --git a/torch_np/tests/numpy_tests/lib/test_function_base.py b/torch_np/tests/numpy_tests/lib/test_function_base.py
index 901ef6bf..544420f7 100644
--- a/torch_np/tests/numpy_tests/lib/test_function_base.py
+++ b/torch_np/tests/numpy_tests/lib/test_function_base.py
@@ -28,8 +28,8 @@
 # FIXME: make from torch_np
 from numpy.lib import (
     add_newdoc_ufunc, angle, bartlett, blackman, corrcoef, cov,
-    delete, diff, digitize, extract, flipud, gradient, hamming, hanning,
-    i0, insert, interp, kaiser, meshgrid, msort, piecewise, place, rot90,
+    delete, diff, digitize, extract, gradient, hamming, hanning,
+    i0, insert, interp, kaiser, meshgrid, msort, piecewise, place,
     select, setxor1d, sinc, trapz, trim_zeros, unwrap, unique, vectorize
     )
 from torch_np._detail._util import normalize_axis_tuple
@@ -37,7 +37,9 @@
 
 def get_mat(n):
     data = np.arange(n)
-    data = np.add.outer(data, data)
+#    data = np.add.outer(data, data)
+    from torch_np._helpers import _outer
+    data = _outer(data, data)
     return data
 
 
@@ -52,14 +54,13 @@ def _make_complex(real, imag):
     return ret
 
 
-@pytest.mark.xfail(reason='TODO: implement')
 class TestRot90:
     def test_basic(self):
-        assert_raises(ValueError, rot90, np.ones(4))
-        assert_raises(ValueError, rot90, np.ones((2,2,2)), axes=(0,1,2))
-        assert_raises(ValueError, rot90, np.ones((2,2)), axes=(0,2))
-        assert_raises(ValueError, rot90, np.ones((2,2)), axes=(1,1))
-        assert_raises(ValueError, rot90, np.ones((2,2,2)), axes=(-2,1))
+        assert_raises(ValueError, np.rot90, np.ones(4))
+        assert_raises((ValueError, RuntimeError), np.rot90, np.ones((2,2,2)), axes=(0,1,2))
+        assert_raises(ValueError, np.rot90, np.ones((2,2)), axes=(0,2))
+        assert_raises(ValueError, np.rot90, np.ones((2,2)), axes=(1,1))
+        assert_raises(ValueError, np.rot90, np.ones((2,2,2)), axes=(-2,1))
 
         a = [[0, 1, 2],
              [3, 4, 5]]
@@ -75,22 +76,22 @@ def test_basic(self):
               [3, 4, 5]]
 
         for k in range(-3, 13, 4):
-            assert_equal(rot90(a, k=k), b1)
+            assert_equal(np.rot90(a, k=k), b1)
         for k in range(-2, 13, 4):
-            assert_equal(rot90(a, k=k), b2)
+            assert_equal(np.rot90(a, k=k), b2)
         for k in range(-1, 13, 4):
-            assert_equal(rot90(a, k=k), b3)
+            assert_equal(np.rot90(a, k=k), b3)
         for k in range(0, 13, 4):
-            assert_equal(rot90(a, k=k), b4)
+            assert_equal(np.rot90(a, k=k), b4)
 
-        assert_equal(rot90(rot90(a, axes=(0,1)), axes=(1,0)), a)
-        assert_equal(rot90(a, k=1, axes=(1,0)), rot90(a, k=-1, axes=(0,1)))
+        assert_equal(np.rot90(np.rot90(a, axes=(0,1)), axes=(1,0)), a)
+        assert_equal(np.rot90(a, k=1, axes=(1,0)), np.rot90(a, k=-1, axes=(0,1)))
 
     def test_axes(self):
         a = np.ones((50, 40, 3))
-        assert_equal(rot90(a).shape, (40, 50, 3))
-        assert_equal(rot90(a, axes=(0,2)), rot90(a, axes=(0,-1)))
-        assert_equal(rot90(a, axes=(1,2)), rot90(a, axes=(-2,-1)))
+        assert_equal(np.rot90(a).shape, (40, 50, 3))
+        assert_equal(np.rot90(a, axes=(0,2)), np.rot90(a, axes=(0,-1)))
+        assert_equal(np.rot90(a, axes=(1,2)), np.rot90(a, axes=(-2,-1)))
 
     def test_rotation_axes(self):
         a = np.arange(8).reshape((2,2,2))
@@ -112,16 +113,15 @@ def test_rotation_axes(self):
                       [[6, 7],
                        [2, 3]]]
 
-        assert_equal(rot90(a, axes=(0, 1)), a_rot90_01)
-        assert_equal(rot90(a, axes=(1, 0)), a_rot90_10)
-        assert_equal(rot90(a, axes=(1, 2)), a_rot90_12)
+        assert_equal(np.rot90(a, axes=(0, 1)), a_rot90_01)
+        assert_equal(np.rot90(a, axes=(1, 0)), a_rot90_10)
+        assert_equal(np.rot90(a, axes=(1, 2)), a_rot90_12)
 
         for k in range(1,5):
-            assert_equal(rot90(a, k=k, axes=(2, 0)),
-                         rot90(a_rot90_20, k=k-1, axes=(2, 0)))
+            assert_equal(np.rot90(a, k=k, axes=(2, 0)),
+                         np.rot90(a_rot90_20, k=k-1, axes=(2, 0)))
 
 
-@pytest.mark.xfail(reason='TODO: implement')
 class TestFlip:
 
     def test_axes(self):
@@ -130,6 +130,7 @@ def test_axes(self):
         assert_raises(np.AxisError, np.flip, np.ones((4, 4)), axis=-3)
         assert_raises(np.AxisError, np.flip, np.ones((4, 4)), axis=(0, 3))
 
+    @pytest.mark.xfail(reason='no [::-1] indexing')
     def test_basic_lr(self):
         a = get_mat(4)
         b = a[:, ::-1]
@@ -140,6 +141,7 @@ def test_basic_lr(self):
              [5, 4, 3]]
         assert_equal(np.flip(a, 1), b)
 
+    @pytest.mark.xfail(reason='no [::-1] indexing')
     def test_basic_ud(self):
         a = get_mat(4)
         b = a[::-1, :]

From 0297cdc8a49b6d7361351a1a3ff553ef4c461eed Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 4 Feb 2023 13:38:15 +0300
Subject: [PATCH 2/8] MAINT: moveaxis

---
 torch_np/_wrapper.py                          |  2 ++
 .../tests/numpy_tests/core/test_numeric.py    | 21 +++++++------------
 2 files changed, 9 insertions(+), 14 deletions(-)

diff --git a/torch_np/_wrapper.py b/torch_np/_wrapper.py
index 5e9f5282..ea419a65 100644
--- a/torch_np/_wrapper.py
+++ b/torch_np/_wrapper.py
@@ -476,6 +476,8 @@ def broadcast_arrays(*args, subok=False):
 
 @asarray_replacer()
 def moveaxis(a, source, destination):
+    source = _util.normalize_axis_tuple(source, a.ndim, 'source')
+    destination = _util.normalize_axis_tuple(destination, a.ndim, 'destination')
     return asarray(torch.moveaxis(a, source, destination))
 
 
diff --git a/torch_np/tests/numpy_tests/core/test_numeric.py b/torch_np/tests/numpy_tests/core/test_numeric.py
index 57cd4d65..bf0e84d7 100644
--- a/torch_np/tests/numpy_tests/core/test_numeric.py
+++ b/torch_np/tests/numpy_tests/core/test_numeric.py
@@ -3164,7 +3164,6 @@ def test_results(self):
             assert_(not res.flags['OWNDATA'])
 
 
-@pytest.mark.xfail(reason="TODO")
 class TestMoveaxis:
     def test_move_to_end(self):
         x = np.random.randn(5, 6, 7)
@@ -3212,27 +3211,21 @@ def test_move_multiples(self):
 
     def test_errors(self):
         x = np.random.randn(1, 2, 3)
-        assert_raises_regex(np.AxisError, 'source.*out of bounds',
+        assert_raises(np.AxisError, #'source.*out of bounds',
                             np.moveaxis, x, 3, 0)
-        assert_raises_regex(np.AxisError, 'source.*out of bounds',
+        assert_raises(np.AxisError, #'source.*out of bounds',
                             np.moveaxis, x, -4, 0)
-        assert_raises_regex(np.AxisError, 'destination.*out of bounds',
+        assert_raises(np.AxisError, #'destination.*out of bounds',
                             np.moveaxis, x, 0, 5)
-        assert_raises_regex(ValueError, 'repeated axis in `source`',
+        assert_raises(ValueError, #'repeated axis in `source`',
                             np.moveaxis, x, [0, 0], [0, 1])
-        assert_raises_regex(ValueError, 'repeated axis in `destination`',
+        assert_raises(ValueError, #'repeated axis in `destination`',
                             np.moveaxis, x, [0, 1], [1, 1])
-        assert_raises_regex(ValueError, 'must have the same number',
+        assert_raises((ValueError, RuntimeError), #'must have the same number',
                             np.moveaxis, x, 0, [0, 1])
-        assert_raises_regex(ValueError, 'must have the same number',
+        assert_raises((ValueError, RuntimeError), #'must have the same number',
                             np.moveaxis, x, [0, 1], [0])
 
-    def test_array_likes(self):
-        x = np.ma.zeros((1, 2, 3))
-        result = np.moveaxis(x, 0, 0)
-        assert_(x.shape, result.shape)
-        assert_(isinstance(result, np.ma.MaskedArray))
-
         x = [1, 2, 3]
         result = np.moveaxis(x, 0, 0)
         assert_(x, list(result))

From 1ac48e3b129dd359a5bc3bb9eaadeadd933e10d9 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 4 Feb 2023 14:19:38 +0300
Subject: [PATCH 3/8] ENH: add rollaxis

---
 torch_np/_detail/_flips.py                    | 19 +++++++++++++++++++
 torch_np/_wrapper.py                          |  3 +++
 .../tests/numpy_tests/core/test_numeric.py    |  2 +-
 3 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/torch_np/_detail/_flips.py b/torch_np/_detail/_flips.py
index 13cb29fa..a46aca5d 100644
--- a/torch_np/_detail/_flips.py
+++ b/torch_np/_detail/_flips.py
@@ -30,3 +30,22 @@ def rot90(m_tensor, k=1, axes=(0, 1)):
 def swapaxes(tensor, axis1, axis2):
     return torch.swapaxes(tensor, axis1, axis2)
 
+
+# https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numeric.py#L1259
+def rollaxis(tensor, axis, start=0):
+    n = tensor.ndim
+    axis = _util.normalize_axis_index(axis, n)
+    if start < 0:
+        start += n
+    msg = "'%s' arg requires %d <= %s < %d, but %d was passed in"
+    if not (0 <= start < n + 1):
+        raise _util.AxisError(msg % ('start', -n, 'start', n + 1, start))
+    if axis < start:
+        # it's been removed
+        start -= 1
+    if axis == start:
+        return tensor[...]
+    axes = list(range(0, n))
+    axes.remove(axis)
+    axes.insert(start, axis)
+    return tensor.view(axes)
diff --git a/torch_np/_wrapper.py b/torch_np/_wrapper.py
index ea419a65..a48b9c85 100644
--- a/torch_np/_wrapper.py
+++ b/torch_np/_wrapper.py
@@ -485,6 +485,9 @@ def swapaxis(a, axis1, axis2):
     arr = asarray(a)
     return arr.swapaxes(axis1, axis2)
 
+@asarray_replacer()
+def rollaxis(a, axis, start=0):
+    return _flips.rollaxis(a, axis, start)
 
 
 def unravel_index(indices, shape, order="C"):
diff --git a/torch_np/tests/numpy_tests/core/test_numeric.py b/torch_np/tests/numpy_tests/core/test_numeric.py
index bf0e84d7..8cd6a0f4 100644
--- a/torch_np/tests/numpy_tests/core/test_numeric.py
+++ b/torch_np/tests/numpy_tests/core/test_numeric.py
@@ -3101,7 +3101,6 @@ def test_roll_empty(self):
         assert_equal(np.roll(x, 1), np.array([]))
 
 
-@pytest.mark.xfail(reason="TODO")
 class TestRollaxis:
 
     # expected shape indexed by (axis, start) for array of
@@ -3126,6 +3125,7 @@ def test_exceptions(self):
         assert_raises(np.AxisError, np.rollaxis, a, 4, 0)
         assert_raises(np.AxisError, np.rollaxis, a, 0, 5)
 
+    @pytest.mark.xfail(reason="needs np.indices")
     def test_results(self):
         a = np.arange(1*2*3*4).reshape(1, 2, 3, 4).copy()
         aind = np.indices(a.shape)

From 513163c335fe796912bdd194dc034b1a10c50155 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 4 Feb 2023 14:20:56 +0300
Subject: [PATCH 4/8] TST: move TestArgwhere back to
 numpy_tests/core/test_numeric.py

---
 .../tests/numpy_tests/core/test_numeric.py    |  5 ++-
 torch_np/tests/test_reductions.py             | 32 -------------------
 2 files changed, 2 insertions(+), 35 deletions(-)

diff --git a/torch_np/tests/numpy_tests/core/test_numeric.py b/torch_np/tests/numpy_tests/core/test_numeric.py
index 8cd6a0f4..2c516ed7 100644
--- a/torch_np/tests/numpy_tests/core/test_numeric.py
+++ b/torch_np/tests/numpy_tests/core/test_numeric.py
@@ -2988,7 +2988,6 @@ def test_mode(self):
             np.convolve(d, k, mode=None)
 
 
-@pytest.mark.xfail(reason="TODO")
 class TestArgwhere:
 
     @pytest.mark.parametrize('nd', [0, 1, 2])
@@ -3002,12 +3001,12 @@ def test_nd(self, nd):
 
         # only one
         x[...] = False
-        x.flat[0] = True
+        x.ravel()[0] = True
         assert_equal(np.argwhere(x).shape, (1, nd))
 
         # all but one
         x[...] = True
-        x.flat[0] = False
+        x.ravel()[0] = False
         assert_equal(np.argwhere(x).shape, (x.size - 1, nd))
 
         # all
diff --git a/torch_np/tests/test_reductions.py b/torch_np/tests/test_reductions.py
index 4772e155..86c69313 100644
--- a/torch_np/tests/test_reductions.py
+++ b/torch_np/tests/test_reductions.py
@@ -142,38 +142,6 @@ def test_basic(self):
         assert_equal(np.flatnonzero(x), [0, 1, 3, 4])
 
 
-class TestArgwhere:
-    @pytest.mark.parametrize("nd", [0, 1, 2])
-    def test_nd(self, nd):
-        # get an nd array with multiple elements in every dimension
-        x = np.empty((2,) * nd, bool)
-
-        # none
-        x[...] = False
-        assert_equal(np.argwhere(x).shape, (0, nd))
-
-        # only one
-        x[...] = False
-        x.ravel()[0] = True
-        assert_equal(np.argwhere(x).shape, (1, nd))
-
-        # all but one
-        x[...] = True
-        x.ravel()[0] = False
-        assert_equal(np.argwhere(x).shape, (x.size - 1, nd))
-
-        # all
-        x[...] = True
-        assert_equal(np.argwhere(x).shape, (x.size, nd))
-
-    def test_2D(self):
-        x = np.arange(6).reshape((2, 3))
-        assert_array_equal(np.argwhere(x > 1), [[0, 2], [1, 0], [1, 1], [1, 2]])
-
-    def test_list(self):
-        assert_equal(np.argwhere([4, 0, 2, 1, 3]), [[0], [2], [3], [4]])
-
-
 class TestAny:
     def test_basic(self):
         y1 = [0, 0, 1, 0]

From 353698f3c800b2ccdaff43934c90e560d7c4b31f Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 4 Feb 2023 14:23:07 +0300
Subject: [PATCH 5/8] TST: reenable TestStdVar

---
 torch_np/tests/numpy_tests/core/test_numeric.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/torch_np/tests/numpy_tests/core/test_numeric.py b/torch_np/tests/numpy_tests/core/test_numeric.py
index 2c516ed7..c7a0f60b 100644
--- a/torch_np/tests/numpy_tests/core/test_numeric.py
+++ b/torch_np/tests/numpy_tests/core/test_numeric.py
@@ -2593,7 +2593,6 @@ def test_non_finite_scalar(self):
         assert_(type(np.isclose(0, np.inf)) is np.bool_)
 
 
-@pytest.mark.xfail(reason="TODO")
 class TestStdVar:
     def setup_method(self):
         self.A = np.array([1, -1, 1, -1])

From 3da48a2c7f0e446fef04b7fbc917596b6ca6f38a Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 4 Feb 2023 15:27:30 +0300
Subject: [PATCH 6/8] TST: move TestNonzero to numpy_tests/core/test_numeric.py

---
 .../tests/numpy_tests/core/test_numeric.py    | 337 ++++--------------
 torch_np/tests/test_reductions.py             | 125 -------
 2 files changed, 65 insertions(+), 397 deletions(-)

diff --git a/torch_np/tests/numpy_tests/core/test_numeric.py b/torch_np/tests/numpy_tests/core/test_numeric.py
index c7a0f60b..303bd250 100644
--- a/torch_np/tests/numpy_tests/core/test_numeric.py
+++ b/torch_np/tests/numpy_tests/core/test_numeric.py
@@ -1251,31 +1251,36 @@ def test_failed_itemsetting(self):
             np.fromiter(iterable, dtype=np.dtype((int, 2)))
 
 
-@pytest.mark.xfail(reason="TODO")
-class TestNonzero:
+class TestNonzeroAndCountNonzero:
+    def test_count_nonzero_list(self):
+        lst = [[0, 1, 2, 3], [1, 0, 0, 6]]
+        assert np.count_nonzero(lst) == 5
+        assert_array_equal(np.count_nonzero(lst, axis=0), np.array([1, 1, 1, 2]))
+        assert_array_equal(np.count_nonzero(lst, axis=1), np.array([3, 2]))
+
     def test_nonzero_trivial(self):
         assert_equal(np.count_nonzero(np.array([])), 0)
-        assert_equal(np.count_nonzero(np.array([], dtype='?')), 0)
+        assert_equal(np.count_nonzero(np.array([], dtype="?")), 0)
         assert_equal(np.nonzero(np.array([])), ([],))
 
         assert_equal(np.count_nonzero(np.array([0])), 0)
-        assert_equal(np.count_nonzero(np.array([0], dtype='?')), 0)
+        assert_equal(np.count_nonzero(np.array([0], dtype="?")), 0)
         assert_equal(np.nonzero(np.array([0])), ([],))
 
         assert_equal(np.count_nonzero(np.array([1])), 1)
-        assert_equal(np.count_nonzero(np.array([1], dtype='?')), 1)
+        assert_equal(np.count_nonzero(np.array([1], dtype="?")), 1)
         assert_equal(np.nonzero(np.array([1])), ([0],))
 
+        assert isinstance(np.count_nonzero([]), np.ndarray)
+
     def test_nonzero_zerod(self):
         assert_equal(np.count_nonzero(np.array(0)), 0)
-        assert_equal(np.count_nonzero(np.array(0, dtype='?')), 0)
-        with assert_warns(DeprecationWarning):
-            assert_equal(np.nonzero(np.array(0)), ([],))
+        assert_equal(np.count_nonzero(np.array(0, dtype="?")), 0)
 
         assert_equal(np.count_nonzero(np.array(1)), 1)
-        assert_equal(np.count_nonzero(np.array(1, dtype='?')), 1)
-        with assert_warns(DeprecationWarning):
-            assert_equal(np.nonzero(np.array(1)), ([0],))
+        assert_equal(np.count_nonzero(np.array(1, dtype="?")), 1)
+
+        assert isinstance(np.count_nonzero(np.array(1)), np.ndarray)
 
     def test_nonzero_onedim(self):
         x = np.array([1, 0, 2, -1, 0, 0, 8])
@@ -1283,46 +1288,23 @@ def test_nonzero_onedim(self):
         assert_equal(np.count_nonzero(x), 4)
         assert_equal(np.nonzero(x), ([0, 2, 3, 6],))
 
-        # x = np.array([(1, 2), (0, 0), (1, 1), (-1, 3), (0, 7)],
-        #              dtype=[('a', 'i4'), ('b', 'i2')])
-        x = np.array([(1, 2, -5, -3), (0, 0, 2, 7), (1, 1, 0, 1), (-1, 3, 1, 0), (0, 7, 0, 4)],
-                     dtype=[('a', 'i4'), ('b', 'i2'), ('c', 'i1'), ('d', 'i8')])
-        assert_equal(np.count_nonzero(x['a']), 3)
-        assert_equal(np.count_nonzero(x['b']), 4)
-        assert_equal(np.count_nonzero(x['c']), 3)
-        assert_equal(np.count_nonzero(x['d']), 4)
-        assert_equal(np.nonzero(x['a']), ([0, 2, 3],))
-        assert_equal(np.nonzero(x['b']), ([0, 2, 3, 4],))
+        assert isinstance(np.count_nonzero(x), np.ndarray)
 
     def test_nonzero_twodim(self):
         x = np.array([[0, 1, 0], [2, 0, 3]])
-        assert_equal(np.count_nonzero(x.astype('i1')), 3)
-        assert_equal(np.count_nonzero(x.astype('i2')), 3)
-        assert_equal(np.count_nonzero(x.astype('i4')), 3)
-        assert_equal(np.count_nonzero(x.astype('i8')), 3)
+        assert_equal(np.count_nonzero(x.astype("i1")), 3)
+        assert_equal(np.count_nonzero(x.astype("i2")), 3)
+        assert_equal(np.count_nonzero(x.astype("i4")), 3)
+        assert_equal(np.count_nonzero(x.astype("i8")), 3)
         assert_equal(np.nonzero(x), ([0, 1, 1], [1, 0, 2]))
 
         x = np.eye(3)
-        assert_equal(np.count_nonzero(x.astype('i1')), 3)
-        assert_equal(np.count_nonzero(x.astype('i2')), 3)
-        assert_equal(np.count_nonzero(x.astype('i4')), 3)
-        assert_equal(np.count_nonzero(x.astype('i8')), 3)
+        assert_equal(np.count_nonzero(x.astype("i1")), 3)
+        assert_equal(np.count_nonzero(x.astype("i2")), 3)
+        assert_equal(np.count_nonzero(x.astype("i4")), 3)
+        assert_equal(np.count_nonzero(x.astype("i8")), 3)
         assert_equal(np.nonzero(x), ([0, 1, 2], [0, 1, 2]))
 
-        x = np.array([[(0, 1), (0, 0), (1, 11)],
-                   [(1, 1), (1, 0), (0, 0)],
-                   [(0, 0), (1, 5), (0, 1)]], dtype=[('a', 'f4'), ('b', 'u1')])
-        assert_equal(np.count_nonzero(x['a']), 4)
-        assert_equal(np.count_nonzero(x['b']), 5)
-        assert_equal(np.nonzero(x['a']), ([0, 1, 1, 2], [2, 0, 1, 1]))
-        assert_equal(np.nonzero(x['b']), ([0, 0, 1, 2, 2], [0, 2, 0, 1, 2]))
-
-        assert_(not x['a'].T.flags.aligned)
-        assert_equal(np.count_nonzero(x['a'].T), 4)
-        assert_equal(np.count_nonzero(x['b'].T), 5)
-        assert_equal(np.nonzero(x['a'].T), ([0, 1, 1, 2], [1, 1, 2, 0]))
-        assert_equal(np.nonzero(x['b'].T), ([0, 0, 1, 2, 2], [0, 1, 2, 0, 2]))
-
     def test_sparse(self):
         # test special sparse condition boolean code path
         for i in range(20):
@@ -1331,256 +1313,67 @@ def test_sparse(self):
             assert_equal(np.nonzero(c)[0], np.arange(i, 200 + i, 20))
 
             c = np.zeros(400, dtype=bool)
-            c[10 + i:20 + i] = True
-            c[20 + i*2] = True
-            assert_equal(np.nonzero(c)[0],
-                         np.concatenate((np.arange(10 + i, 20 + i), [20 + i*2])))
-
-    def test_return_type(self):
-        class C(np.ndarray):
-            pass
-
-        for view in (C, np.ndarray):
-            for nd in range(1, 4):
-                shape = tuple(range(2, 2+nd))
-                x = np.arange(np.prod(shape)).reshape(shape).view(view)
-                for nzx in (np.nonzero(x), x.nonzero()):
-                    for nzx_i in nzx:
-                        assert_(type(nzx_i) is np.ndarray)
-                        assert_(nzx_i.flags.writeable)
+            c[10 + i : 20 + i] = True
+            c[20 + i * 2] = True
+            assert_equal(
+                np.nonzero(c)[0],
+                np.concatenate((np.arange(10 + i, 20 + i), [20 + i * 2])),
+            )
 
     def test_count_nonzero_axis(self):
         # Basic check of functionality
         m = np.array([[0, 1, 7, 0, 0], [3, 0, 0, 2, 19]])
 
         expected = np.array([1, 1, 1, 1, 1])
-        assert_equal(np.count_nonzero(m, axis=0), expected)
+        assert_array_equal(np.count_nonzero(m, axis=0), expected)
 
         expected = np.array([2, 3])
-        assert_equal(np.count_nonzero(m, axis=1), expected)
+        assert_array_equal(np.count_nonzero(m, axis=1), expected)
+
+        assert isinstance(np.count_nonzero(m, axis=1), np.ndarray)
 
         assert_raises(ValueError, np.count_nonzero, m, axis=(1, 1))
-        assert_raises(TypeError, np.count_nonzero, m, axis='foo')
+        assert_raises(TypeError, np.count_nonzero, m, axis="foo")
         assert_raises(np.AxisError, np.count_nonzero, m, axis=3)
-        assert_raises(TypeError, np.count_nonzero,
-                      m, axis=np.array([[1], [2]]))
+        assert_raises(TypeError, np.count_nonzero, m, axis=np.array([[1], [2]]))
 
-    def test_count_nonzero_axis_all_dtypes(self):
+    @pytest.mark.parametrize("typecode", np.typecodes["All"])
+    def test_count_nonzero_axis_all_dtypes(self, typecode):
         # More thorough test that the axis argument is respected
         # for all dtypes and responds correctly when presented with
         # either integer or tuple arguments for axis
-        msg = "Mismatch for dtype: %s"
-
-        def assert_equal_w_dt(a, b, err_msg):
-            assert_equal(a.dtype, b.dtype, err_msg=err_msg)
-            assert_equal(a, b, err_msg=err_msg)
-
-        for dt in np.typecodes['All']:
-            err_msg = msg % (np.dtype(dt).name,)
-
-            if dt != 'V':
-                if dt != 'M':
-                    m = np.zeros((3, 3), dtype=dt)
-                    n = np.ones(1, dtype=dt)
-
-                    m[0, 0] = n[0]
-                    m[1, 0] = n[0]
-
-                else:  # np.zeros doesn't work for np.datetime64
-                    m = np.array(['1970-01-01'] * 9)
-                    m = m.reshape((3, 3))
-
-                    m[0, 0] = '1970-01-12'
-                    m[1, 0] = '1970-01-12'
-                    m = m.astype(dt)
-
-                expected = np.array([2, 0, 0], dtype=np.intp)
-                assert_equal_w_dt(np.count_nonzero(m, axis=0),
-                                  expected, err_msg=err_msg)
-
-                expected = np.array([1, 1, 0], dtype=np.intp)
-                assert_equal_w_dt(np.count_nonzero(m, axis=1),
-                                  expected, err_msg=err_msg)
-
-                expected = np.array(2)
-                assert_equal(np.count_nonzero(m, axis=(0, 1)),
-                             expected, err_msg=err_msg)
-                assert_equal(np.count_nonzero(m, axis=None),
-                             expected, err_msg=err_msg)
-                assert_equal(np.count_nonzero(m),
-                             expected, err_msg=err_msg)
-
-            if dt == 'V':
-                # There are no 'nonzero' objects for np.void, so the testing
-                # setup is slightly different for this dtype
-                m = np.array([np.void(1)] * 6).reshape((2, 3))
-
-                expected = np.array([0, 0, 0], dtype=np.intp)
-                assert_equal_w_dt(np.count_nonzero(m, axis=0),
-                                  expected, err_msg=err_msg)
-
-                expected = np.array([0, 0], dtype=np.intp)
-                assert_equal_w_dt(np.count_nonzero(m, axis=1),
-                                  expected, err_msg=err_msg)
-
-                expected = np.array(0)
-                assert_equal(np.count_nonzero(m, axis=(0, 1)),
-                             expected, err_msg=err_msg)
-                assert_equal(np.count_nonzero(m, axis=None),
-                             expected, err_msg=err_msg)
-                assert_equal(np.count_nonzero(m),
-                             expected, err_msg=err_msg)
-
-    def test_count_nonzero_axis_consistent(self):
-        # Check that the axis behaviour for valid axes in
-        # non-special cases is consistent (and therefore
-        # correct) by checking it against an integer array
-        # that is then casted to the generic object dtype
-        from itertools import combinations, permutations
-
-        axis = (0, 1, 2, 3)
-        size = (5, 5, 5, 5)
-        msg = "Mismatch for axis: %s"
-
-        rng = np.random.RandomState(1234)
-        m = rng.randint(-100, 100, size=size)
-        n = m.astype(object)
-
-        for length in range(len(axis)):
-            for combo in combinations(axis, length):
-                for perm in permutations(combo):
-                    assert_equal(
-                        np.count_nonzero(m, axis=perm),
-                        np.count_nonzero(n, axis=perm),
-                        err_msg=msg % (perm,))
+
+        m = np.zeros((3, 3), dtype=typecode)
+        n = np.ones(1, dtype=typecode)
+
+        m[0, 0] = n[0]
+        m[1, 0] = n[0]
+
+        expected = np.array([2, 0, 0], dtype=np.intp)
+        result = np.count_nonzero(m, axis=0)
+        assert_array_equal(result, expected)
+        assert expected.dtype == result.dtype
+
+        expected = np.array([1, 1, 0], dtype=np.intp)
+        result = np.count_nonzero(m, axis=1)
+        assert_array_equal(result, expected)
+        assert expected.dtype == result.dtype
+
+        expected = np.array(2)
+        assert_array_equal(np.count_nonzero(m, axis=(0, 1)), expected)
+        assert_array_equal(np.count_nonzero(m, axis=None), expected)
+        assert_array_equal(np.count_nonzero(m), expected)
 
     def test_countnonzero_axis_empty(self):
         a = np.array([[0, 0, 1], [1, 0, 1]])
         assert_equal(np.count_nonzero(a, axis=()), a.astype(bool))
 
     def test_countnonzero_keepdims(self):
-        a = np.array([[0, 0, 1, 0],
-                      [0, 3, 5, 0],
-                      [7, 9, 2, 0]])
-        assert_equal(np.count_nonzero(a, axis=0, keepdims=True),
-                     [[1, 2, 3, 0]])
-        assert_equal(np.count_nonzero(a, axis=1, keepdims=True),
-                     [[1], [2], [3]])
-        assert_equal(np.count_nonzero(a, keepdims=True),
-                     [[6]])
-
-    def test_array_method(self):
-        # Tests that the array method
-        # call to nonzero works
-        m = np.array([[1, 0, 0], [4, 0, 6]])
-        tgt = [[0, 1, 1], [0, 0, 2]]
-
-        assert_equal(m.nonzero(), tgt)
-
-        class BoolErrors:
-            def __bool__(self):
-                raise ValueError("Not allowed")
-
-        assert_raises(ValueError, np.nonzero, np.array([BoolErrors()]))
-
-    def test_nonzero_sideeffect_safety(self):
-        # gh-13631
-        class FalseThenTrue:
-            _val = False
-            def __bool__(self):
-                try:
-                    return self._val
-                finally:
-                    self._val = True
-
-        class TrueThenFalse:
-            _val = True
-            def __bool__(self):
-                try:
-                    return self._val
-                finally:
-                    self._val = False
-
-        # result grows on the second pass
-        a = np.array([True, FalseThenTrue()])
-        assert_raises(RuntimeError, np.nonzero, a)
-
-        a = np.array([[True], [FalseThenTrue()]])
-        assert_raises(RuntimeError, np.nonzero, a)
-
-        # result shrinks on the second pass
-        a = np.array([False, TrueThenFalse()])
-        assert_raises(RuntimeError, np.nonzero, a)
-
-        a = np.array([[False], [TrueThenFalse()]])
-        assert_raises(RuntimeError, np.nonzero, a)
-
-    def test_nonzero_sideffects_structured_void(self):
-        # Checks that structured void does not mutate alignment flag of
-        # original array.
-        arr = np.zeros(5, dtype="i1,i8,i8")  # `ones` may short-circuit
-        assert arr.flags.aligned  # structs are considered "aligned"
-        assert not arr["f2"].flags.aligned
-        # make sure that nonzero/count_nonzero do not flip the flag:
-        np.nonzero(arr)
-        assert arr.flags.aligned
-        np.count_nonzero(arr)
-        assert arr.flags.aligned
-
-    def test_nonzero_exception_safe(self):
-        # gh-13930
-
-        class ThrowsAfter:
-            def __init__(self, iters):
-                self.iters_left = iters
-
-            def __bool__(self):
-                if self.iters_left == 0:
-                    raise ValueError("called `iters` times")
-
-                self.iters_left -= 1
-                return True
-
-        """
-        Test that a ValueError is raised instead of a SystemError
-
-        If the __bool__ function is called after the error state is set,
-        Python (cpython) will raise a SystemError.
-        """
-
-        # assert that an exception in first pass is handled correctly
-        a = np.array([ThrowsAfter(5)]*10)
-        assert_raises(ValueError, np.nonzero, a)
-
-        # raise exception in second pass for 1-dimensional loop
-        a = np.array([ThrowsAfter(15)]*10)
-        assert_raises(ValueError, np.nonzero, a)
-
-        # raise exception in second pass for n-dimensional loop
-        a = np.array([[ThrowsAfter(15)]]*10)
-        assert_raises(ValueError, np.nonzero, a)
-
-    @pytest.mark.skipif(IS_WASM, reason="wasm doesn't have threads")
-    def test_structured_threadsafety(self):
-        # Nonzero (and some other functions) should be threadsafe for
-        # structured datatypes, see gh-15387. This test can behave randomly.
-        from concurrent.futures import ThreadPoolExecutor
-
-        # Create a deeply nested dtype to make a failure more likely:
-        dt = np.dtype([("", "f8")])
-        dt = np.dtype([("", dt)])
-        dt = np.dtype([("", dt)] * 2)
-        # The array should be large enough to likely run into threading issues
-        arr = np.random.uniform(size=(5000, 4)).view(dt)[:, 0]
-        def func(arr):
-            arr.nonzero()
-
-        tpe = ThreadPoolExecutor(max_workers=8)
-        futures = [tpe.submit(func, arr) for _ in range(10)]
-        for f in futures:
-            f.result()
-
-        assert arr.dtype is dt
+        a = np.array([[0, 0, 1, 0], [0, 3, 5, 0], [7, 9, 2, 0]])
+        assert_array_equal(np.count_nonzero(a, axis=0, keepdims=True), [[1, 2, 3, 0]])
+        assert_array_equal(np.count_nonzero(a, axis=1, keepdims=True), [[1], [2], [3]])
+        assert_array_equal(np.count_nonzero(a, keepdims=True), [[6]])
+        assert isinstance(np.count_nonzero(a, axis=1, keepdims=True), np.ndarray)
 
 
 @pytest.mark.xfail(reason="TODO")
diff --git a/torch_np/tests/test_reductions.py b/torch_np/tests/test_reductions.py
index 86c69313..e7456810 100644
--- a/torch_np/tests/test_reductions.py
+++ b/torch_np/tests/test_reductions.py
@@ -11,131 +11,6 @@
 )
 
 
-class TestNonzeroAndCountNonzero:
-    def test_count_nonzero_list(self):
-        lst = [[0, 1, 2, 3], [1, 0, 0, 6]]
-        assert np.count_nonzero(lst) == 5
-        assert_array_equal(np.count_nonzero(lst, axis=0), np.array([1, 1, 1, 2]))
-        assert_array_equal(np.count_nonzero(lst, axis=1), np.array([3, 2]))
-
-    def test_nonzero_trivial(self):
-        assert_equal(np.count_nonzero(np.array([])), 0)
-        assert_equal(np.count_nonzero(np.array([], dtype="?")), 0)
-        assert_equal(np.nonzero(np.array([])), ([],))
-
-        assert_equal(np.count_nonzero(np.array([0])), 0)
-        assert_equal(np.count_nonzero(np.array([0], dtype="?")), 0)
-        assert_equal(np.nonzero(np.array([0])), ([],))
-
-        assert_equal(np.count_nonzero(np.array([1])), 1)
-        assert_equal(np.count_nonzero(np.array([1], dtype="?")), 1)
-        assert_equal(np.nonzero(np.array([1])), ([0],))
-
-        assert isinstance(np.count_nonzero([]), np.ndarray)
-
-    def test_nonzero_zerod(self):
-        assert_equal(np.count_nonzero(np.array(0)), 0)
-        assert_equal(np.count_nonzero(np.array(0, dtype="?")), 0)
-
-        assert_equal(np.count_nonzero(np.array(1)), 1)
-        assert_equal(np.count_nonzero(np.array(1, dtype="?")), 1)
-
-        assert isinstance(np.count_nonzero(np.array(1)), np.ndarray)
-
-    def test_nonzero_onedim(self):
-        x = np.array([1, 0, 2, -1, 0, 0, 8])
-        assert_equal(np.count_nonzero(x), 4)
-        assert_equal(np.count_nonzero(x), 4)
-        assert_equal(np.nonzero(x), ([0, 2, 3, 6],))
-
-        assert isinstance(np.count_nonzero(x), np.ndarray)
-
-    def test_nonzero_twodim(self):
-        x = np.array([[0, 1, 0], [2, 0, 3]])
-        assert_equal(np.count_nonzero(x.astype("i1")), 3)
-        assert_equal(np.count_nonzero(x.astype("i2")), 3)
-        assert_equal(np.count_nonzero(x.astype("i4")), 3)
-        assert_equal(np.count_nonzero(x.astype("i8")), 3)
-        assert_equal(np.nonzero(x), ([0, 1, 1], [1, 0, 2]))
-
-        x = np.eye(3)
-        assert_equal(np.count_nonzero(x.astype("i1")), 3)
-        assert_equal(np.count_nonzero(x.astype("i2")), 3)
-        assert_equal(np.count_nonzero(x.astype("i4")), 3)
-        assert_equal(np.count_nonzero(x.astype("i8")), 3)
-        assert_equal(np.nonzero(x), ([0, 1, 2], [0, 1, 2]))
-
-    def test_sparse(self):
-        # test special sparse condition boolean code path
-        for i in range(20):
-            c = np.zeros(200, dtype=bool)
-            c[i::20] = True
-            assert_equal(np.nonzero(c)[0], np.arange(i, 200 + i, 20))
-
-            c = np.zeros(400, dtype=bool)
-            c[10 + i : 20 + i] = True
-            c[20 + i * 2] = True
-            assert_equal(
-                np.nonzero(c)[0],
-                np.concatenate((np.arange(10 + i, 20 + i), [20 + i * 2])),
-            )
-
-    def test_count_nonzero_axis(self):
-        # Basic check of functionality
-        m = np.array([[0, 1, 7, 0, 0], [3, 0, 0, 2, 19]])
-
-        expected = np.array([1, 1, 1, 1, 1])
-        assert_array_equal(np.count_nonzero(m, axis=0), expected)
-
-        expected = np.array([2, 3])
-        assert_array_equal(np.count_nonzero(m, axis=1), expected)
-
-        assert isinstance(np.count_nonzero(m, axis=1), np.ndarray)
-
-        assert_raises(ValueError, np.count_nonzero, m, axis=(1, 1))
-        assert_raises(TypeError, np.count_nonzero, m, axis="foo")
-        assert_raises(np.AxisError, np.count_nonzero, m, axis=3)
-        assert_raises(TypeError, np.count_nonzero, m, axis=np.array([[1], [2]]))
-
-    @pytest.mark.parametrize("typecode", np.typecodes["All"])
-    def test_count_nonzero_axis_all_dtypes(self, typecode):
-        # More thorough test that the axis argument is respected
-        # for all dtypes and responds correctly when presented with
-        # either integer or tuple arguments for axis
-
-        m = np.zeros((3, 3), dtype=typecode)
-        n = np.ones(1, dtype=typecode)
-
-        m[0, 0] = n[0]
-        m[1, 0] = n[0]
-
-        expected = np.array([2, 0, 0], dtype=np.intp)
-        result = np.count_nonzero(m, axis=0)
-        assert_array_equal(result, expected)
-        assert expected.dtype == result.dtype
-
-        expected = np.array([1, 1, 0], dtype=np.intp)
-        result = np.count_nonzero(m, axis=1)
-        assert_array_equal(result, expected)
-        assert expected.dtype == result.dtype
-
-        expected = np.array(2)
-        assert_array_equal(np.count_nonzero(m, axis=(0, 1)), expected)
-        assert_array_equal(np.count_nonzero(m, axis=None), expected)
-        assert_array_equal(np.count_nonzero(m), expected)
-
-    def test_countnonzero_axis_empty(self):
-        a = np.array([[0, 0, 1], [1, 0, 1]])
-        assert_equal(np.count_nonzero(a, axis=()), a.astype(bool))
-
-    def test_countnonzero_keepdims(self):
-        a = np.array([[0, 0, 1, 0], [0, 3, 5, 0], [7, 9, 2, 0]])
-        assert_array_equal(np.count_nonzero(a, axis=0, keepdims=True), [[1, 2, 3, 0]])
-        assert_array_equal(np.count_nonzero(a, axis=1, keepdims=True), [[1], [2], [3]])
-        assert_array_equal(np.count_nonzero(a, keepdims=True), [[6]])
-        assert isinstance(np.count_nonzero(a, axis=1, keepdims=True), np.ndarray)
-
-
 class TestFlatnonzero:
     def test_basic(self):
         x = np.arange(-2, 3)

From 679e0535319673e10fd1bf10f9d3c8627eebd597 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 4 Feb 2023 15:33:27 +0300
Subject: [PATCH 7/8] lint

---
 torch_np/_detail/_flips.py | 3 ++-
 torch_np/_ndarray.py       | 2 +-
 torch_np/_wrapper.py       | 9 ++++++---
 3 files changed, 9 insertions(+), 5 deletions(-)

diff --git a/torch_np/_detail/_flips.py b/torch_np/_detail/_flips.py
index a46aca5d..bf12841d 100644
--- a/torch_np/_detail/_flips.py
+++ b/torch_np/_detail/_flips.py
@@ -5,6 +5,7 @@
 
 from . import _scalar_types, _util
 
+
 def flip(m_tensor, axis=None):
     # XXX: semantic difference: np.flip returns a view, torch.flip copies
     if axis is None:
@@ -39,7 +40,7 @@ def rollaxis(tensor, axis, start=0):
         start += n
     msg = "'%s' arg requires %d <= %s < %d, but %d was passed in"
     if not (0 <= start < n + 1):
-        raise _util.AxisError(msg % ('start', -n, 'start', n + 1, start))
+        raise _util.AxisError(msg % ("start", -n, "start", n + 1, start))
     if axis < start:
         # it's been removed
         start -= 1
diff --git a/torch_np/_ndarray.py b/torch_np/_ndarray.py
index 45b79af7..7af5da80 100644
--- a/torch_np/_ndarray.py
+++ b/torch_np/_ndarray.py
@@ -10,7 +10,7 @@
     dtype_to_torch,
     emulate_out_arg,
 )
-from ._detail import _reductions, _util, _flips
+from ._detail import _flips, _reductions, _util
 
 newaxis = None
 
diff --git a/torch_np/_wrapper.py b/torch_np/_wrapper.py
index a48b9c85..d9768159 100644
--- a/torch_np/_wrapper.py
+++ b/torch_np/_wrapper.py
@@ -8,7 +8,7 @@
 import torch
 
 from . import _dtypes, _helpers
-from ._detail import _reductions, _util, _flips
+from ._detail import _flips, _reductions, _util
 from ._ndarray import (
     array,
     asarray,
@@ -476,8 +476,8 @@ def broadcast_arrays(*args, subok=False):
 
 @asarray_replacer()
 def moveaxis(a, source, destination):
-    source = _util.normalize_axis_tuple(source, a.ndim, 'source')
-    destination = _util.normalize_axis_tuple(destination, a.ndim, 'destination')
+    source = _util.normalize_axis_tuple(source, a.ndim, "source")
+    destination = _util.normalize_axis_tuple(destination, a.ndim, "destination")
     return asarray(torch.moveaxis(a, source, destination))
 
 
@@ -485,6 +485,7 @@ def swapaxis(a, axis1, axis2):
     arr = asarray(a)
     return arr.swapaxes(axis1, axis2)
 
+
 @asarray_replacer()
 def rollaxis(a, axis, start=0):
     return _flips.rollaxis(a, axis, start)
@@ -665,8 +666,10 @@ def prod(
         axis=axis, dtype=dtype, out=out, keepdims=keepdims, initial=initial, where=where
     )
 
+
 product = prod
 
+
 def cumprod(a, axis=None, dtype=None, out=None):
     arr = asarray(a)
     return arr.cumprod(axis=axis, dtype=dtype, out=out)

From ded65f293131c7101636e681ae788c8f9b26b500 Mon Sep 17 00:00:00 2001
From: Evgeni Burovski <evgeny.burovskiy@gmail.com>
Date: Sat, 4 Feb 2023 17:20:26 +0300
Subject: [PATCH 8/8] MAINT: address review comments

---
 torch_np/_detail/_flips.py                           | 9 ++++++++-
 torch_np/tests/numpy_tests/lib/test_function_base.py | 4 ++--
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/torch_np/_detail/_flips.py b/torch_np/_detail/_flips.py
index bf12841d..c343571d 100644
--- a/torch_np/_detail/_flips.py
+++ b/torch_np/_detail/_flips.py
@@ -32,7 +32,12 @@ def swapaxes(tensor, axis1, axis2):
     return torch.swapaxes(tensor, axis1, axis2)
 
 
+# Straight vendor from:
 # https://github.com/numpy/numpy/blob/v1.24.0/numpy/core/numeric.py#L1259
+#
+# Also note this function in NumPy is mostly retained for backwards compat
+# (https://stackoverflow.com/questions/29891583/reason-why-numpy-rollaxis-is-so-confusing)
+# so let's not touch it unless hard pressed.
 def rollaxis(tensor, axis, start=0):
     n = tensor.ndim
     axis = _util.normalize_axis_index(axis, n)
@@ -45,7 +50,9 @@ def rollaxis(tensor, axis, start=0):
         # it's been removed
         start -= 1
     if axis == start:
-        return tensor[...]
+        # numpy returns a view, here we try returning the tensor itself
+        # return tensor[...]
+        return tensor
     axes = list(range(0, n))
     axes.remove(axis)
     axes.insert(start, axis)
diff --git a/torch_np/tests/numpy_tests/lib/test_function_base.py b/torch_np/tests/numpy_tests/lib/test_function_base.py
index 544420f7..df14dbbf 100644
--- a/torch_np/tests/numpy_tests/lib/test_function_base.py
+++ b/torch_np/tests/numpy_tests/lib/test_function_base.py
@@ -130,7 +130,7 @@ def test_axes(self):
         assert_raises(np.AxisError, np.flip, np.ones((4, 4)), axis=-3)
         assert_raises(np.AxisError, np.flip, np.ones((4, 4)), axis=(0, 3))
 
-    @pytest.mark.xfail(reason='no [::-1] indexing')
+    @pytest.mark.skip(reason='no [::-1] indexing')
     def test_basic_lr(self):
         a = get_mat(4)
         b = a[:, ::-1]
@@ -141,7 +141,7 @@ def test_basic_lr(self):
              [5, 4, 3]]
         assert_equal(np.flip(a, 1), b)
 
-    @pytest.mark.xfail(reason='no [::-1] indexing')
+    @pytest.mark.skip(reason='no [::-1] indexing')
     def test_basic_ud(self):
         a = get_mat(4)
         b = a[::-1, :]