jax

5cad02c1 - [pjrt] Use the `PjRtMemorySpace` version of `BufferFromHostLiteral`

[pjrt] Use the `PjRtMemorySpace` version of `BufferFromHostLiteral` Buffers live in memory spaces and not on devices. The `PjRtDevice` version of `BufferFromHostLiteral` is deprecated and will be removed once the migration is complete. PiperOrigin-RevId: 721860983