fix: use jax's __setitem__ syntax in ak.merge_union_of_records to make it work with the jax backend

[ikrommyd]

ak.merge_union_of_records doesn't currently handle jax's setitem syntax. I'm adding that

[codecov]

Codecov Report

❌ Patch coverage is 63.63636% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.72%. Comparing base (b749e49) to head (90f6670).
⚠️ Report is 452 commits behind head on main.

Files with missing lines	Patch %	Lines
...rc/awkward/operations/ak_merge_union_of_records.py	63.63%	4 Missing ⚠️

Additional details and impacted files

Files with missing lines	Coverage Δ
...rc/awkward/operations/ak_merge_union_of_records.py	82.11% <63.63%> (-3.48%)	⬇️

... and 199 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]

@ikrommyd - thanks for looking into it. As we have discussed during the last meeting, we are refactoring the JAX, VirualArrays, PlaceholerArrays, and TypeTracerArrays code. IMHO, it would be better to have it all in one place. Please, add this to the list of issues to be fixed. Thanks!

[ikrommyd]

Well refactoring anything doesn’t touch this. This is just an operation that somewhat accidentally works with cpu but doesn’t work with anything else so I just fixed that. It’s not a super common operation though so please take over the PR as you please. It’s not critical.

[ikrommyd]

Given #3598, I'm chaging this to target the jax branch and setting it to draft mode to be merged later.

[ianna]

Given #3598, I'm chaging this to target the jax branch and setting it to draft mode to be merged later.

Thanks, as discussed, we are removing both Jax nplike and Jax backend. Please, move this to your fork if your development is based on it. Thanks.

[ikrommyd]

It already lives on my fork. It now targets the awkward-jax branch and not main which I believe makes this a valid PR if jax backend is only to live there after the removal.

[ianna]

It already lives on my fork. It now targets the awkward-jax branch and not main which I believe makes this a valid PR if jax backend is only to live there after the removal.

AFAIK we have no plans to maintain the awkward-jax branch. You can merge it to your fork.

[ianna]

It already lives on my fork. It now targets the awkward-jax branch and not main which I believe makes this a valid PR if jax backend is only to live there after the removal.

It should not be on this list. Please move it to your fork. I’m closing here. Thanks!

[github-actions]

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

[ikrommyd]

@ianna this is the PR I meant today. Just fixes these couple of setitem calls.

[ianna]

@ikrommyd - great! Thank you.

[ianna]

@pfackeldey - ping :-)

[pfackeldey]

LGTM! 👍