jax

cd019717 - [pallas:mgpu] Fix handling of `AbstractRef` in `_commute_transform`.