Change for `DeviceLocalLayout` in JAX 0.6.3.#126

hertschuh merged 1 commit intokeras-team:mainkeras-team/keras-rs:mainfrom

hertschuh:jax_layouthertschuh/keras-rs:jax_layoutCopy head branch name to clipboard