Merge branch 'main' into aybchan/gke-update

Author: aybchan

.github/workflows/_ci.yaml

@@ -739,7 +739,7 @@ jobs:
         job-config-file:  ".github/eks-workflow-files/axlearn/axlearn-fuji-model.yml"
         job-name: ${{ env.JOB_NAME }}
 
-  test-alphafold-inference-eks:
+  test-alphafold-inference:
     needs: build-alphafold
     if: >-
       inputs.ARCHITECTURE == 'amd64' &&
@@ -757,17 +757,18 @@ jobs:
           NUM_GPUS: 1
           IMAGE: ${{ needs.build-alphafold.outputs.DOCKER_TAG_FINAL }}
           COMMAND: |
-            set -e
+            set -e;
 
-            python -c "from alphafold3.build_data import build_data; build_data()"
+            python -c 'from alphafold3.build_data import build_data; build_data()';
 
             python /usr/local/bin/generate_af_parameters.py \
               --input-parameters-file /opt/alphafold_input_data/model_parameters.md \
-              --out-model-dir /opt/alphafold/model
+              --out-model-dir /opt/alphafold/model;
 
             python /usr/local/bin/af3_inference_benchmark.py \
               --input-json-path /opt/alphafold_input_data/input_sequences/bench_L00512.json \
               --model-dir /opt/alphafold/model \
               --output-json /opt/output/output.json
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           NVCR_TOKEN: ${{ secrets.NVCR_TOKEN }}
+          STATISTICS: true

README.md

@@ -68,7 +68,6 @@ We support and test the following JAX frameworks and model architectures. More d
         <a href="https://gist.github.com/nvjax/913c2af68649fe568e9711c2dabb23ae/#file-final-jax-md"><img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-jax-build-arm64.json&logo=docker&label=arm64"></a>
       </td>
       <td>
-        <br>
         <a href="https://gist.github.com/nvjax/913c2af68649fe568e9711c2dabb23ae/#file-badge-jax-unit-backend-independent-h100-json">
           <img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-jax-unit-backend-independent-H100.json&logo=nvidia&label=JAX%20-%20backend%20independent%20(H100)">
         </a>
@@ -196,7 +195,9 @@ We support and test the following JAX frameworks and model architectures. More d
         </a>
       </td>
       <td>
-        [no tests]
+        <a href="https://gist.github.com/nvjax/913c2af68649fe568e9711c2dabb23ae/#file-badge-alphafold-inference-h100-json">
+          <img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-alphafold-inference-H100.json&logo=nvidia&label=AlphaFold%20inference%20(H100)">
+        </a>
       </td>
     </tr>
   </tbody>
@@ -273,12 +274,14 @@ For this initiative, we are publishing `scale-training` tagged containers - `jax
   </tbody>
 </table>
 
-The dedicated scale-training workflow runs every Saturday at 00:00 UTC and publishes containers every other Saturday, starting on April 18, 2026.
+The table belowed shows all the produced `scale-training` containers.
+From Saturday 23rd May, `scale-training` container will be generated and published every other Saturday at 00:00 UTC.
 Refer to this page [STAGING.md](https://github.com/openxla/xla/blob/nv-staging/latest/STAGING.md) for more information on the underlying XLA staging branch, the pending PRs included. The page also lists the corresponding JAX commit used.
 
 Staging releases:
 | Release date | Container | XLA branch (includes pending PRs) |
 | ------------------------------------- | -------------- |-------------- |
+| 2026-05-11 | [ghcr.io/nvidia/jax:jax-scale-training-2026-05-11](https://github.com/NVIDIA/JAX-Toolbox/pkgs/container/jax/856082965?tag=jax-scale-training-2026-05-11) | [5dfe2147](https://github.com/openxla/xla/blob/552b0a3ef06453c74de9f62f778a0f3a960d7e6d/STAGING.md)
 | 2026-04-24 | [ghcr.io/nvidia/jax:jax-scale-training-2026-04-24](https://github.com/NVIDIA/JAX-Toolbox/pkgs/container/jax/820220988?tag=jax-scale-training-2026-04-24) | [5dfe2147](https://github.com/openxla/xla/blob/5dfe2147cbdd54b2fa1d76da817c64a1847373ca/STAGING.md)
 | 2026-04-18 | [ghcr.io/nvidia/jax:jax-scale-training-2026-04-18](https://github.com/NVIDIA/JAX-Toolbox/pkgs/container/jax/805313885?tag=jax-scale-training-2026-04-18) | [8147118](https://github.com/sfvaroglu/xla/tree/8147118a7b9707d26dcb747767a9c0dd9081325f)