Expose memory space in general to JAX users

[yliu120]

Prototype Reference: https://github.com/openxla/xla/pull/23149/files

This is the internal annotation custom calls we might use to actually annotate a buffer to a specific memory space. However, we actually think the memory space annotations are too heavy and not necessary.

To get great performance and allow maximal flexibility, the users should be able to pin specific buffers onto some specific memory spaces. Let's consider the following case, a JAX user want to get the buffer as a persistent temp buffer or a pinned host buffer so that they can use them inside their FFI custom ops.

We should in general do the following,

1. Fixes SPMD partitioning to recognize heterogenous memory spaces.
2. Don't drop user-set memory space in layout assignment
3. Fixes layout normalization
4. Adds a bit memory space propagation.

[nouiz]

Can you confirm that by "persistent temp buffer" you want an activation to always have the same buffer?
Why that is useful?

JAX started to support supporting activation offloading to the host memory via the jax.remat API. Instead of recomputing, it offload. Is that what you want to do here?

Here is some recent documentation about memory space for input and outputs of a JAX function:
https://docs.jax.dev/en/latest/sharded-computation.html#sharding-transformation-between-memory-types

Here is some doc to do activation offloading instead of recomputation: https://docs.jax.dev/en/latest/gradient-checkpointing.html#custom-policies-for-offload