Skip to content

Fix np.clip returning wrong shape when bounds broadcast larger than a - #10767

Open
VenkateswarluNagineni wants to merge 5 commits into
numba:mainfrom
VenkateswarluNagineni:fix/clip-broadcast-shape
Open

VenkateswarluNagineni wants to merge 5 commits into
numba:mainfrom
VenkateswarluNagineni:fix/clip-broadcast-shape

Conversation

@VenkateswarluNagineni

Copy link
Copy Markdown

Problem

np.clip(a, a_min, a_max) returns the wrong shape when a_min or a_max broadcasts to a larger shape than a. NumPy returns the broadcast shape; Numba returns a.shape.

import numpy as np
from numba import njit

@njit
def clip(a, lo, hi): return np.clip(a, lo, hi)

a   = np.arange(3.0).reshape(3, 1)
big = np.zeros((3, 4))

print(np.clip(a, big, 3.0).shape)  # (3, 4)  — correct
print(clip(a, big, 3.0).shape)     # (3, 1)  — wrong

Beyond returning the wrong shape, the fill loop (np.ndindex(a_b.shape)) iterates over the broadcast shape and writes past the end of the result buffer when the broadcast is larger than a, which is a memory-safety issue.

This is the same class of defect as #10682 on the out=None allocation path. PR #10761 fixed the out= argument path and explicitly left this case out of scope.

Root cause

Three code paths allocate the result with np.empty_like(a) (shape = a.shape) before calling np.broadcast_arrays:

  • _np_clip_impl — both bounds are arrays (covers np_clip_aa, np_clip_sa, np_clip_as)
  • np_clip_na — a_min is None, a_max is an array
  • np_clip_an — a_max is None, a_min is an array

Fix

Move the np.broadcast_arrays(...) call before the result allocation in each path, then allocate with np.empty(a_b.shape, a.dtype).

Verification

Three new regression tests in numba/tests/test_array_methods.py (test_clip_broadcast_shape), one per affected code path. Each confirms that the result shape and values match NumPy. All three fail on main and pass with this change.

Closes #10760

VenkateswarluNagineni added a commit to VenkateswarluNagineni/numba that referenced this pull request Aug 8, 2026
@esc

esc commented Aug 11, 2026

Copy link
Copy Markdown
Member

@VenkateswarluNagineni this looks like an LLM was used to asssist with this PR? Please label the model used to conform with our AI Policy: https://numba.readthedocs.io/en/latest/reference/ai_tools_policy.html thanks!

@esc esc added more info needed This issue needs more information 2 - In Progress labels Aug 11, 2026
@esc
esc self-requested a review August 11, 2026 14:23
@VenkateswarluNagineni

Copy link
Copy Markdown
Author

Thanks for the flag — I used Claude Sonnet 4.6 (claude-sonnet-4-6) throughout this PR: to identify the root cause in arrayobj.py, write the three-site fix, and draft the regression tests and changelog. All code was reviewed and submitted by me.

@esc

esc commented Aug 19, 2026

Copy link
Copy Markdown
Member

@VenkateswarluNagineni can you rebase this on main and fixup the conflict? #10761 has been merged and so this PT now conflicts, thank you!

@esc esc removed the more info needed This issue needs more information label Aug 19, 2026
@esc

esc commented Aug 19, 2026

Copy link
Copy Markdown
Member

Thanks for the flag — I used Claude Sonnet 4.6 (claude-sonnet-4-6) throughout this PR: to identify the root cause in arrayobj.py, write the three-site fix, and draft the regression tests and changelog. All code was reviewed and submitted by me.

Awesome, thank you for updating!

VenkateswarluNagineni added a commit to VenkateswarluNagineni/numba that referenced this pull request Aug 20, 2026
Comment thread docs/upcoming_changes/10767.bug_fix.rst Outdated
@@ -0,0 +1,4 @@
PR #10767: Fix ``np.clip`` returning wrong shape when ``a_min`` or ``a_max``

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The snippet needs a title with a hyphen underscore, as per:

https://github.com/numba/numba/blob/main/docs/upcoming_changes/README.rst

When a_min or a_max is an array that broadcasts to a shape larger than a,
np.clip was allocating the result with np.empty_like(a) (shape = a.shape)
before computing the broadcast, while the fill loop ran over the broadcast
shape a_b.shape. When the broadcast is larger, the loop wrote past the end
of the result buffer and the returned array had the wrong shape.

Fix: compute np.broadcast_arrays first in all three affected code paths,
then allocate the result with the broadcast shape.

Fixes numba#10760
@esc

esc commented Sep 10, 2026

Copy link
Copy Markdown
Member

@VenkateswarluNagineni are you still interested in completing this PR?

@VenkateswarluNagineni

Copy link
Copy Markdown
Author

Yes, still on it — just rebased on main. Force pushing the branch now.

@esc

esc commented Sep 11, 2026

Copy link
Copy Markdown
Member

Yes, still on it — just rebased on main. Force pushing the branch now.

Excellent!

@esc esc added 4 - Waiting on reviewer Waiting for reviewer to respond to author and removed 2 - In Progress labels Sep 11, 2026
@esc esc added this to the 0.69.0rc1 milestone Sep 11, 2026
@esc

esc commented Sep 11, 2026

Copy link
Copy Markdown
Member

@VenkateswarluNagineni looks like the release notes snippet still has the wrong underline, can you fix?

@esc esc left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR has some serious issues and introduces regressions in it's current state:

import numpy as np
from numba import jit, float64

@jit
def f(a):
    return np.clip(a, 0., 1.)

a = np.asfortranarray(np.zeros((3, 4)))
r = f(a)
print("np.clip F-order preserved:", np.clip(a, 0., 1.).flags.f_contiguous,
      "(NumPy reference)")
print("numba  F-order preserved:", r.flags.f_contiguous,
      "| return type:", f.nopython_signatures[0].return_type)

try:
    @jit(float64[::1, :](float64[::1, :]))
    def g(a):
        return np.clip(a, 0., 1.)
    g(a)
    print("layout-pinned caller float64[::1, :]: compiles")
except Exception:
    print("layout-pinned caller float64[::1, :]: TypingError")

Before (main, 0.69.0dev0):

np.clip F-order preserved: True (NumPy reference)
numba  F-order preserved: True | return type: array(float64, 2d, F)
layout-pinned caller float64[::1, :]: compiles

After (PR #10767 @ e48ce6b):

np.clip F-order preserved: True (NumPy reference)
numba  F-order preserved: False | return type: array(float64, 2d, C)
layout-pinned caller float64[::1, :]: TypingError

Looks like the order would not be preserved after update?

The previous fix (np.empty_like → np.empty) always produced C-order
output, breaking F-contiguous inputs (reported by esc in review).

Replace the @register_jitable helper with an @overload that dispatches
at the type level: F-order inputs use np.asfortranarray(np.empty(...))
to allocate an F-order result of the correct broadcast shape; C/A-order
inputs keep the existing np.empty path.

Add test_clip_fortran_order covering scalar bounds, one-None-bound, and
broadcast-larger-than-a paths.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
@esc esc added 4 - Waiting on author Waiting for author to respond to review and removed 4 - Waiting on reviewer Waiting for reviewer to respond to author labels Sep 15, 2026
@esc

esc commented Sep 21, 2026

Copy link
Copy Markdown
Member

@VenkateswarluNagineni ping here, are you still available to fix this up?

@VenkateswarluNagineni

VenkateswarluNagineni commented Sep 21, 2026 •

Copy link
Copy Markdown
Author

Fixed, thanks for the detailed repro — that made it easy to reason through and verify.

Root cause: switching _np_clip_prepare_out from np.empty_like(a) to np.empty(shape, a.dtype) fixed the shape bug but lost layout preservation, since np.empty's overload always returns C-order at the type level.

Fix: _np_clip_prepare_out is now an @overload that dispatches on a.layout at compile time — F-contiguous inputs go through np.asfortranarray(np.empty(shape, a.dtype)) to get an F-order result of the correct (possibly broadcast) shape, everything else keeps the plain np.empty path.

Built numba locally against ffd7f754d and ran your repro verbatim:

np.clip F-order preserved: True (NumPy reference)
numba  F-order preserved: True | return type: array(float64, 2d, F)
layout-pinned caller float64[::1, :]: compiles

Matches pre-regression behaviour. Also ran the full test_array_methods.py -k clip suite (14 tests, all passing), including the new test_clip_fortran_order covering the scalar-bounds, one-None-bound, and broadcast-larger-than-a paths for F-contiguous input.

@esc esc left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The changelog entry makes it sound like the PR fixed a layout regression. This is clearly not the case as I demonstrated during my previous review.

Please fix this and take more care when reviewing the output of your coding agent before git push in future.

@esc esc removed this from the 0.69.0rc1 milestone Sep 22, 2026
@VenkateswarluNagineni

Copy link
Copy Markdown
Author

Thanks for the correction. I updated the changelog so it no longer frames the PR as fixing a layout regression; it now describes the user-facing broadcast-shape fix and notes that NumPy-compatible F-contiguous layout behavior is preserved.

I also tightened the Fortran-order regression test so the array-bounds case actually broadcasts to a larger shape with a non-degenerate F-contiguous input.

Validation run locally on commit 8d5ea14:

  • python -m pytest numba/tests/test_array_methods.py -k clip -q: 14 passed, 74 deselected
  • python maint/towncrier_rst_validator.py --pull_request_id 10767: passed, including rstcheck
  • esc's F-order/layout-pinned repro: Numba preserves F-order, return type is array(float64, 2d, F), and float64[::1, :] caller compiles

For AI-policy completeness: this latest cleanup/check pass was assisted by OpenAI Codex (GPT-5); I reviewed the diff and validation output before pushing.

@esc esc left a comment

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@VenkateswarluNagineni thank you for taking the time to fix this up.

I have thought about this and I believe that this feature can not be implemented in a clean way due a gap in the current Numba support of the NumPy API.

What would make this much easier, cleaner and straightforward to implement is:

return np.empty_like(a, shape=shape)

But:

In [2]: import numpy as np
   ...: from numba import jit
   ...:
   ...: @jit
   ...: def f(a):
   ...:     return np.empty_like(a, shape=(3, 4))
   ...:
   ...: f(np.zeros((3, 1)))
---------------------------------------------------------------------------
TypingError                               Traceback (most recent call last)
Cell In[2], line 8
      4 @jit
      5 def f(a):
      6     return np.empty_like(a, shape=(3, 4))
----> 8 f(np.zeros((3, 1)))

File ~/git/numba/numba/core/dispatcher.py:422, in _DispatcherBase._compile_for_args(self, *args, **kws)
    418         msg = (f"{str(e).rstrip()} \n\nThis error may have been caused "
    419                f"by the following argument(s):\n{args_str}\n")
    420         e.patch_message(msg)
--> 422     error_rewrite(e, 'typing')
    423 except errors.UnsupportedError as e:
    424     # Something unsupported is present in the user code, add help info
    425     error_rewrite(e, 'unsupported_error')

File ~/git/numba/numba/core/dispatcher.py:363, in _DispatcherBase._compile_for_args.<locals>.error_rewrite(e, issue_type)
    361     raise e
    362 else:
--> 363     raise e.with_traceback(None)

TypingError: Failed in nopython mode pipeline (step: nopython frontend)
No implementation of function Function(<built-in function empty_like>) found for signature:

 >>> empty_like(array(float64, 2d, C), shape=Tuple(Literal[int](3), Literal[int](4)))

There are 2 candidate implementations:
  - Of which 2 did not match due to:
  Overload in function 'ol_np_empty_like': File: numba/np/arrayobj.py: Line 4817.
    With argument(s): '(array(float64, 2d, C), shape=UniTuple(int64 x 2))':
   Rejected as the implementation raised a specific error:
     TypingError: got an unexpected keyword argument 'shape'
  raised from /Users/esc/git/numba/numba/core/typing/templates.py:791

During: resolving callee type: Function(<built-in function empty_like>)
During: typing of call at <ipython-input-2-937834765156> (6)

File "<ipython-input-2-937834765156>", line 6:
def f(a):
    return np.empty_like(a, shape=(3, 4))
    ^

During: Pass nopython_type_inference

So my suggestion now has changed a bit and actually, I would put this PR on hold and add shape to empty_like instead or wait until someone else does this. It does appear to more involved that what was done for this PR. The rebase this PR such that the diff then becomes minimal.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

4 - Waiting on author Waiting for author to respond to review

Projects

None yet

Development

Successfully merging this pull request may close these issues.

np.clip returns an array with the wrong shape when a_min or a_max broadcasts to a larger result

2 participants