Add GPU CI #2

	name: ROCM CI

	on:
	# Trigger the workflow on push or pull request,
	# but only for the rocm-main branch
	push:
	branches:
	- rocm-main
	pull_request:
	branches:
	- rocm-main

	jobs:
	build-docker: # strategy and matrix come here
	runs-on: mi-250
	steps:
	- name: Print Environment Variables
	run: printenv
	- uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
	- name: Build Docker
	env:
	DOCKER_IMG_NAME: rocm_jax_r6_1_3_py3_10_id${{ github.run_id }}
	# XLA_CLONE_DIR:
	run: \|
	./build/rocm/ci_build.sh --rocm_version 6.1.3 \
	--keep_image --py_version 3.10
	- name: Archive jax wheels
	uses: actions/upload-artifact@v4
	with:
	name: rocm_jax_r6_1_3_py3_10_id${{ github.run_id }}
	path: ${{ github.workspace }}/wheelhouse/*.whl
	- name: Detect GPUs
	run: \|
	docker run \
	--group-add video \
	$RENDER_DEVICES \
	--cap-add=SYS_PTRACE \
	--security-opt seccomp=unconfined \
	--shm-size=64G \
	rocm/rocm-terminal rocm-smi
	- name: Run tests
	env:
	DOCKER_IMG_NAME: rocm_jax_r6_1_3_py3_10_id${{ github.run_id }}
	run: \|
	./build/rocm/ci_build test ${DOCKER_IMG_NAME}

Provide feedback