jax

c028eb9f - [Pallas:MGPU] Fix `plgpu.inline_mgpu` support for scalar inputs under WG semantics.