BUG: ensure i0 does not change the shape.

3 files changed, 27 insertions, 16 deletions

diff --git a/doc/release/1.17.0-notes.rst b/doc/release/1.17.0-notes.rst
index 0007fe1a6..58d38cba9 100644
--- a/doc/release/1.17.0-notes.rst
+++ b/doc/release/1.17.0-notes.rst
@@ -124,6 +124,17 @@ the stream changes for any given seed is extremely small. If a 0 is encountered
 underlying generator, then the incorrect value produced (either ``np.inf``
 or ``np.nan``) is now dropped.
 
+``i0`` now always returns a result with the same shape as the input
+-------------------------------------------------------------------
+Previously, the output was squeezed, such that, e.g., input with just a single
+element would lead to an array scalar being returned, and inputs with shapes
+such as ``(10, 1)`` would yield results that would not broadcast against the
+input.
+
+Note that we generally recommend the scipy implementation over the numpy one:
+it is a proper ufunc written in C, and more than an order of magnitude faster.
+
+
 C API changes
 =============
 
diff --git a/numpy/lib/function_base.py b/numpy/lib/function_base.py
index fabb87adc..9690c305f 100644
--- a/numpy/lib/function_base.py
+++ b/numpy/lib/function_base.py
@@ -3070,10 +3070,13 @@ def i0(x):
 
     See Also
     --------
-    scipy.special.iv, scipy.special.ive
+    scipy.special.i0, scipy.special.iv, scipy.special.ive
 
     Notes
     -----
+    The scipy implementation is recommended over this function: it is a
+    proper ufunc written in C, and more than an order of magnitude faster.
+
     We use the algorithm published by Clenshaw [1]_ and referenced by
     Abramowitz and Stegun [2]_, for which the function domain is
     partitioned into the two intervals [0,8] and (8,inf), and Chebyshev
@@ -3093,21 +3096,14 @@ def i0(x):
 
     Examples
     --------
-    >>> np.i0([0.])
-    array(1.0) # may vary
+    >>> np.i0(0.)
+    array(1.0)  # may vary
     >>> np.i0([0., 1. + 2j])
-    array([ 1.00000000+0.j        ,  0.18785373+0.64616944j]) # may vary
+    array([ 1.00000000+0.j        ,  0.18785373+0.64616944j])  # may vary
 
     """
-    x = atleast_1d(x).copy()
-    y = empty_like(x)
-    ind = (x < 0)
-    x[ind] = -x[ind]
-    ind = (x <= 8.0)
-    y[ind] = _i0_1(x[ind])
-    ind2 = ~ind
-    y[ind2] = _i0_2(x[ind2])
-    return y.squeeze()
+    x = np.abs(x)
+    return piecewise(x, [x <= 8.0], [_i0_1, _i0_2])
 
 ## End of cephes code for i0
 
diff --git a/numpy/lib/tests/test_function_base.py b/numpy/lib/tests/test_function_base.py
index a3d4c6efb..93ebabae0 100644
--- a/numpy/lib/tests/test_function_base.py
+++ b/numpy/lib/tests/test_function_base.py
@@ -1984,9 +1984,9 @@ class Test_I0(object):
             np.array(1.0634833707413234))
 
         A = np.array([0.49842636, 0.6969809, 0.22011976, 0.0155549])
-        assert_almost_equal(
-            i0(A),
-            np.array([1.06307822, 1.12518299, 1.01214991, 1.00006049]))
+        expected = np.array([1.06307822, 1.12518299, 1.01214991, 1.00006049])
+        assert_almost_equal(i0(A), expected)
+        assert_almost_equal(i0(-A), expected)
 
         B = np.array([[0.827002, 0.99959078],
                       [0.89694769, 0.39298162],
@@ -2000,6 +2000,10 @@ class Test_I0(object):
                       [1.03633899, 1.00067775],
                       [1.03352052, 1.13557954],
                       [1.05884290, 1.06432317]]))
+        # Regression test for gh-11205
+        i0_0 = np.i0([0.])
+        assert_equal(i0_0.shape, (1,))
+        assert_array_equal(np.i0([0.]), np.array([1.]))
 
 
 class TestKaiser(object):