update linux arm64 container to ml-build

Author: nitins17

.github/workflows/build_artifacts.yml

@@ -1,9 +1,9 @@
 name: Build JAX Artifacts
 
 on:
-  # pull_request:
-  #   branches:
-  #     - main
+  pull_request:
+    branches:
+      - main
   workflow_dispatch:
     inputs:
       halt-for-connection:
@@ -54,7 +54,7 @@ jobs:
     runs-on: ${{ matrix.runner }}
 
     container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
-                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') ||
+                   (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') ||
                    (contains(matrix.runner, 'windows-x86') && null) }}
 
     env:

.github/workflows/pytest_cpu.yml

@@ -29,7 +29,7 @@ jobs:
 
     runs-on: ${{ matrix.runner }}
     container: ${{ (contains(matrix.runner, 'linux-x86') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest') ||
-                (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest') ||
+                (contains(matrix.runner, 'linux-arm64') && 'us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest') ||
                 (contains(matrix.runner, 'windows-x86') && null) }}
 
     env:

ci/envs/docker.env

@@ -29,7 +29,7 @@ fi
 
 # Linux Aarch64 specifc settings
 if [[ $os == "linux" ]] && [[ $arch == "aarch64" ]]; then
-  export JAXCI_DOCKER_IMAGE="us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/linux-arm64-arc-container:latest"
+  export JAXCI_DOCKER_IMAGE="us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build-arm64:latest"
 fi
 
 # Windows specific settings