merge branch 'master' into n-arity-v4

Author: MilesCranmer

.github/workflows/benchmark_pr.yml

@@ -1,64 +1,14 @@
-name: Benchmark a pull request
-
+name: Benchmark this PR
 on:
   pull_request_target:
-    branches:
-      - master
-
+    branches: [master]
 permissions:
   pull-requests: write
 
 jobs:
-    generate_plots:
-        runs-on: ubuntu-latest
-
-        steps:
-            - uses: actions/checkout@v2
-            - uses: julia-actions/setup-julia@v1
-              with:
-                version: "1"
-            - uses: julia-actions/cache@v1
-            - name: Extract Package Name from Project.toml
-              id: extract-package-name
-              run: |
-                PACKAGE_NAME=$(grep "^name" Project.toml | sed 's/^name = "\(.*\)"$/\1/')
-                echo "::set-output name=package_name::$PACKAGE_NAME"
-            - name: Build AirspeedVelocity
-              env:
-                JULIA_NUM_THREADS: 2
-              run: |
-                # Lightweight build step, as sometimes the runner runs out of memory:
-                julia -e 'ENV["JULIA_PKG_PRECOMPILE_AUTO"]=0; import Pkg; Pkg.add(;url="https://github.com/MilesCranmer/AirspeedVelocity.jl.git")'
-                julia -e 'ENV["JULIA_PKG_PRECOMPILE_AUTO"]=0; import Pkg; Pkg.build("AirspeedVelocity")'
-            - name: Add ~/.julia/bin to PATH
-              run: |
-                echo "$HOME/.julia/bin" >> $GITHUB_PATH
-            - name: Run benchmarks
-              run: |
-                echo $PATH
-                ls -l ~/.julia/bin
-                mkdir results
-                benchpkg ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --url=${{ github.event.repository.clone_url }} --bench-on="${{github.event.pull_request.head.sha}}" --output-dir=results/ --tune --exeflags="-O3 --threads=auto"
-            - name: Create markdown table from benchmarks
-              run: |
-                benchpkgtable ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --input-dir=results/ --ratio > table.md
-                echo '### Benchmark Results' > body.md
-                echo '' >> body.md
-                echo '' >> body.md
-                cat table.md >> body.md
-                echo '' >> body.md
-            - name: Find Comment
-              uses: peter-evans/find-comment@v2
-              id: fcbenchmark
-              with:
-                issue-number: ${{ github.event.pull_request.number }}
-                comment-author: 'github-actions[bot]'
-                body-includes: Benchmark Results
-
-            - name: Comment on PR
-              uses: peter-evans/create-or-update-comment@v3
-              with:
-                comment-id: ${{ steps.fcbenchmark.outputs.comment-id }}
-                issue-number: ${{ github.event.pull_request.number }}
-                body-path: body.md
-                edit-mode: replace
+  bench:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: MilesCranmer/AirspeedVelocity.jl@action-v1
+        with:
+          julia-version: '1'

Project.toml

@@ -38,7 +38,7 @@ Optim = "0.19, 1"
 PrecompileTools = "1"
 Reexport = "1"
 SymbolicUtils = "0.19, ^1.0.5, 2, 3"
-Zygote = "0.6"
+Zygote = "0.7"
 julia = "1.10"
 
 [extras]

src/ChainRules.jl

@@ -7,6 +7,7 @@ using ChainRulesCore:
     ZeroTangent,
     Tangent,
     @thunk,
+    unthunk,
     canonicalize
 using ..OperatorEnumModule: OperatorEnum
 using ..NodeModule: AbstractExpressionNode, with_type_parameters, tree_mapreduce
@@ -52,7 +53,8 @@ struct EvalPullback{N,A,O} <: Function
 end
 
 # TODO: Preferable to use the primal in the pullback somehow
-function (e::EvalPullback)((dY, _))
+function (e::EvalPullback)((thunked_dY, _))
+    dY = unthunk(thunked_dY)
     _, dX_constants_dY, complete = eval_grad_tree_array(
         e.tree, e.X, e.operators; variable=Val(:both)
     )