py-jaxlib: add spack-built ROCm support (#49611)

* py-jaxlib: add spack-built ROCm support

* fix style

* py-jaxlib 0.4.38 rocm support

* py-jaxlib 0.4.38 rocm support

* add comgr dependency

* changes for ROCm external and enable till 0.4.38

* enable version of py-jax

* add jax+rocm to ci

* add conflict for cuda and remove py-jaxlib from aarch64 pipeline

* Update var/spack/repos/builtin/packages/py-jaxlib/package.py

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

* add conflict for aarch64

---------

Co-authored-by: Adam J. Stewart <ajstewart426@gmail.com>

share/spack/gitlab/cloud_pipelines/stacks/ml-linux-x86_64-rocm/spack.yaml

@@ -27,9 +27,8 @@ spack:
     - py-transformers
 
     # JAX
-    # Does not yet support Spack-installed ROCm
-    # - py-jax
-    # - py-jaxlib
+    - py-jax
+    - py-jaxlib
 
     # Keras
     - py-keras backend=tensorflow