Skip to content

Improve clipping time with np_compat.clip#447

Open
Lasloruhberg wants to merge 6 commits into
data-apis:mainfrom
Lasloruhberg:improve_clipping_time
Open

Improve clipping time with np_compat.clip#447
Lasloruhberg wants to merge 6 commits into
data-apis:mainfrom
Lasloruhberg:improve_clipping_time

Conversation

@Lasloruhberg

Copy link
Copy Markdown

This PR introduces a numpy specific alias to xp.clip to improve performance, adressing Issue #444.
With respect to the basic implementation I removed array copies and new creations where possible and revert to the intrinsic np.clip.

The benchmark also implies that dask clip is quite below its possible performance, as it remains well below the other programs.

As far as I can tell this rewrite is compliant with the API standard, but please advise if not.
I am happy to receive feedback.

Current on main (run on a laptop, no GPU):

Batch size: 1024
Image shape: (128, 128, 3)
Measurement runs: 8

Skipped backends: cupy (unavailable)

backend dtype avg ms std ms images/s std images/s
numpy-native float32 230.097 20.311 4450.3 357.8
numpy-native uint8 60.985 8.060 16790.9 1877.6
numpy float32 4765.966 298.603 214.9 13.6
numpy uint8 3194.393 266.443 320.6 27.0
dask float32 3279.046 513.540 312.3 49.4
dask uint8 1657.283 168.447 617.9 67.7
jax float32 197.261 71.699 5191.1 1549.5
jax uint8 38.581 7.265 26541.4 4950.1
pytorch float32 160.713 54.031 6371.6 1940.3
pytorch uint8 16.838 3.103 60814.3 13022.6

Now on this branch

Batch size: 1024
Image shape: (128, 128, 3)
Measurement runs: 15

backend dtype avg ms std ms images/s std images/s
numpy-native float32 156.005 27.627 6563.9 1260.4
numpy-native uint8 34.172 5.294 29965.8 4482.5
numpy float32 157.614 30.625 6496.9 1226.8
numpy uint8 41.640 7.193 24592.0 4832.1
dask float32 1675.356 204.533 611.2 81.3
dask uint8 1544.973 493.518 662.8 169.1
jax float32 81.755 30.772 12525.2 3866.9
jax uint8 42.276 66.549 24221.9 13959.8
pytorch float32 86.618 34.658 11822.1 3924.1
pytorch uint8 16.298 3.716 62829.0 12721.9

As evident the new clip version for numpy restores its native performance, while ensuring the same broadcasting and dtype preservation behaviour

Benchmark:
benchmark.py

Copilot AI review requested due to automatic review settings July 3, 2026 08:46

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Pull request overview

This PR replaces the NumPy backend’s clip alias with a NumPy-specific implementation that calls np.clip directly (after casting bounds to x.dtype) to restore near-native NumPy performance, and adds targeted regression tests for dtype/broadcasting/out behavior.

Changes:

  • Implement array_api_compat.numpy.clip as a dedicated wrapper around np.clip with dtype-preserving bound casting and reduced intermediate allocations.
  • Add NumPy-specific tests covering broadcasting, out= behavior, and dtype preservation (including edge cases with bounds outside representable ranges).
  • Minor formatting cleanups in numpy/_aliases.py (line wrapping / commas / __all__ updates).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 3 comments.

File Description
tests/test_numpy.py Adds tests for NumPy clip behavior (broadcasting, out, dtype preservation, and edge cases).
src/array_api_compat/numpy/_aliases.py Replaces the generic clip alias with a NumPy-specialized implementation to improve performance and preserve dtype.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Comment thread tests/test_numpy.py Outdated
Comment on lines +108 to +115
def clip(
x: Array,
/,
min: float | Array | None = None,
max: float | Array | None = None,
out: Array | None = None,
**kwargs,
) -> Array:

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

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

Removed out, if in kwargs pop it from there and use internally in function only

Comment thread src/array_api_compat/numpy/_aliases.py Outdated
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants