Add update_reference.yml workflow for manual HLO updates

Author: darisoy

.github/workflows/update_reference_hlo.yml

@@ -0,0 +1,32 @@
+name: Update HLO Reference
+
+on:
+  workflow_dispatch: # This creates the "Run workflow" button
+
+jobs:
+  update-hlo:
+    runs-on: ["linux-x86-ct6e-180-4tpu"] # Needs to be a TPU runner
+    container:
+      image: ubuntu:latest
+      options: -v /var/run/docker.sock:/var/run/docker.sock
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.ref }} # Checkout the PR branch
+
+      - name: Install Docker CLI
+        run: |
+          apt update && apt install docker.io -y
+
+      - name: Run update script
+        run: |
+          docker run -v ${{ github.workspace }}:/workspace -w /workspace maxtext-unit-test-tpu:py312 python3 tools/update_hlo_references.py
+
+      - name: Commit and Push changes
+        run: |
+          git config --global user.name "github-actions[bot]"
+          git config --global user.email "github-actions[bot]@users.noreply.github.com"
+          git add tests/integration/reference_hlo.txt
+          git commit -m "Update reference HLO from manual trigger"
+          git push

tests/integration/hlo_diff_test.py

@@ -19,7 +19,7 @@ def filter_line(line):
     return None
   # line = STACK_FRAME_REGEX.sub("stack_frame_id=DUMMY", line)
   # line = SHARDING_REGEX.sub("S(DUMMY)", line)
-  line = OP_NUM_REGEX.sub(".DUMMY", line)
+  # line = OP_NUM_REGEX.sub(".DUMMY", line)
   # Normalize 'square' to 'mul' to ignore naming differences (e.g. %square.157 vs %mul.1252)
   # line = line.replace("square", "mul")
   return line