diff --git a/.github/workflows/_test_upstream_maxtext.yaml b/.github/workflows/_test_upstream_maxtext.yaml
index 61bee91d6..77589d479 100644
--- a/.github/workflows/_test_upstream_maxtext.yaml
+++ b/.github/workflows/_test_upstream_maxtext.yaml
@@ -11,13 +11,13 @@ on:
       EXTRA_TEST_ARGS:
         type: string
         description: Extra command line args to pass to test-maxtext.sh
-        default: ""
+        default: ''
         required: false
       BADGE_FILENAME:
         type: string
         description: 'Name of the endpoint JSON file for shields.io badge'
         required: false
-        default: 'badge-maxtext-test.json'
+        default: 'badge-upstream-maxtext-test.json'
       ARTIFACT_NAME:
         type: string
         description: 'Name of the artifact zip file'
@@ -34,12 +34,11 @@ on:
         value: ${{ jobs.sitrep.outputs.STATUS }}
 
 jobs:
-
   single-process-multi-device:
     strategy:
       matrix:
         PARALLEL_CONFIG:
-        - [1, 1, 2, 4]
+          - [1, 1, 2, 4]
         # - [1, 1, 1, 8] # PP, DP, FSDP, TP
       fail-fast: false
 
@@ -183,12 +182,12 @@ jobs:
     strategy:
       matrix:
         PARALLEL_CONFIG:
-        - [1, 1, 1, 1]
-        - [1, 1, 8, 1]
-        - [1, 1, 1, 8]
-        - [1, 1, 4, 2]
-        - [1, 2, 2, 2]
-        - [1, 4, 2, 2]
+          - [1, 1, 1, 1]
+          - [1, 1, 8, 1]
+          - [1, 1, 1, 8]
+          - [1, 1, 4, 2]
+          - [1, 2, 2, 2]
+          - [1, 4, 2, 2]
       fail-fast: false
 
     runs-on: ubuntu-22.04
@@ -366,7 +365,7 @@ jobs:
 
   sitrep:
     needs: [single-process-multi-device, maxtext-multinode, metrics]
-    if: "!cancelled()"
+    if: '!cancelled()'
     uses: ./.github/workflows/_sitrep_mgmn.yaml
     secrets: inherit
     with:
@@ -377,7 +376,7 @@ jobs:
   summary:
     runs-on: ubuntu-22.04
     needs: [single-process-multi-device, maxtext-multinode]
-    if: "!cancelled()"
+    if: '!cancelled()'
     steps:
       - name: Generate TensorBoard query URL
         run: |
@@ -394,7 +393,7 @@ jobs:
   outcome:
     needs: sitrep
     runs-on: ubuntu-22.04
-    if: "!cancelled()"
+    if: '!cancelled()'
     steps:
       - name: Sets workflow status based on test outputs
         run: |
diff --git a/README.md b/README.md
index fb4ae772d..fe7a960ca 100644
--- a/README.md
+++ b/README.md
@@ -222,7 +222,7 @@
         <code>ghcr.io/nvidia/jax:upstream-maxtext</code>
       </td>
       <td>
-        <a href="https://gist.github.com/nvjax/913c2af68649fe568e9711c2dabb23ae/#file-final-maxtext-md"><img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-upstream-maxtext-build-amd64.json&logo=docker&label=amd64"></a>
+        <a href="https://gist.github.com/nvjax/913c2af68649fe568e9711c2dabb23ae/#file-final-upstream-maxtext-md"><img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-upstream-maxtext-build-amd64.json&logo=docker&label=amd64"></a>
          <img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-upstream-maxtext-build-arm64.json&logo=docker&label=arm64">
       </td>
       <td>
@@ -241,8 +241,8 @@
         <code>ghcr.io/nvidia/jax:maxtext</code>
       </td>
       <td>
-        <a href="https://gist.github.com/nvjax/913c2af68649fe568e9711c2dabb23ae/#file-final-maxtext-md"><img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-rosetta-maxtext-build-amd64.json&logo=docker&label=amd64"></a>
-         <img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-rosetta-maxtext-build-arm64.json&logo=docker&label=arm64">
+        <a href="https://gist.github.com/nvjax/913c2af68649fe568e9711c2dabb23ae/#file-final-md"><img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-rosetta-build-amd64.json&logo=docker&label=amd64"></a>
+         <img style="height:1em;" src="https://img.shields.io/endpoint?url=https%3A%2F%2Fgist.githubusercontent.com%2Fnvjax%2F913c2af68649fe568e9711c2dabb23ae%2Fraw%2Fbadge-rosetta-build-arm64.json&logo=docker&label=arm64">
       </td>
       <td>
         <picture>