use ml build image

Author: nitins17

.github/workflows/asan.yaml

@@ -14,14 +14,16 @@ jobs:
   asan:
     runs-on: linux-x86-n2-64
     container:
-      image: index.docker.io/library/ubuntu@sha256:b359f1067efa76f37863778f7b6d0e8d911e3ee8efa807ad01fbf5dc1ef9006b # ratchet:ubuntu:24.04
+      image: us-central1-docker.pkg.dev/tensorflow-sigs/tensorflow/ml-build:latest # ratchet:ubuntu:24.04
     strategy:
       fail-fast: false
     defaults:
       run:
         shell: bash -l {0}
     steps:
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
+        with:
+          path: jax
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
         with:
           repository: python/cpython