Merge branch 'main' into tessera_pass_draft

Author: jessicacotturone21

.github/workflows/build-reactantjll.yml

@@ -31,10 +31,30 @@ concurrency:
   cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 
 jobs:
+  get-reactant-commit:
+    name: 'Get Reactant commit'
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+
+    steps:
+      - name: Get latest Reactant commit
+        id: get-reactant-main
+        shell: bash
+        run: |
+          REACTANT_MAIN="$(git ls-remote https://github.com/EnzymeAD/Reactant.jl.git main | awk '{print $1}')"
+          echo "reactant_main=${REACTANT_MAIN}"
+          echo "reactant_main=${REACTANT_MAIN}" >> "${GITHUB_OUTPUT}"
+    outputs:
+      reactant_main: ${{ steps.get-reactant-main.outputs.reactant_main }}
+
   build-jll:
     name: Build Reactant_jll
     if: github.event.pull_request.draft == false
     uses: EnzymeAD/ReactantBuilder/.github/workflows/build-reactant-reusable.yml@main
+    needs: get-reactant-commit
     with:
       reactantbuilder_ref: 'main'
       enzyme_jax_commit: ${{ github.sha }}
+      # If you want to use a different commit than `main`, put here the sha256
+      # hash, but it must always be a hash, can't use branch name here.
+      reactant_commit: ${{ needs.get-reactant-commit.outputs.reactant_main }}

.github/workflows/pr-jax-main.yml

@@ -54,6 +54,7 @@ jobs:
           # - '0123456789abcdef0123456789abcdef01234567'
         reactant_commit:
           - 'main'
+          # - ap/symbol_ref_kernel_call
 
     steps:
       - name: Check GPUs

.github/workflows/test-gb-25.yml

@@ -0,0 +1,102 @@
+name: "Update upstream dependency"
+
+on:
+  workflow_call:
+    secrets:
+      ENZYMEAD_BOT_ID:
+        required: true
+      ENZYMEAD_BOT_PRIVATE_KEY:
+        required: true
+    inputs:
+      upstream_repo:
+        description: 'The GitHub repository of the upstream dependency'
+        required: true
+        type: 'string'
+      upstream_commit:
+        description: 'The commit of the upstream dependency to update to (optional)'
+        required: false
+        type: 'string'
+        default: ''
+      variable_name:
+        description: 'The name of the Bazel variable referencing the commit to update'
+        required: true
+        type: 'string'
+      base_branch:
+        description: 'Name of the base branch against which to open the PR (optional, defaults to main)'
+        required: false
+        type: 'string'
+        default: 'main'
+
+jobs:
+  pr-update-dep:
+    name: 'Update dependency'
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+
+    steps:
+      - uses: actions/create-github-app-token@v2
+        id: generate_token
+        with:
+          app-id: "${{ secrets.ENZYMEAD_BOT_ID }}"
+          private-key: "${{ secrets.ENZYMEAD_BOT_PRIVATE_KEY }}"
+      - uses: actions/checkout@v5
+      - name: Set branch name
+        run: |
+          echo "BRANCH_NAME=update-$(echo ${{ inputs.upstream_repo }} | tr '/A-Z' '-a-z')" >> "${GITHUB_ENV}"
+      - name: Prepare commit message
+        run: |
+          echo "COMMIT_MSG_PREFIX=Update ${{ inputs.upstream_repo }} to commit " >> "${GITHUB_ENV}"
+      - name: Get upstream commit
+        shell: bash
+        run: |
+          OLD_UPSTREAM_COMMIT=$(grep '${{ inputs.variable_name }} *= *".*"' workspace.bzl | cut -d= -f2 | tr -d '" ')
+
+          if [[ -n '${{ inputs.upstream_commit }}' ]]; then
+              NEW_UPSTREAM_COMMIT="${{ inputs.upstream_commit }}"
+          else
+              NEW_UPSTREAM_COMMIT="$(git ls-remote https://github.com/${{ inputs.upstream_repo }}.git main | awk '{print $1}')"
+          fi
+
+          echo "OLD_UPSTREAM_COMMIT=${OLD_UPSTREAM_COMMIT}" >> "${GITHUB_ENV}"
+          echo "NEW_UPSTREAM_COMMIT=${NEW_UPSTREAM_COMMIT}" >> "${GITHUB_ENV}"
+          echo "DIFF_MSG=Diff: https://github.com/${{ inputs.upstream_repo }}/compare/${OLD_UPSTREAM_COMMIT}...${NEW_UPSTREAM_COMMIT}" >> "${GITHUB_ENV}"
+      - name: Modify upstream commit
+        run: |
+          sed -i 's/${{ inputs.variable_name }} *= *".*"/${{ inputs.variable_name }} = "${{ env.NEW_UPSTREAM_COMMIT }}"/' workspace.bzl
+      - name: Decide whether to open the pull request
+        id: decide
+        run: |
+          # Fetch latest commit on the branch
+          LATEST_COMMIT_MESSAGE="$(curl -s -H 'Authorization: token ${{ secrets.GITHUB_TOKEN }}' 'https://api.github.com/repos/${{ github.repository }}/commits/${{ env.BRANCH_NAME }}' | jq -r '.commit.message')"
+          echo "Latest commit message on the branch '${{ env.BRANCH_NAME }}' is:"
+          echo "----------"
+          echo "${LATEST_COMMIT_MESSAGE}"
+          echo "----------"
+          if [[ "${LATEST_COMMIT_MESSAGE}" == "null" || "${LATEST_COMMIT_MESSAGE}" == "${{ env.COMMIT_MSG_PREFIX }}"* ]]; then
+              # Open the PR only if the target branch doesn't exist or latest commit message
+              # starts with the expected prefix.
+              echo "We're going to open the pull request..."
+              echo "should_pr=true" >> "${GITHUB_OUTPUT}"
+          else
+              # Otherwise there are other commits to the branch and we don't want to
+              # override them!
+              echo "There seem to be extra commits, we won't open the pull request"
+              echo "should_pr=false" >> "${GITHUB_OUTPUT}"
+          fi
+      - name: Create Pull Request
+        if: ${{ steps.decide.outputs.should_pr == 'true' }}
+        uses: peter-evans/create-pull-request@v7
+        with:
+          token: ${{ steps.generate_token.outputs.token }}
+          commit-message: |
+            ${{ env.COMMIT_MSG_PREFIX }}${{ env.NEW_UPSTREAM_COMMIT }}
+
+            ${{ env.DIFF_MSG }}
+          body: ${{ env.DIFF_MSG }}
+          base: ${{ inputs.base_branch }}
+          branch: ${{ env.BRANCH_NAME }}
+          title: '${{ env.COMMIT_MSG_PREFIX }}${{ env.NEW_UPSTREAM_COMMIT }}'
+          delete-branch: true
+          draft: false
+          author: 'enzymead-bot[bot] <238314553+enzymead-bot[bot]@users.noreply.github.com>'
+          labels: 'dependencies'

.github/workflows/update-dependency.yml

@@ -0,0 +1,27 @@
+name: "Open PR to update Enzyme commit"
+
+on:
+  schedule:
+    - cron: '19 10 * * *'
+  workflow_dispatch:
+    inputs:
+      enzyme_commit:
+        description: 'The Enzyme commit to update to (optional)'
+        default: ''
+        type: 'string'
+
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+jobs:
+  pr-latest-enzyme:
+    name: "Update Enzyme"
+    uses: ./.github/workflows/update-dependency.yml
+    with:
+      upstream_repo: "EnzymeAD/Enzyme"
+      upstream_commit: ${{ inputs.enzyme_commit }}
+      variable_name: "ENZYME_COMMIT"
+    secrets: inherit

.github/workflows/update-enzyme.yml

@@ -0,0 +1,27 @@
+name: "Open PR to update JAX commit"
+
+on:
+  schedule:
+    - cron: '19 22 * * *'
+  workflow_dispatch:
+    inputs:
+      jax_commit:
+        description: 'The JAX commit to update to (optional)'
+        default: ''
+        type: 'string'
+
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
+
+jobs:
+  pr-latest-jax:
+    name: "Update JAX"
+    uses: ./.github/workflows/update-dependency.yml
+    with:
+      upstream_repo: "jax-ml/jax"
+      upstream_commit: ${{ inputs.jax_commit }}
+      variable_name: "JAX_COMMIT"
+    secrets: inherit

.github/workflows/update-jax.yml

@@ -33,6 +33,10 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 #     ],
 # )
 
+load("//third_party/ml_toolchain:workspace.bzl", ml_toolchain_workspace = "repo")
+
+ml_toolchain_workspace()
+
 load("//third_party/jax:workspace.bzl", jax_workspace = "repo")
 
 jax_workspace()