np.linalg.vector_norm: return correct shape for keepdims

jakevdp · jakevdp · commit e22c1b8d2a63 · 2024-01-09T12:11:03.000-08:00
diff --git a/numpy/linalg/_linalg.py b/numpy/linalg/_linalg.py
@@ -3306,6 +3306,7 @@ def vector_norm(x, /, *, axis=None, keepdims=False, ord=2):
 
     """
     x = asanyarray(x)
+    shape = list(x.shape)
     if axis is None:
         # Note: np.linalg.norm() doesn't handle 0-D arrays
         x = x.ravel()
@@ -3331,9 +3332,8 @@ def vector_norm(x, /, *, axis=None, keepdims=False, ord=2):
     if keepdims:
         # We can't reuse np.linalg.norm(keepdims) because of the reshape hacks
         # above to avoid matrix norm logic.
-        shape = list(x.shape)
         _axis = normalize_axis_tuple(
-            range(x.ndim) if axis is None else axis, x.ndim
+            range(len(shape)) if axis is None else axis, len(shape)
         )
         for i in _axis:
             shape[i] = 1
diff --git a/numpy/linalg/tests/test_linalg.py b/numpy/linalg/tests/test_linalg.py
@@ -2355,3 +2355,8 @@ def test_vector_norm():
     assert_almost_equal(
         actual, np.array([6.7082, 8.124, 9.6436]), double_decimal=3
     )
+
+    actual = np.linalg.vector_norm(x, keepdims=True)
+    expected = np.full((1, 1), 14.2828, dtype='float64')
+    assert_equal(actual.shape, expected.shape)
+    assert_almost_equal(actual, expected, double_decimal=3)

Original file line number	Diff line number	Diff line change
`@@ -2355,3 +2355,8 @@ def test_vector_norm():`
`2355`	`2355`	`assert_almost_equal(`
`2356`	`2356`	`actual, np.array([6.7082, 8.124, 9.6436]), double_decimal=3`
`2357`	`2357`	`)`
	`2358`	`+`
	`2359`	`+ actual = np.linalg.vector_norm(x, keepdims=True)`
	`2360`	`+ expected = np.full((1, 1), 14.2828, dtype='float64')`
	`2361`	`+ assert_equal(actual.shape, expected.shape)`
	`2362`	`+ assert_almost_equal(actual, expected, double_decimal=3)`