feat: add axis=None specialization for Min reducer + small refactoring for existing specializatons

[ikrommyd]

This adds axis=None fast path specialization for Min.
We also do some cleanup at the same time and remove the existing sum/max specializations for jax.
(I am sneaking in one more deprecation warning here for jax so that it's also raised when one does ak.jax.register_and_check()).

[codecov]

Codecov Report

❌ Patch coverage is 88.05970% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.65%. Comparing base (970f5d0) to head (8662964).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/awkward/_nplikes/typetracer.py	25.00%	6 Missing ⚠️
src/awkward/_connect/jax/reducers.py	93.10%	2 Missing ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
src/awkward/_do.py	84.70% <ø> (-0.18%)	⬇️
src/awkward/_kernels.py	70.00% <ø> (-0.84%)	⬇️
src/awkward/_nplikes/array_module.py	95.34% <100.00%> (+0.35%)	⬆️
src/awkward/_nplikes/numpy_like.py	100.00% <ø> (ø)
src/awkward/_reducers.py	98.31% <100.00%> (+0.09%)	⬆️
src/awkward/jax.py	85.24% <100.00%> (+0.50%)	⬆️
src/awkward/_connect/jax/reducers.py	92.28% <93.10%> (-0.64%)	⬇️
src/awkward/_nplikes/typetracer.py	77.17% <25.00%> (+1.51%)	⬆️

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

The documentation preview is ready to be viewed at http://preview.awkward-array.org.s3-website.us-east-1.amazonaws.com/PR3896

[ikrommyd]

@ianna if you could please check that the same speedup is there for max and that I didn't ruin anything, that'd be great. And that also there is a speedup for min too now.

[ianna]

@ikrommyd - please keep the JAX code in - we need it to run the test in preparation for the next week blueprint. As for the other changes - they are not random and have been tested on the previous reducer and did contribute to the performance gains. Thanks.

[ikrommyd]

I can keep the jax code in and also add the an axis none min reducer for jax too, but I don't understand how it has anything to do with the blueprint. jax backend is slow, unjittable, and deprecated. How does this performance optimization help? Reductions work for jax, they just do not use this optimization. Also with the previous code, they were not being used for jax anyways because they jax reducers do not implement the axis_none_reducer method.

[ianna]

I can keep the jax code in and also add the an axis none min reducer for jax too, but I don't understand how it has anything to do with the blueprint. jax backend is slow, unjittable, and deprecated. How does this performance optimization help? Reductions work for jax, they just do not use this optimization. Also with the previous code, they were not being used for jax anyways because they jax reducers do not implement the axis_none_reducer method.

Thanks! I will be taking about JAX and want to cross the tees — just to double check performance statements we made in the past.

[ianna]

I can keep the jax code in and also add the an axis none min reducer for jax too, but I don't understand how it has anything to do with the blueprint. jax backend is slow, unjittable, and deprecated. How does this performance optimization help? Reductions work for jax, they just do not use this optimization. Also with the previous code, they were not being used for jax anyways because they jax reducers do not implement the axis_none_reducer method.

Thanks! I will be taking about JAX and want to cross the tees — just to double check performance statements we made in the past.

Oh, missed that part about implementing extra reducers for JAX — I don’t think it’s necessary at this point. Thanks

[ikrommyd]

I'm not following perfectly. In the last PR, you attempted adding a fast-path axis none specialization for sum and max for jax? Do you want those in or not?

[ianna]

I'm not following perfectly. In the last PR, you attempted adding a fast-path axis none specialization for sum and max for jax? Do you want those in or not?

Yes, I want sum and max for JAX, but not min. Thanks.

[ikrommyd]

Then we'll have min too. Once the final release with the jax backend happens, it makes no sense for max to be much faster than min. I very highly want symmetry there and it's the same as max with changing max->min inside. All the jax changes are ready in this PR.

[ianna]

Then we'll have min too. Once the final release with the jax backend happens, it makes no sense for max to be much faster than min. I very highly want symmetry there and it's the same as max with changing max->min inside. All the jax changes are ready in this PR.

I completely agree that min and max should eventually be symmetrical for the final JAX release. The current asymmetry is actually a deliberate part of my 'max specialization' testing; I'm using min as the control group to measure the exact latency reduction and throughput improvement we get from this specialized path compared to the generic one.

[ikrommyd]

Then we'll have min too. Once the final release with the jax backend happens, it makes no sense for max to be much faster than min. I very highly want symmetry there and it's the same as max with changing max->min inside. All the jax changes are ready in this PR.

I completely agree that min and max should eventually be symmetrical for the final JAX release. The current asymmetry is actually a deliberate part of my 'max specialization' testing; I'm using min as the control group to measure the exact latency reduction and throughput improvement we get from this specialized path compared to the generic one.

Ah I see, okay then I would suggest you do this locally or on another remote branch by just deleting the AxisNoneMin that I'm adding here for jax. I think we should treat main branch as "production" in general. Testing like this is very easy to do locally by just deleting a few lines. Developers may often run on "main". I run my analysis on "main" for example. That's why I think it should be kept clean like production. Like anyone should be able to install from main and not expect "weird" behaviors at any time, that would be the goal.

[ikrommyd]

@ianna removed the tests as those code paths were being hit anyways. Is this good to go now?

[ianna]

@ikrommyd - great! Thanks!