jax

f8e26bdb - Merge pull request #28056 from apaszke:mgpu-memspace-picture