Skip to content

feat: remove jax from main#3598

Closed
ianna wants to merge 9 commits intoscikit-hep:mainfrom
ianna:ianna/cleanup_jax
Closed

feat: remove jax from main#3598
ianna wants to merge 9 commits intoscikit-hep:mainfrom
ianna:ianna/cleanup_jax

Conversation

@ianna
Copy link
Copy Markdown
Member

@ianna ianna commented Aug 3, 2025

use awkward-jax branch for further development

@ianna ianna requested a review from pfackeldey August 3, 2025 14:00
Copy link
Copy Markdown
Collaborator

@pfackeldey pfackeldey left a comment

Choose a reason for hiding this comment

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

Thanks!

-8k lines of code 🎉

This and the fact that Jax backend was cumbersome will definitely help to simplify development in the future.

I've added my comments to the review. One thing that I'm missing so far is removal of JAX from the docs, and pointing people to the last awkward release in the docs that contains the JAX backend, and/or the awkward-jax branch.



@high_level_function()
def from_jax(
Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

I don't think we should get rid of this function. We still want to support converting a jax array to an awkward array with e.g. cpu backend, we're just getting rid of jax internally in awkward.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Thanks, I agree. I was debating if it is useful to keep these high level functions.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

restored!

Copy link
Copy Markdown
Collaborator

@pfackeldey pfackeldey Aug 3, 2025

Choose a reason for hiding this comment

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

This and the ak.to_jax implementation do not need the jax backend/nplike. We can restore their behavior without reintroducing all of that code. (we should for sure get rid of the jax backend & nplike)

@ianna ianna force-pushed the ianna/cleanup_jax branch from cd5755c to f29d29b Compare August 3, 2025 15:14
@codecov
Copy link
Copy Markdown

codecov bot commented Aug 8, 2025

Codecov Report

❌ Patch coverage is 95.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.78%. Comparing base (b749e49) to head (1607c74).
⚠️ Report is 394 commits behind head on main.

Files with missing lines Patch % Lines
src/awkward/_slicing.py 75.00% 2 Missing ⚠️
Additional details and impacted files
Files with missing lines Coverage Δ
src/awkward/_backends/dispatch.py 87.80% <ø> (-5.06%) ⬇️
src/awkward/_backends/jax.py 92.30% <ø> (-7.70%) ⬇️
src/awkward/_broadcasting.py 92.41% <100.00%> (-2.94%) ⬇️
src/awkward/_nplikes/__init__.py 67.85% <ø> (-13.97%) ⬇️
src/awkward/_nplikes/jax.py 56.70% <ø> (-25.91%) ⬇️
src/awkward/_nplikes/virtual.py 91.01% <100.00%> (ø)
src/awkward/contents/indexedoptionarray.py 88.43% <100.00%> (+0.21%) ⬆️
src/awkward/contents/listarray.py 89.88% <100.00%> (+0.45%) ⬆️
src/awkward/contents/unionarray.py 85.89% <100.00%> (+0.63%) ⬆️
src/awkward/highlevel.py 78.51% <ø> (+1.84%) ⬆️
... and 14 more

... and 175 files with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

@ianna
Copy link
Copy Markdown
Member Author

ianna commented Aug 11, 2025

@pfackeldey – Let’s postpone this until after the +1 release (around the end of October). I’ll close my PR for now.

@ianna ianna closed this Aug 11, 2025
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