jax

9a0eddda - PR #35550: [ROCm] Add ROCm wheel build and test pipeline to continuous CI

PR #35550: [ROCm] Add ROCm wheel build and test pipeline to continuous CI Imported from GitHub PR https://github.com/jax-ml/jax/pull/35550 - Add jax-rocm-plugin and jax-rocm-pjrt to allowed artifacts in build_artifacts.sh with ROCm version flag passthrough. - Create build_rocm_artifacts.yml reusable workflow that builds ROCm wheels in a ROCm container and uploads them to S3 via OIDC. - Extend wheel_tests_continuous.yml with build-rocm-artifacts, run-pytest-rocm, and run-bazel-test-rocm jobs. <!-- Thanks for taking the time to contribute to JAX! A couple things to keep in mind: * Contributing to JAX requires signing the Contributor License Agreement: https://docs.jax.dev/en/latest/contributing.html#google-contributor-license-agreement * Please run lint checks and tests locally: https://docs.jax.dev/en/latest/contributing.html#contributing-code-using-pull-requests * If applicable, read our policy on AI generated code: https://docs.jax.dev/en/latest/contributing.html#can-i-contribute-ai-generated-code --> Copybara import of the project: -- 43c3c5daa66c389059e6f81be7e6a5dcebcd3b83 by Alexandros Theodoridis <atheodor@amd.com>: Add ROCm wheel build and test pipeline to continuous CI - Add jax-rocm-plugin and jax-rocm-pjrt to allowed artifacts in build_artifacts.sh with ROCm version flag passthrough. - Create build_rocm_artifacts.yml reusable workflow that builds ROCm wheels in a ROCm container and uploads them to S3 via OIDC. - Extend wheel_tests_continuous.yml with build-rocm-artifacts, run-pytest-rocm, and run-bazel-test-rocm jobs. -- e4e0d96c031fa4317747eb6f3caa9887bbaf4590 by Alexandros Theodoridis <atheodor@amd.com>: Trigger CI/CD pipeline Merging this change closes #35550 COPYBARA_INTEGRATE_REVIEW=https://github.com/jax-ml/jax/pull/35550 from ROCm:implement_periodic_build_wheels_pipeline_rocm_v2 e4e0d96c031fa4317747eb6f3caa9887bbaf4590 PiperOrigin-RevId: 879796713