PR #11583: Weight offloading of Jax memories: support memory_kind for GPUs
PR #11583: Weight offloading of Jax memories: support memory_kind for GPUs

Imported from GitHub PR https://github.com/openxla/xla/pull/11583

The first step to implement weight offloading is to support the existing host memory kind for GPUs: pinned_host.  

Below are examples of the supported Jax APIs:
```
NamedSharding(mesh, P('x', 'y'), memory_kind='pinned_host')
s_host.with_memory_kind('pinned_host')
```
The next step will be supporting Jax API of device_put() via PjRtStreamExecutorClient::BufferFromHostBuffer() in XLA. Each tensor can only have one memory space or memory kind. This will be implemented in PjRtBuffer to get the mapped memory space from NamedSharding() or with_memory_kind().
Copybara import of the project:

--
7ac05a767e23755f8505bef9d1f077d1a293ab0c by Jane Liu <janeliu@nvidia.com>:

Support NamedSharding with memory_kind by adding host memory spaces.

Merging this change closes #11583

FUTURE_COPYBARA_INTEGRATE_REVIEW=https://github.com/openxla/xla/pull/11583 from zhenying-liu:weight-offloading 7ac05a767e23755f8505bef9d1f077d1a293ab0c
