Add update_reference.yml workflow for manual HLO updates

Author: darisoy

.github/workflows/update_reference_hlo.yml

@@ -0,0 +1,29 @@
+name: Update HLO Reference
+
+on:
+  pull_request:
+    branches:
+      - main
+
+jobs:
+  update-hlo:
+    runs-on: ["linux-x86-ct6e-180-4tpu"] # Needs to be a TPU runner
+    container:
+      image: maxtext-unit-test-tpu:py312 # Same image as CI tests
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+        with:
+          ref: ${{ github.ref }} # Checkout the PR branch
+
+      - name: Run update script
+        run: |
+          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 PR 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