BUG: i0 Bessel function regression on array-likes supporting ufuncs

For array likes supporting UFuncs, `np.abs` would return an array-like,
and this is currently not compatible with the use of `np.piecewise`.
The simplest fix seems to be to just call asanyarray (which piecewise
calls anyway on the array) beforehand. This way we ensure the conditions
are also an array.

Fixes gh-13894

2 files changed, 17 insertions, 0 deletions

diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index 1fcb6137c..cf7246402 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -3102,6 +3102,7 @@ def i0(x):
     array([ 1.00000000+0.j        ,  0.18785373+0.64616944j])  # may vary
 
     """
+    x = np.asanyarray(x)
     x = np.abs(x)
     return piecewise(x, [x <= 8.0], [_i0_1, _i0_2])
 
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index c0b8ad6b8..eae52c002 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -2005,6 +2005,22 @@ class Test_I0(object):
         assert_equal(i0_0.shape, (1,))
         assert_array_equal(np.i0([0.]), np.array([1.]))
 
+    def test_non_array(self):
+        a = np.arange(4)
+
+        class array_like:
+            __array_interface__ = a.__array_interface__
+
+            def __array_wrap__(self, arr):
+                return self
+
+        # E.g. pandas series survive ufunc calls through array-wrap:
+        assert isinstance(np.abs(array_like()), array_like)
+        exp = np.i0(a)
+        res = np.i0(array_like())
+
+        assert_array_equal(exp, res)
+
 
 class TestKaiser(object):