Fix FFI example test in CI

Author: charleshofer

.github/workflows/ci-build.yaml

@@ -216,9 +216,7 @@ jobs:
 
   ffi:
     name: FFI example
-    runs-on: linux-x86-g2-16-l4-1gpu
-    container:
-      image: index.docker.io/tensorflow/build:latest-python3.12@sha256:48e99608fe9434ada5b14e19fdfd8e64f4cfc83aacd328b9c2101b210e984295 # ratchet:index.docker.io/tensorflow/build:latest-python3.12
+    runs-on: ROCM-Ubuntu
     timeout-minutes: 30
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
@@ -237,7 +235,7 @@ jobs:
         path: ${{ steps.pip-cache.outputs.dir }}
         key: ${{ runner.os }}-pip-ffi-examples-${{ hashFiles('**/setup.py', '**/requirements.txt', '**/test-requirements.txt', 'examples/**/pyproject.toml') }}
     - name: Install JAX
-      run: pip install .[cuda12]
+      run: pip install .
     - name: Build and install example project
       run: python -m pip install -v ./examples/ffi[test]
       env:
@@ -246,7 +244,7 @@ jobs:
         # a different toolchain. GCC is the default compiler on the
         # 'ubuntu-latest' runner, but we still set this explicitly just to be
         # clear.
-        CMAKE_ARGS: -DCMAKE_CXX_COMPILER=g++ -DJAX_FFI_EXAMPLE_ENABLE_CUDA=ON
+        CMAKE_ARGS: -DCMAKE_CXX_COMPILER=g++ #-DJAX_FFI_EXAMPLE_ENABLE_CUDA=ON
     - name: Run CPU tests
       run: python -m pytest examples/ffi/tests
       env: