[rocm] is jaxlib meant to be cpu specific?

[Ruturaj4]

we observe that ->

pip list | grep jax
jax                          0.4.30
jax-cuda12-pjrt              0.4.30
jax-cuda12-plugin            0.4.30
jaxlib                       0.4.30

with pjrt support and --enable_cuda --build_gpu_plugin --gpu_plugin_cuda_version=12, we get three wheels. isn't jaxlib suuposed
to be cpu specific? if so, why are there gpu components in that.

Length Date Time Name

---

  629  2024-07-11 12:58   jaxlib/__init__.py
 8056  2024-07-11 12:58   jaxlib/cpu_feature_guard.so
  905  2024-07-11 12:58   jaxlib/gpu_common_utils.py
 3771  2024-07-11 12:58   jaxlib/gpu_linalg.py
 4760  2024-07-11 12:58   jaxlib/gpu_prng.py
 5576  2024-07-11 12:58   jaxlib/gpu_rnn.py
19445  2024-07-11 12:58   jaxlib/gpu_solver.py
13417  2024-07-11 12:58   jaxlib/gpu_sparse.py
 2307  2024-07-11 12:58   jaxlib/gpu_triton.py
10310  2024-07-11 12:58   jaxlib/hlo_helpers.py
24819  2024-07-11 12:58   jaxlib/lapack.py
    0  2024-07-11 12:58   jaxlib/py.typed

159968 2024-07-11 12:58 jaxlib/utils.so
....

[hawkinsp]

Yes, there are still a few python-only GPU components in jaxlib, mostly things like lowering rules (as here). But it's the same jaxlib used on all hardware platforms.

We might consider moving them into the plugins at some point, although I don't know of a pressing reason to do it.