Repository navigation
Fix np.clip returning wrong shape when bounds broadcast larger than a - #10767
VenkateswarluNagineni wants to merge 5 commits into
Conversation
|
@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! |
|
Thanks for the flag — I used Claude Sonnet 4.6 (claude-sonnet-4-6) throughout this PR: to identify the root cause in |
|
@VenkateswarluNagineni can you rebase this on |
Awesome, thank you for updating! |
bf68b8c to
d5ddb06
Compare
| @@ -0,0 +1,4 @@ | |||
| PR #10767: Fix ``np.clip`` returning wrong shape when ``a_min`` or ``a_max`` | |||
There was a problem hiding this comment.
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
|
@VenkateswarluNagineni are you still interested in completing this PR? |
|
Yes, still on it — just rebased on main. Force pushing the branch now. |
d5ddb06 to
d1d3aef
Compare
Excellent! |
|
@VenkateswarluNagineni looks like the release notes snippet still has the wrong underline, can you fix? |
esc
left a comment
There was a problem hiding this comment.
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
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>
|
@VenkateswarluNagineni ping here, are you still available to fix this up? |
|
Fixed, thanks for the detailed repro — that made it easy to reason through and verify. Root cause: switching Fix: Built numba locally against Matches pre-regression behaviour. Also ran the full |
esc
left a comment
There was a problem hiding this comment.
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.
|
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:
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
left a comment
There was a problem hiding this comment.
@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.
Problem
np.clip(a, a_min, a_max)returns the wrong shape whena_minora_maxbroadcasts to a larger shape thana. NumPy returns the broadcast shape; Numba returnsa.shape.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 thana, which is a memory-safety issue.This is the same class of defect as #10682 on the
out=Noneallocation path. PR #10761 fixed theout=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 callingnp.broadcast_arrays:_np_clip_impl— both bounds are arrays (coversnp_clip_aa,np_clip_sa,np_clip_as)np_clip_na—a_minisNone,a_maxis an arraynp_clip_an—a_maxisNone,a_minis an arrayFix
Move the
np.broadcast_arrays(...)call before the result allocation in each path, then allocate withnp.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 onmainand pass with this change.Closes #10760