diff --git a/.JuliaFormatter.toml b/.JuliaFormatter.toml deleted file mode 100644 index 3494a9f..0000000 --- a/.JuliaFormatter.toml +++ /dev/null @@ -1,3 +0,0 @@ -style = "sciml" -format_markdown = true -format_docstrings = true diff --git a/.github/dependabot.yml b/.github/dependabot.yml index ec3b005..f901dd8 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -5,6 +5,14 @@ updates: directory: "/" # Location of package manifests schedule: interval: "weekly" - ignore: - - dependency-name: "crate-ci/typos" - update-types: ["version-update:semver-patch", "version-update:semver-minor"] + - package-ecosystem: "julia" + directories: + - "/" + - "/docs" + - "/test" + schedule: + interval: "daily" + groups: + all-julia-packages: + patterns: + - "*" diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml deleted file mode 100644 index cba9134..0000000 --- a/.github/workflows/CompatHelper.yml +++ /dev/null @@ -1,16 +0,0 @@ -name: CompatHelper -on: - schedule: - - cron: 0 0 * * * - workflow_dispatch: -jobs: - CompatHelper: - runs-on: ubuntu-latest - steps: - - name: Pkg.add("CompatHelper") - run: julia -e 'using Pkg; Pkg.add("CompatHelper")' - - name: CompatHelper.main() - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - COMPATHELPER_PRIV: ${{ secrets.DOCUMENTER_KEY }} - run: julia -e 'using CompatHelper; CompatHelper.main()' diff --git a/.github/workflows/DependabotAutoMerge.yml b/.github/workflows/DependabotAutoMerge.yml new file mode 100644 index 0000000..44b814c --- /dev/null +++ b/.github/workflows/DependabotAutoMerge.yml @@ -0,0 +1,9 @@ +name: "Dependabot Auto-merge" +on: pull_request +permissions: + contents: write + pull-requests: write +jobs: + automerge: + uses: "SciML/.github/.github/workflows/dependabot-automerge.yml@v1" + secrets: "inherit" diff --git a/.github/workflows/DocPreviewCleanup.yml b/.github/workflows/DocPreviewCleanup.yml new file mode 100644 index 0000000..960a549 --- /dev/null +++ b/.github/workflows/DocPreviewCleanup.yml @@ -0,0 +1,10 @@ +name: "Doc Preview Cleanup" +on: + pull_request: + types: [closed] +permissions: + contents: write +jobs: + cleanup: + uses: "SciML/.github/.github/workflows/docs-preview-cleanup.yml@v1" + secrets: "inherit" diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml new file mode 100644 index 0000000..3755851 --- /dev/null +++ b/.github/workflows/Documentation.yml @@ -0,0 +1,21 @@ +name: "Documentation" + +on: + pull_request: + branches: + - main + - 'release-' + push: + branches: + - main + tags: '*' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref_name != github.event.repository.default_branch || github.ref != 'refs/tags/v*' }} + +jobs: + docs: + name: Documentation + uses: "SciML/.github/.github/workflows/documentation.yml@v1" + secrets: "inherit" diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml new file mode 100644 index 0000000..ee944a5 --- /dev/null +++ b/.github/workflows/Downgrade.yml @@ -0,0 +1,32 @@ +name: Downgrade + +on: + pull_request: + branches: + - main + - 'release-' + paths-ignore: + - 'docs/**' + push: + branches: + - main + paths-ignore: + - 'docs/**' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref_name != github.event.repository.default_branch || github.ref != 'refs/tags/v*' }} + +jobs: + downgrade: + name: Downgrade + strategy: + fail-fast: false + matrix: + group: + - Core + uses: "SciML/.github/.github/workflows/downgrade.yml@v1" + with: + group: "${{ matrix.group }}" + julia-version: "lts" + secrets: "inherit" diff --git a/.github/workflows/FormatCheck.yml b/.github/workflows/FormatCheck.yml index fe9c128..db91458 100644 --- a/.github/workflows/FormatCheck.yml +++ b/.github/workflows/FormatCheck.yml @@ -1,13 +1,16 @@ -name: "Format Check" +name: format-check on: push: branches: + - 'master' - 'main' + - 'release-' tags: '*' pull_request: jobs: - format-check: - name: "Format Check" - uses: "SciML/.github/.github/workflows/format-suggestions-on-pr.yml@v1" + runic: + name: Runic + uses: "SciML/.github/.github/workflows/runic.yml@v1" + secrets: "inherit" diff --git a/.github/workflows/RunicSuggestions.yml b/.github/workflows/RunicSuggestions.yml new file mode 100644 index 0000000..43c0768 --- /dev/null +++ b/.github/workflows/RunicSuggestions.yml @@ -0,0 +1,9 @@ +name: "Runic Suggestions" +on: pull_request +permissions: + contents: read + pull-requests: write +jobs: + runic-suggestions: + uses: "SciML/.github/.github/workflows/runic-suggestions-on-pr.yml@v1" + secrets: "inherit" diff --git a/.github/workflows/SpellCheck.yml b/.github/workflows/SpellCheck.yml new file mode 100644 index 0000000..1223235 --- /dev/null +++ b/.github/workflows/SpellCheck.yml @@ -0,0 +1,9 @@ +name: Spell Check + +on: [pull_request] + +jobs: + typos-check: + name: Spell Check with Typos + uses: "SciML/.github/.github/workflows/spellcheck.yml@v1" + secrets: "inherit" diff --git a/.github/workflows/TagBot.yml b/.github/workflows/TagBot.yml index f49313b..9ed2d3e 100644 --- a/.github/workflows/TagBot.yml +++ b/.github/workflows/TagBot.yml @@ -1,15 +1,9 @@ -name: TagBot +name: "TagBot" on: issue_comment: - types: - - created + types: [created] workflow_dispatch: jobs: - TagBot: - if: github.event_name == 'workflow_dispatch' || github.actor == 'JuliaTagBot' - runs-on: ubuntu-latest - steps: - - uses: JuliaRegistries/TagBot@v1 - with: - token: ${{ secrets.GITHUB_TOKEN }} - ssh: ${{ secrets.DOCUMENTER_KEY }} + tagbot: + uses: "SciML/.github/.github/workflows/tagbot.yml@v1" + secrets: "inherit" diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index ddce0a4..2722dcc 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -18,39 +18,6 @@ concurrency: cancel-in-progress: ${{ github.ref_name != github.event.repository.default_branch || github.ref != 'refs/tags/v*' }} jobs: - test: - name: "Tests" - strategy: - fail-fast: false - matrix: - version: - - '1' - - 'lts' - - 'nightly' - arch: - - x64 - uses: "SciML/.github/.github/workflows/tests.yml@v1" - with: - julia-version: "${{ matrix.version }}" - julia-arch: "${{ matrix.arch }}" + tests: + uses: "SciML/.github/.github/workflows/grouped-tests.yml@v1" secrets: "inherit" - - docs: - name: Documentation - runs-on: ubuntu-latest - steps: - - uses: actions/checkout@v4 - - uses: julia-actions/setup-julia@v2 - with: - version: '1' - - uses: julia-actions/julia-buildpkg@v1 - - uses: julia-actions/julia-docdeploy@v1 - env: - GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} - - run: | - julia --project=docs -e ' - using Documenter: DocMeta, doctest - using FunctionWrappersWrappers - DocMeta.setdocmeta!(FunctionWrappersWrappers, :DocTestSetup, :(using FunctionWrappersWrappers); recursive=true) - doctest(FunctionWrappersWrappers)' diff --git a/.typos.toml b/.typos.toml new file mode 100644 index 0000000..cabb966 --- /dev/null +++ b/.typos.toml @@ -0,0 +1,143 @@ +[default.extend-words] +# Julia-specific functions +indexin = "indexin" +findfirst = "findfirst" +findlast = "findlast" +eachindex = "eachindex" +setp = "setp" +getp = "getp" +setu = "setu" +getu = "getu" + +# Mathematical/scientific terms +jacobian = "jacobian" +hessian = "hessian" +eigenvalue = "eigenvalue" +eigenvector = "eigenvector" +discretization = "discretization" +linearization = "linearization" +parameterized = "parameterized" +discretized = "discretized" +vectorized = "vectorized" +extrapolant = "extrapolant" +# Person names in citations - these are correct spellings +Sigal = "Sigal" # Sigal Gottlieb is a real person/author +Steinebach = "Steinebach" +Protopapa = "Protopapa" +Hairer = "Hairer" +Wanner = "Wanner" +Verwer = "Verwer" +Sandu = "Sandu" +Carmichael = "Carmichael" +Potra = "Potra" +Seinfeld = "Seinfeld" +Dabdub = "Dabdub" +Marzo = "Marzo" +Loon = "Loon" +RODASP = "RODASP" +Rodas4P = "Rodas4P" +Rodas5P = "Rodas5P" +Rodas6P = "Rodas6P" +Scholz = "Scholz" +Veldhuizen = "Veldhuizen" +Tsit5DA = "Tsit5DA" +Kaps = "Kaps" +Rentrop = "Rentrop" +Rang = "Rang" +Darmstadt = "Darmstadt" +"Méthodes" = "Méthodes" +"différentiels" = "différentiels" +"algébriques" = "algébriques" +"problèmes" = "problèmes" +problemes = "problemes" # French word in citation without accent +Ono = "Ono" # Hiroshi Ono, author of Ono10/Ono12 RK methods +padd = "padd" # variable name in analyticless_convergence_tests.jl +Preprint = "Preprint" +Nr = "Nr" +Wouver = "Wouver" +Sauces = "Sauces" +Schiesser = "Schiesser" +Vande = "Vande" +BIT = "BIT" +Numer = "Numer" +Jcon = "Jcon" +arXiv = "arXiv" + +# (Numer already declared above for "BIT Numer" context.) + +# Technical variable names +dorder = "dorder" # Parameter for delta order / order change + + +# Common variable patterns in Julia/SciML +ists = "ists" +ispcs = "ispcs" +osys = "osys" +rsys = "rsys" +usys = "usys" +fsys = "fsys" +eqs = "eqs" +rhs = "rhs" +lhs = "lhs" +ode = "ode" +pde = "pde" +sde = "sde" +dde = "dde" +bvp = "bvp" +ivp = "ivp" + +# Common abbreviations +tol = "tol" +rtol = "rtol" +atol = "atol" +idx = "idx" +jdx = "jdx" +prev = "prev" +curr = "curr" +init = "init" +tmp = "tmp" +vec = "vec" +arr = "arr" +dt = "dt" +du = "du" +dx = "dx" +dy = "dy" +dz = "dz" + +# Algorithm/type suffixes +alg = "alg" +prob = "prob" +sol = "sol" +cb = "cb" +opts = "opts" +args = "args" +kwargs = "kwargs" + +# Scientific abbreviations +ND = "ND" +nd = "nd" +MTK = "MTK" +ODE = "ODE" +PDE = "PDE" +SDE = "SDE" + +# SDE-specific terms +IIF = "IIF" # Integrating Factor method (IIF1M, IIF2M, IIF1Mil) +iif = "iif" +currate = "currate" # current rate variable in tau-leaping +SIE = "SIE" # SIEA/SIEB algorithm name prefix +SME = "SME" # SME algorithm name (Stochastic Modified Euler) +resetted = "resetted" # variable name in stepsize control +strat = "strat" # Stratonovich abbreviation + +# Variable names used in codebase +yhat = "yhat" # estimated y value +accpt = "accpt" # accepted step flag variable +gam = "gam" # gamma abbreviation in tableaux +clos = "clos" # closure abbreviation + +# Journal abbreviations in BibTeX citations +Comput = "Comput" # J. Comput. Phys. (Journal of Computational Physics) + +# Julia spelling convention +inferrable = "inferrable" # Julia uses "inferrable" not "inferable" diff --git a/Project.toml b/Project.toml index dfe9f36..b6d86e5 100644 --- a/Project.toml +++ b/Project.toml @@ -1,19 +1,38 @@ name = "FunctionWrappersWrappers" uuid = "77dc65aa-8811-40c2-897b-53d922fa7daf" authors = ["Chris Elrod and contributors"] -version = "0.1.4" +version = "1.9.4" [deps] FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e" +PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77" +[weakdeps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" + +[extensions] +FunctionWrappersWrappersEnzymeExt = ["Enzyme", "EnzymeCore"] +FunctionWrappersWrappersMooncakeExt = "Mooncake" + [compat] -FunctionWrappers = "1" -TruncatedStacktraces = "1" -julia = "1.6" +Enzyme = "0.13" +EnzymeCore = "0.8" +FunctionWrappers = "1.1.2" +Mooncake = "0.5.36" +PrecompileTools = "1.1" +SafeTestsets = "0.1, 1" +SciMLTesting = "1" +Test = "1" +TruncatedStacktraces = "1.1" +julia = "1.10" [extras] +SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Test"] +test = ["Test", "SafeTestsets", "SciMLTesting"] diff --git a/docs/Manifest.toml b/docs/Manifest.toml deleted file mode 100644 index 98352d8..0000000 --- a/docs/Manifest.toml +++ /dev/null @@ -1,100 +0,0 @@ -# This file is machine-generated - editing it directly is not advised - -julia_version = "1.7.2-pre.0" -manifest_format = "2.0" - -[[deps.ANSIColoredPrinters]] -git-tree-sha1 = "574baf8110975760d391c710b6341da1afa48d8c" -uuid = "a4c015fc-c6ff-483c-b24f-f7ea428134e9" -version = "0.0.1" - -[[deps.Base64]] -uuid = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f" - -[[deps.Dates]] -deps = ["Printf"] -uuid = "ade2ca70-3891-5945-98fb-dc099432e06a" - -[[deps.DocStringExtensions]] -deps = ["LibGit2"] -git-tree-sha1 = "b19534d1895d702889b219c382a6e18010797f0b" -uuid = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" -version = "0.8.6" - -[[deps.Documenter]] -deps = ["ANSIColoredPrinters", "Base64", "Dates", "DocStringExtensions", "IOCapture", "InteractiveUtils", "JSON", "LibGit2", "Logging", "Markdown", "REPL", "Test", "Unicode"] -git-tree-sha1 = "75c6cf9d99e0efc79b724f5566726ad3ad010a01" -uuid = "e30172f5-a6a5-5a46-863b-614d45cd2de4" -version = "0.27.12" - -[[deps.FunctionWrappersWrappers]] -path = ".." -uuid = "77dc65aa-8811-40c2-897b-53d922fa7daf" -version = "0.1.0" - -[[deps.IOCapture]] -deps = ["Logging", "Random"] -git-tree-sha1 = "f7be53659ab06ddc986428d3a9dcc95f6fa6705a" -uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89" -version = "0.2.2" - -[[deps.InteractiveUtils]] -deps = ["Markdown"] -uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240" - -[[deps.JSON]] -deps = ["Dates", "Mmap", "Parsers", "Unicode"] -git-tree-sha1 = "8076680b162ada2a031f707ac7b4953e30667a37" -uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6" -version = "0.21.2" - -[[deps.LibGit2]] -deps = ["Base64", "NetworkOptions", "Printf", "SHA"] -uuid = "76f85450-5226-5b5a-8eaa-529ad045b433" - -[[deps.Logging]] -uuid = "56ddb016-857b-54e1-b83d-db4d58db5568" - -[[deps.Markdown]] -deps = ["Base64"] -uuid = "d6f4376e-aef5-505a-96c1-9c027394607a" - -[[deps.Mmap]] -uuid = "a63ad114-7e13-5084-954f-fe012c677804" - -[[deps.NetworkOptions]] -uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908" - -[[deps.Parsers]] -deps = ["Dates"] -git-tree-sha1 = "0b5cfbb704034b5b4c1869e36634438a047df065" -uuid = "69de0a69-1ddd-5017-9359-2bf0b02dc9f0" -version = "2.2.1" - -[[deps.Printf]] -deps = ["Unicode"] -uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7" - -[[deps.REPL]] -deps = ["InteractiveUtils", "Markdown", "Sockets", "Unicode"] -uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb" - -[[deps.Random]] -deps = ["SHA", "Serialization"] -uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" - -[[deps.SHA]] -uuid = "ea8e919c-243c-51af-8825-aaa63cd721ce" - -[[deps.Serialization]] -uuid = "9e88b42a-f829-5b0c-bbe9-9e923198166b" - -[[deps.Sockets]] -uuid = "6462fe0b-24de-5631-8697-dd941f90decc" - -[[deps.Test]] -deps = ["InteractiveUtils", "Logging", "Random", "Serialization"] -uuid = "8dfed614-e22c-5e08-85e1-65c5234f0b40" - -[[deps.Unicode]] -uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5" diff --git a/docs/Project.toml b/docs/Project.toml index f1a23e2..c279ceb 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,3 +1,7 @@ [deps] Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf" + +[compat] +Documenter = "1.16.1" +FunctionWrappersWrappers = "0.1.3, 1.1" diff --git a/docs/make.jl b/docs/make.jl index bcbf26c..cc6bdcf 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,8 +1,10 @@ using FunctionWrappersWrappers using Documenter -DocMeta.setdocmeta!(FunctionWrappersWrappers, :DocTestSetup, - :(using FunctionWrappersWrappers); recursive = true) +DocMeta.setdocmeta!( + FunctionWrappersWrappers, :DocTestSetup, + :(using FunctionWrappersWrappers); recursive = true +) makedocs(; modules = [FunctionWrappersWrappers], diff --git a/ext/FunctionWrappersWrappersEnzymeExt.jl b/ext/FunctionWrappersWrappersEnzymeExt.jl new file mode 100644 index 0000000..1c05580 --- /dev/null +++ b/ext/FunctionWrappersWrappersEnzymeExt.jl @@ -0,0 +1,329 @@ +module FunctionWrappersWrappersEnzymeExt + +using FunctionWrappersWrappers +using FunctionWrappersWrappers: SingleCacheStorage, DictCacheStorage, NoCacheStorage +using Enzyme +using EnzymeCore +using EnzymeCore.EnzymeRules + +# ============================================================================= +# Mark cache-storage types as inactive +# ============================================================================= +# `SingleCacheStorage` and `DictCacheStorage` are mutable / contain a `Dict`, +# and their cache-miss branches write to that storage. Without these +# declarations Enzyme conservatively treats any closure that *might* touch a +# `FunctionWrappersWrapper` (e.g. via `remake(prob; p = …)` capturing the +# problem in scope) as potentially writing to the wrapper's cache, and +# refuses to prove the captured argument read-only. The cache values are +# `FunctionWrapper`s used purely for dispatch / dynamic call speedup; they +# never hold derivative data. +EnzymeCore.EnzymeRules.inactive_type(::Type{<:SingleCacheStorage}) = true +EnzymeCore.EnzymeRules.inactive_type(::Type{<:DictCacheStorage}) = true +EnzymeCore.EnzymeRules.inactive_type(::Type{NoCacheStorage}) = true + +# ============================================================================= +# Helper: build a Forward mode from FwdConfig flags +# ============================================================================= +# The outer caller may invoke `Enzyme.autodiff(set_runtime_activity(Forward), …)` +# or `set_strong_zero(Forward)` or `ForwardWithPrimal`. Those settings flow +# into the `EnzymeRules.FwdConfig{NeedsPrimal, NeedsShadow, Width, +# RuntimeActivity, StrongZero}` type parameters of the rule's first argument. +# Before this fix the rules hard-coded plain `Forward` in their inner +# `Enzyme.autodiff` delegation, which dropped both `RuntimeActivity` and +# `StrongZero` — breaking users who need `set_runtime_activity(Forward)` to +# avoid `EnzymeRuntimeActivityError` inside the wrapped function (the SciML +# `Rosenbrock23(autodiff = AutoEnzyme(set_runtime_activity(Forward)))` path +# on an in-place time-dependent RHS; see OrdinaryDiffEq.jl PR #3518). +# +# `_fwd_mode(needs_primal, RuntimeActivity, StrongZero)` returns the +# `ForwardMode` matching the outer config so the delegated call inherits +# those flags. +@inline function _fwd_mode( + ::Val{NeedsPrimal}, ::Val{RuntimeActivity}, ::Val{StrongZero} + ) where {NeedsPrimal, RuntimeActivity, StrongZero} + mode = NeedsPrimal ? ForwardWithPrimal : Forward + RuntimeActivity && (mode = Enzyme.set_runtime_activity(mode)) + StrongZero && (mode = Enzyme.set_strong_zero(mode)) + return mode +end + +# ============================================================================= +# Forward mode rules — generalized to arbitrary batch width W +# ============================================================================= + +# Shadow only (Forward mode, no primal) +# +# `func` is `Annotation{<:FunctionWrappersWrapper}` rather than +# `Const{<:FunctionWrappersWrapper}` so that callers passing +# `Duplicated{<:FunctionWrappersWrapper}` also dispatch here. Enzyme drives +# the rule that way when the outer `autodiff` call is differentiating through +# a closure that carries an FWW (e.g. NonlinearSolve + SciMLSensitivity, see +# SciML/FunctionWrappersWrappers.jl#48). The FWW struct itself only carries +# `FunctionWrapper`s plus cache storage — none of those fields have a +# meaningful tangent — so the function shadow is ignored and the inner +# `Enzyme.autodiff` call uses `Const(f_orig)`. +function EnzymeRules.forward( + ::EnzymeRules.FwdConfig{false, true, W, RuntimeActivity, StrongZero}, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + RT::Type{<:EnzymeCore.Annotation{T}}, + args::Vararg{EnzymeCore.Annotation, N} + ) where {T, W, N, RuntimeActivity, StrongZero} + f_orig = unwrap(func.val) + mode = _fwd_mode(Val(false), Val(RuntimeActivity), Val(StrongZero)) + if W == 1 + shadow_result = Enzyme.autodiff(mode, Const(f_orig), Duplicated{T}, args...) + return shadow_result[1]::T + else + shadow_result = Enzyme.autodiff(mode, Const(f_orig), BatchDuplicated{T, W}, args...) + # Enzyme returns the batch shadow as an `AnonymousStruct` — a + # `NamedTuple{(:1, :2, …), NTuple{W, T}}` (see + # `Enzyme.Compiler.AnonymousStruct` in `Enzyme/src/compiler/utils.jl`). + # Convert to a plain tuple so the rule's return matches the + # `BatchDuplicated` shadow contract Enzyme expects from a forward rule. + return Tuple(shadow_result[1])::NTuple{W, T} + end +end + +# Both primal and shadow (ForwardWithPrimal mode) +function EnzymeRules.forward( + ::EnzymeRules.FwdConfig{true, true, W, RuntimeActivity, StrongZero}, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + RT::Type{<:EnzymeCore.Annotation{T}}, + args::Vararg{EnzymeCore.Annotation, N} + ) where {T, W, N, RuntimeActivity, StrongZero} + f_orig = unwrap(func.val) + pargs = ntuple(i -> args[i].val, Val(N)) + primal = f_orig(pargs...)::T + # Use plain Forward (not ForwardWithPrimal) here — we already have the + # primal above, and `Duplicated{T}` / `BatchDuplicated{T,W}` as the RT + # annotation asks only for the shadow. + mode = _fwd_mode(Val(false), Val(RuntimeActivity), Val(StrongZero)) + if W == 1 + shadow_result = Enzyme.autodiff(mode, Const(f_orig), Duplicated{T}, args...) + shadow = shadow_result[1]::T + return Duplicated(primal, shadow) + else + shadow_result = Enzyme.autodiff(mode, Const(f_orig), BatchDuplicated{T, W}, args...) + # See the comment on the {false, true} rule — `shadow_result[1]` is a + # NamedTuple, not an NTuple. + shadows = Tuple(shadow_result[1])::NTuple{W, T} + return BatchDuplicated(primal, shadows) + end +end + +# Primal only (Const return type) — width-independent +function EnzymeRules.forward( + ::EnzymeRules.FwdConfig{true, false, W, RuntimeActivity, StrongZero}, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + RT::Type{<:EnzymeCore.Annotation}, + args::Vararg{EnzymeCore.Annotation, N} + ) where {W, N, RuntimeActivity, StrongZero} + f_orig = unwrap(func.val) + pargs = ntuple(i -> args[i].val, Val(N)) + return f_orig(pargs...) +end + +# Neither primal nor shadow requested in the RETURN. Enzyme dispatches on +# this combo for IIP functions (Const return type) where the caller still +# needs primal and shadow propagation through `Duplicated` args — e.g. SciML +# solvers calling an IIP RHS via `AutoEnzyme(…, function_annotation = Const)`. +# The previous revision ran `f_orig(pargs...)` by hand; that mutated the +# primal IIP buffer but left `Duplicated` shadow buffers untouched, giving +# trivial Jacobians and blowing up Rodas4/5/Veldd4 error tolerances 4–9 +# orders of magnitude in OrdinaryDiffEq.jl v7. Delegate to `Enzyme.autodiff` +# on the unwrapped function with a Const return annotation so the Duplicated +# arg shadows are propagated correctly and no return is produced. +# +# IMPORTANT: forward the `RuntimeActivity` and `StrongZero` flags from the +# outer config into the delegated `Enzyme.autodiff` call. Prior to this +# fix the rule hard-coded `Forward`, silently dropping +# `set_runtime_activity(Forward)` on the way down into `f_orig`. +function EnzymeRules.forward( + ::EnzymeRules.FwdConfig{false, false, W, RuntimeActivity, StrongZero}, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + RT::Type{<:EnzymeCore.Annotation}, + args::Vararg{EnzymeCore.Annotation, N} + ) where {W, N, RuntimeActivity, StrongZero} + f_orig = unwrap(func.val) + mode = _fwd_mode(Val(false), Val(RuntimeActivity), Val(StrongZero)) + Enzyme.autodiff(mode, Const(f_orig), Const, args...) + return nothing +end + +# ============================================================================= +# Reverse mode rules +# ============================================================================= + +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfig, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + RT::Type{<:EnzymeCore.Active{T}}, + args::Vararg{EnzymeCore.Annotation, N} + ) where {T, N} + f_orig = unwrap(func.val) + pargs = ntuple(i -> args[i].val, Val(N)) + result = f_orig(pargs...)::T + + if EnzymeRules.needs_primal(config) + return EnzymeRules.AugmentedReturn(result, nothing, nothing) + else + return EnzymeRules.AugmentedReturn(nothing, nothing, nothing) + end +end + +# Const return (e.g. IIP functions returning Nothing, or any non-differentiated +# return). Just run the primal for its side effects; no tape is needed because +# the reverse pass has nothing to propagate back from the return. +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfig, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + RT::Type{<:EnzymeCore.Const}, + args::Vararg{EnzymeCore.Annotation, N} + ) where {N} + f_orig = unwrap(func.val) + pargs = ntuple(i -> args[i].val, Val(N)) + f_orig(pargs...) + return EnzymeRules.AugmentedReturn(nothing, nothing, nothing) +end + +# Duplicated / BatchDuplicated return: record the primal so that reverse has +# it available when propagating dret through the arguments. +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfig, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + RT::Type{<:EnzymeCore.Duplicated{T}}, + args::Vararg{EnzymeCore.Annotation, N} + ) where {T, N} + f_orig = unwrap(func.val) + pargs = ntuple(i -> args[i].val, Val(N)) + primal = f_orig(pargs...)::T + if EnzymeRules.needs_primal(config) + return EnzymeRules.AugmentedReturn(primal, zero(primal), nothing) + else + return EnzymeRules.AugmentedReturn(nothing, zero(primal), nothing) + end +end + +function EnzymeRules.augmented_primal( + config::EnzymeRules.RevConfig, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + RT::Type{<:EnzymeCore.BatchDuplicated{T, W}}, + args::Vararg{EnzymeCore.Annotation, N} + ) where {T, W, N} + f_orig = unwrap(func.val) + pargs = ntuple(i -> args[i].val, Val(N)) + primal = f_orig(pargs...)::T + shadows = ntuple(_ -> zero(primal), Val(W)) + if EnzymeRules.needs_primal(config) + return EnzymeRules.AugmentedReturn(primal, shadows, nothing) + else + return EnzymeRules.AugmentedReturn(nothing, shadows, nothing) + end +end + +# Helper: build a Forward mode reflecting a RevConfig's runtime_activity / +# strong_zero flags so the internal forward-mode delegation inside reverse +# rules inherits the user's outer config. +@inline function _fwd_mode_from_rev(config::EnzymeRules.RevConfig) + mode = Forward + EnzymeRules.runtime_activity(config) && (mode = Enzyme.set_runtime_activity(mode)) + EnzymeRules.strong_zero(config) && (mode = Enzyme.set_strong_zero(mode)) + return mode +end + +# Varargs reverse: compute each partial via forward-mode AD on the unwrapped +# function, then scale by dret. This avoids type-inference issues that arise +# from calling autodiff(Reverse, Const{Any}(...), ...). +@generated function _fww_reverse_grads( + mode, f_orig, dret_val::T, args::Vararg{EnzymeCore.Active, N} + ) where {T, N} + # Build forward-mode calls for each partial derivative + exprs = [] + for i in 1:N + seeds = [j == i ? :(one(eltype(typeof(args[$j])))) : :(zero(eltype(typeof(args[$j])))) for j in 1:N] + dups = [:(Duplicated(args[$j].val, $(seeds[j]))) for j in 1:N] + Ti = :(eltype(typeof(args[$i]))) + push!( + exprs, quote + fwd = Enzyme.autodiff(mode, Const(f_orig), Duplicated{$T}, $(dups...)) + $Ti(fwd[1] * dret_val)::$Ti + end + ) + end + return Expr(:tuple, exprs...) +end + +function EnzymeRules.reverse( + config::EnzymeRules.RevConfig, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + dret::EnzymeCore.Active{T}, + tape, + args::Vararg{EnzymeCore.Active, N} + ) where {T, N} + f_orig = unwrap(func.val) + return _fww_reverse_grads(_fwd_mode_from_rev(config), f_orig, dret.val, args...) +end + +# Handle mixed Active/Const args: return nothing for Const, gradient for Active +function EnzymeRules.reverse( + config::EnzymeRules.RevConfig, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + dret::EnzymeCore.Active, + tape, + args::Vararg{EnzymeCore.Annotation, N} + ) where {N} + f_orig = unwrap(func.val) + dret_val = dret.val + mode = _fwd_mode_from_rev(config) + return ntuple(Val(N)) do i + if args[i] isa EnzymeCore.Const + nothing + else + # Use forward-mode to compute partial derivative + dup_args = ntuple(Val(N)) do j + if j == i + Duplicated(args[j].val, one(eltype(typeof(args[j])))) + else + Duplicated(args[j].val, zero(eltype(typeof(args[j])))) + end + end + fwd = Enzyme.autodiff(mode, Const(f_orig), Duplicated, dup_args...) + fwd[1] * dret_val + end + end +end + +# Const return — Enzyme passes the RT as a `Type{<:Const}` to `reverse`, not +# as an instance. Delegate the reverse pass to +# `Enzyme.autodiff(Reverse, Const(f_orig), Const, args...)` so gradients +# accumulate into any `Duplicated` arg shadow buffers (the SciML IIP +# pattern). Simply returning `nothing` left Duplicated shadows at zero. +# +# Per Enzyme's rule return-type protocol, `Active` args require a concrete +# scalar gradient (not `nothing`). Under a `Const` return there is no +# gradient source, so Active arg gradients are zero. `Duplicated` / +# `BatchDuplicated` args return `nothing` because their gradients are +# accumulated in-place by the `Enzyme.autodiff(Reverse, …)` call above. +function EnzymeRules.reverse( + config::EnzymeRules.RevConfig, + func::EnzymeCore.Annotation{<:FunctionWrappersWrapper}, + dret::Type{<:EnzymeCore.Const}, + tape, + args::Vararg{EnzymeCore.Annotation, N} + ) where {N} + f_orig = unwrap(func.val) + # Only worth invoking Enzyme.autodiff when at least one arg is + # Duplicated/BatchDuplicated — otherwise there's nothing to accumulate. + if any(a -> a isa EnzymeCore.Duplicated || a isa EnzymeCore.BatchDuplicated, args) + Enzyme.autodiff(Reverse, Const(f_orig), Const, args...) + end + return ntuple(Val(N)) do i + if args[i] isa EnzymeCore.Active + zero(eltype(typeof(args[i]))) + else + nothing + end + end +end + +end diff --git a/ext/FunctionWrappersWrappersMooncakeExt.jl b/ext/FunctionWrappersWrappersMooncakeExt.jl new file mode 100644 index 0000000..3e2aa17 --- /dev/null +++ b/ext/FunctionWrappersWrappersMooncakeExt.jl @@ -0,0 +1,36 @@ +module FunctionWrappersWrappersMooncakeExt + +using FunctionWrappersWrappers +import Mooncake +using Mooncake: @is_primitive, MinimalCtx, CoDual, NoRData, zero_tangent, NoTangent, fdata + +# Make calling a FunctionWrappersWrapper a Mooncake primitive. +# Instead of differentiating through the FunctionWrapper dispatch machinery +# (which fails because the tuple of differently-typed FunctionWrappers produces +# incompatible FunctionWrapperTangent types), unwrap to the original function +# and differentiate through that directly. + +@is_primitive MinimalCtx Tuple{<:FunctionWrappersWrapper, Vararg} + +function Mooncake.rrule!!( + f::CoDual{<:FunctionWrappersWrapper}, args::Vararg{CoDual}, + ) + f_orig = unwrap(f.x) + # Build a derived rule for calling the unwrapped function with these arg types. + # We can't use rrule!! directly since the unwrapped function (e.g. SciMLBase.Void) + # is generally not a Mooncake primitive — it needs a derived (compiled) rule. + sig = Tuple{typeof(f_orig), map(Core.Typeof ∘ Mooncake.primal, args)...} + rule = Mooncake.build_rrule(sig) + # Use fdata to get the correct tangent component for the CoDual — zero_tangent + # returns NoTangent for singleton callables but derived rules expect NoFData. + f_orig_codual = CoDual(f_orig, fdata(zero_tangent(f_orig))) + y, pb = rule(f_orig_codual, args...) + fww_pb(dy) = (NoRData(), Base.tail(pb(dy))...) + return y, fww_pb +end + +# FunctionWrappersWrapper is not differentiable data itself — the wrapped function +# is what carries the derivative information, and we handle that in the rrule above. +Mooncake.tangent_type(::Type{<:FunctionWrappersWrapper}) = NoTangent + +end diff --git a/src/FunctionWrappersWrappers.jl b/src/FunctionWrappersWrappers.jl index 7da5b9b..a99d624 100644 --- a/src/FunctionWrappersWrappers.jl +++ b/src/FunctionWrappersWrappers.jl @@ -1,51 +1,332 @@ module FunctionWrappersWrappers -using FunctionWrappers +using FunctionWrappers: FunctionWrappers import TruncatedStacktraces -export FunctionWrappersWrapper +export FunctionWrappersWrapper, unwrap, wrapped_signatures, wrapped_return_types +export NoCache, SingleCache, DictCache +export Strict, AllowAll, AllowNonIsBits -struct FunctionWrappersWrapper{FW, FB} +# ============================================================================ +# Cache modes: control how fallback FunctionWrappers are cached +# ============================================================================ +abstract type AbstractCacheMode end + +""" + NoCache() + +No caching — every fallback call goes through dynamic dispatch (`obj[](arg...)`), +incurring 1 allocation per call. +""" +struct NoCache <: AbstractCacheMode end + +""" + SingleCache() + +Cache a single `FunctionWrapper` for the last-seen argument types. After the first +fallback call, subsequent calls with the same types are zero-allocation. If called with +different types, the cache is replaced (1 alloc on miss). This is the recommended default. +""" +struct SingleCache <: AbstractCacheMode end + +""" + DictCache() + +Cache `FunctionWrapper`s in a `Dict` keyed by argument type. Handles multiple +non-isbits types without thrashing. Slightly higher lookup overhead than `SingleCache`. +""" +struct DictCache <: AbstractCacheMode end + +# ============================================================================ +# Fallback policies: control when fallback is allowed +# ============================================================================ +abstract type AbstractFallbackPolicy end + +""" + Strict() + +Never fall back — throw `NoFunctionWrapperFoundError` if no wrapper matches. +""" +struct Strict <: AbstractFallbackPolicy end + +""" + AllowAll() + +Always fall back to the original function when no wrapper matches. +""" +struct AllowAll <: AbstractFallbackPolicy end + +""" + AllowNonIsBits() + +Fall back only when argument types contain non-isbits elements (e.g., `BigFloat`, +`SparseConnectivityTracer` types). Throws `NoFunctionWrapperFoundError` for isbits +type mismatches (e.g., `Float32` when `Float64` was expected), which catches bugs. +This is the recommended default. +""" +struct AllowNonIsBits <: AbstractFallbackPolicy end + +# ============================================================================ +# Cache storage types +# ============================================================================ +struct NoCacheStorage end +mutable struct SingleCacheStorage + cached::Any # Union{Nothing, FunctionWrapper} + SingleCacheStorage() = new(nothing) +end +struct DictCacheStorage + cache::Dict{DataType, Any} + DictCacheStorage() = new(Dict{DataType, Any}()) +end + +_make_cache_storage(::NoCache) = NoCacheStorage() +_make_cache_storage(::SingleCache) = SingleCacheStorage() +_make_cache_storage(::DictCache) = DictCacheStorage() + +# ============================================================================ +# Main type +# ============================================================================ + +""" + FunctionWrappersWrapper{FW, P, CS} + +A wrapper around a tuple of `FunctionWrapper`s that dispatches calls to the +matching wrapper based on argument types. When no wrapper matches, behavior is +controlled by the fallback policy `P` and cache mode `CS`. + +# Type parameters +- `FW`: Tuple type of `FunctionWrapper`s +- `P`: Fallback policy (`Strict`, `AllowAll`, or `AllowNonIsBits`) +- `CS`: Cache storage type (`NoCacheStorage`, `SingleCacheStorage`, `DictCacheStorage`) +""" +struct FunctionWrappersWrapper{FW, P, CS} fw::FW + cache_storage::CS + function FunctionWrappersWrapper{FW, P, CS}( + fw::FW, cs::CS + ) where {FW, P, CS} + return new{FW, P, CS}(fw, cs) + end end TruncatedStacktraces.@truncate_stacktrace FunctionWrappersWrapper -function (fww::FunctionWrappersWrapper{FW, FB})(args::Vararg{Any, K}) where {FW, K, FB} - _call(fww.fw, args, fww) +""" + FunctionWrappersWrapper{FW, P, CS}(f) + +Create a `FunctionWrappersWrapper` when the type parameters are specified. + +# Arguments +- `f`: The function to wrap + +# Type parameters +- `FW`: Tuple type of `FunctionWrapper`s +- `P`: Fallback policy (`Strict`, `AllowAll`, or `AllowNonIsBits`) +- `CS`: Cache storage type (`NoCacheStorage`, `SingleCacheStorage`, `DictCacheStorage`) +""" +function FunctionWrappersWrapper{FW, P, CS}(f) where {K, FW <: NTuple{K, Any}, P, CS} + fw = ntuple(i -> FW.parameters[i](f), Val(K)) + cs = CS() + return FunctionWrappersWrapper{FW, P, CS}(fw, cs) +end + +""" + FunctionWrappersWrapper(f, argtypes, rettypes; cache=SingleCache(), policy=AllowNonIsBits()) + +Create a `FunctionWrappersWrapper` with configurable fallback behavior. + +# Arguments +- `f`: The function to wrap +- `argtypes`: Tuple of argument type signatures (e.g., `(Tuple{Float64, Float64},)`) +- `rettypes`: Tuple of return types (e.g., `(Float64,)`) + +# Keywords +- `cache`: Cache mode for fallback path — `NoCache()`, `SingleCache()` (default), or `DictCache()` +- `policy`: Fallback policy — `Strict()`, `AllowAll()`, or `AllowNonIsBits()` (default) +""" +function FunctionWrappersWrapper( + f::F, argtypes::Tuple{Vararg{Any, K}}, rettypes::Tuple{Vararg{Type, K}}; + cache::AbstractCacheMode = SingleCache(), + policy::AbstractFallbackPolicy = AllowNonIsBits() + ) where {F, K} + fwt = map(argtypes, rettypes) do A, R + FunctionWrappers.FunctionWrapper{R, A}(f) + end + cs = _make_cache_storage(cache) + return FunctionWrappersWrapper{typeof(fwt), typeof(policy), typeof(cs)}(fwt, cs) +end + +Base.convert(::Type{T}, obj) where {T <: FunctionWrappersWrapper} = T(obj) +Base.convert(::Type{T}, obj::T) where {T <: FunctionWrappersWrapper} = obj + +# ============================================================================ +# Call dispatch — entry point +# ============================================================================ + +function (fww::FunctionWrappersWrapper{FW, P, CS})( + args::Vararg{Any, K} + ) where {FW, K, P, CS} + return _call(fww.fw, args, fww) end -function _call(fw::Tuple{FunctionWrappers.FunctionWrapper{R, A}, Vararg}, - arg::A, fww::FunctionWrappersWrapper) where {R, A} - first(fw)(arg...) +# Match path: try each FunctionWrapper in order +function _call( + fw::Tuple{FunctionWrappers.FunctionWrapper{R, A}, Vararg}, + arg::A, fww::FunctionWrappersWrapper + ) where {R, A} + return first(fw)(arg...) end -function _call(fw::Tuple{FunctionWrappers.FunctionWrapper{R, A1}, Vararg}, - arg::A2, fww::FunctionWrappersWrapper) where {R, A1, A2} - _call(Base.tail(fw), arg, fww) +function _call( + fw::Tuple{FunctionWrappers.FunctionWrapper{R, A1}, Vararg}, + arg::A2, fww::FunctionWrappersWrapper + ) where {R, A1, A2} + return _call(Base.tail(fw), arg, fww) end +# ============================================================================ +# Fallback — Strict: always error +# ============================================================================ + const NO_FUNCTIONWRAPPER_FOUND_MESSAGE = "No matching function wrapper was found!" struct NoFunctionWrapperFoundError <: Exception end function Base.showerror(io::IO, e::NoFunctionWrapperFoundError) - print(io, NO_FUNCTIONWRAPPER_FOUND_MESSAGE) + return print(io, NO_FUNCTIONWRAPPER_FOUND_MESSAGE) end -function _call(::Tuple{}, arg, fww::FunctionWrappersWrapper{<:Any, false}) +function _call(::Tuple{}, arg, fww::FunctionWrappersWrapper{<:Any, Strict}) throw(NoFunctionWrapperFoundError()) end -function _call(::Tuple{}, arg, fww::FunctionWrappersWrapper{<:Any, true}) - first(fww.fw).obj[](arg...) + +# ============================================================================ +# Fallback — AllowAll: always fall back +# ============================================================================ + +function _call(::Tuple{}, arg, fww::FunctionWrappersWrapper{<:Any, AllowAll}) + return _fallback(arg, fww) end -function FunctionWrappersWrapper( - f::F, argtypes::Tuple{Vararg{Any, K}}, rettypes::Tuple{Vararg{DataType, K}}, - fallback::Val{FB} = Val{false}()) where {F, K, FB} - fwt = map(argtypes, rettypes) do A, R - FunctionWrappers.FunctionWrapper{R, A}(f) +# ============================================================================ +# Fallback — AllowNonIsBits: fall back only for non-isbits arg types +# ============================================================================ + +function _call( + ::Tuple{}, arg::A, fww::FunctionWrappersWrapper{<:Any, AllowNonIsBits} + ) where {A} + if _has_non_isbits_args(A) + return _fallback(arg, fww) + end + throw(NoFunctionWrapperFoundError()) +end + +@generated function _has_non_isbits_args(::Type{T}) where {T <: Tuple} + checks = [] + for P in T.parameters + if P <: AbstractArray + push!(checks, :(!(isbitstype($(eltype(P)))))) + else + push!(checks, :(!(isbitstype($P)))) + end + end + isempty(checks) && return :(false) + return Expr(:||, checks...) +end + +# ============================================================================ +# Fallback execution — dispatch on cache storage type +# ============================================================================ + +# --- NoCache: direct dynamic dispatch every time --- +function _fallback(arg, fww::FunctionWrappersWrapper{<:Any, <:Any, NoCacheStorage}) + return first(fww.fw).obj[](arg...) +end + +# --- SingleCache: cache one FunctionWrapper for the last arg types --- +function _fallback( + arg::A, fww::FunctionWrappersWrapper{<:Any, <:Any, SingleCacheStorage} + ) where {A} + cached = fww.cache_storage.cached + if cached isa FunctionWrappers.FunctionWrapper{Any, A} + return cached(arg...) + end + f = first(fww.fw).obj[] + new_fw = FunctionWrappers.FunctionWrapper{Any, A}(f) + fww.cache_storage.cached = new_fw + return new_fw(arg...) +end + +# --- DictCache: cache FunctionWrappers keyed by arg type --- +function _fallback( + arg::A, fww::FunctionWrappersWrapper{<:Any, <:Any, DictCacheStorage} + ) where {A} + cached = get(fww.cache_storage.cache, A, nothing) + if cached isa FunctionWrappers.FunctionWrapper{Any, A} + return cached(arg...) + end + f = first(fww.fw).obj[] + new_fw = FunctionWrappers.FunctionWrapper{Any, A}(f) + fww.cache_storage.cache[A] = new_fw + return new_fw(arg...) +end + +# ============================================================================ +# Introspection +# ============================================================================ + +""" + unwrap(fww::FunctionWrappersWrapper) + +Return the original function that was wrapped. +""" +unwrap(fww::FunctionWrappersWrapper) = first(fww.fw).obj[] + +""" + wrapped_signatures(fww::FunctionWrappersWrapper) + +Return a tuple of the argument type signatures that the wrapper can dispatch on. +""" +function wrapped_signatures(fww::FunctionWrappersWrapper) + return map(fw -> typeof(fw).parameters[2], fww.fw) +end + +""" + wrapped_return_types(fww::FunctionWrappersWrapper) + +Return a tuple of the return types for each wrapped function signature. +""" +function wrapped_return_types(fww::FunctionWrappersWrapper) + return map(fw -> typeof(fw).parameters[1], fww.fw) +end + +# ============================================================================ +# Precompilation +# ============================================================================ + +using PrecompileTools: @compile_workload, @setup_workload + +@setup_workload begin + @compile_workload begin + fw_binary = FunctionWrappersWrapper( + +, + (Tuple{Float64, Float64}, Tuple{Int, Int}), + (Float64, Int) + ) + fw_binary(1.0, 2.0) + fw_binary(1, 2) + + fw_unary = FunctionWrappersWrapper( + abs, + (Tuple{Float64}, Tuple{Int}), + (Float64, Int) + ) + fw_unary(1.0) + fw_unary(1) + + unwrap(fw_unary) + wrapped_signatures(fw_binary) + wrapped_return_types(fw_binary) end - FunctionWrappersWrapper{typeof(fwt), FB}(fwt) end end diff --git a/test/Enzyme/Project.toml b/test/Enzyme/Project.toml new file mode 100644 index 0000000..a36c643 --- /dev/null +++ b/test/Enzyme/Project.toml @@ -0,0 +1,15 @@ +[deps] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" +FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf" +SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +FunctionWrappersWrappers = {path = "../.."} + +[compat] +Enzyme = "0.13" +EnzymeCore = "0.8" +SafeTestsets = "0.0.1, 0.1, 1" +Test = "1" diff --git a/test/Enzyme/enzyme_tests.jl b/test/Enzyme/enzyme_tests.jl new file mode 100644 index 0000000..5a8f5f9 --- /dev/null +++ b/test/Enzyme/enzyme_tests.jl @@ -0,0 +1,517 @@ +using FunctionWrappersWrappers +using Enzyme +using EnzymeCore +using Test + +@testset "Enzyme forward mode" begin + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + + # Forward mode (shadow only) + result = Enzyme.autodiff(Forward, Const(fww), Duplicated, Duplicated(3.0, 1.0)) + @test result[1] ≈ 6.0 + + # ForwardWithPrimal (both primal and shadow) + result = Enzyme.autodiff(ForwardWithPrimal, Const(fww), Duplicated, Duplicated(3.0, 1.0)) + @test result[1] ≈ 6.0 # shadow + @test result[2] ≈ 9.0 # primal +end + +@testset "Enzyme reverse mode - single arg" begin + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + + # Reverse mode + result = Enzyme.autodiff(Reverse, Const(fww), Active, Active(3.0)) + @test result[1][1] ≈ 6.0 + + # ReverseWithPrimal + result = Enzyme.autodiff(ReverseWithPrimal, Const(fww), Active, Active(3.0)) + @test result[1][1] ≈ 6.0 # gradient + @test result[2] ≈ 9.0 # primal +end + +@testset "Enzyme reverse mode - multi arg" begin + g(x, y) = x * y + x^2 + fww = FunctionWrappersWrapper(g, (Tuple{Float64, Float64},), (Float64,)) + + # g(x,y) = x*y + x^2 → ∂g/∂x = y + 2x, ∂g/∂y = x + result = Enzyme.autodiff(Reverse, Const(fww), Active, Active(3.0), Active(4.0)) + @test result[1][1] ≈ 10.0 # ∂g/∂x at (3,4) = 4 + 6 + @test result[1][2] ≈ 3.0 # ∂g/∂y at (3,4) = 3 +end + +@testset "Enzyme with trig functions" begin + fww_sin = FunctionWrappersWrapper(sin, (Tuple{Float64},), (Float64,)) + + # sin'(x) = cos(x) + result = Enzyme.autodiff(Forward, Const(fww_sin), Duplicated, Duplicated(1.0, 1.0)) + @test result[1] ≈ cos(1.0) + + result = Enzyme.autodiff(Reverse, Const(fww_sin), Active, Active(1.0)) + @test result[1][1] ≈ cos(1.0) +end + +@testset "Enzyme batch forward mode (width > 1)" begin + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + + # Batch width = 2: compute derivatives for two tangent directions simultaneously. + # f(x) = x^2 → f'(x) = 2x; at x=3.0 with tangents (1.0, 2.0) → shadows (6.0, 12.0) + result = Enzyme.autodiff( + Forward, Const(fww), BatchDuplicated, + BatchDuplicated(3.0, (1.0, 2.0)) + ) + shadows = result[1] + @test shadows[1] ≈ 6.0 # f'(3) * 1.0 + @test shadows[2] ≈ 12.0 # f'(3) * 2.0 + + # ForwardWithPrimal, batch width = 2 + result_wp = Enzyme.autodiff( + ForwardWithPrimal, Const(fww), BatchDuplicated, + BatchDuplicated(3.0, (1.0, 2.0)) + ) + @test result_wp[1][1] ≈ 6.0 # shadow 1 + @test result_wp[1][2] ≈ 12.0 # shadow 2 + @test result_wp[2] ≈ 9.0 # primal f(3) = 9 +end + +@testset "Enzyme batch forward rule return type is NTuple, not NamedTuple" begin + # Regression test for the typeassert bug: the inner + # `Enzyme.autodiff(Forward, …, BatchDuplicated{T,W}, …)` returns the + # batch shadow wrapped in `Enzyme.Compiler.AnonymousStruct` — a + # `NamedTuple{(:1, :2, …), NTuple{W, T}}`. The rule must convert it + # to a plain `NTuple{W, T}` before returning, otherwise the + # `::NTuple{W, T}` typeassert fires and surfaces as: + # TypeError: in typeassert, expected Tuple{Float64, Float64}, + # got a value of type @NamedTuple{1::Float64, 2::Float64} + # + # The outer `Enzyme.autodiff` testset above doesn't catch this on its + # own because the outer call ALSO wraps the result in + # `AnonymousStruct`, and `shadow[1] / shadow[2]` indexing works on + # both `NamedTuple` and `Tuple`. Call `EnzymeRules.forward` + # directly so we observe the rule's actual return value and can + # assert its concrete type. + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + + # {NeedsPrimal=false, NeedsShadow=true, W=2, RuntimeActivity=false, + # StrongZero=false} — the shadow-only batch branch. + config_shadow = EnzymeCore.EnzymeRules.FwdConfig{false, true, 2, false, false}() + shadow = EnzymeCore.EnzymeRules.forward( + config_shadow, Const(fww), EnzymeCore.BatchDuplicated{Float64, 2}, + BatchDuplicated(3.0, (1.0, 2.0)) + ) + @test shadow isa NTuple{2, Float64} + @test !(shadow isa NamedTuple) + @test shadow == (6.0, 12.0) + + # {NeedsPrimal=true, NeedsShadow=true, W=2, …} — ForwardWithPrimal + # batch branch. Same conversion bug existed on this path. + config_primal = EnzymeCore.EnzymeRules.FwdConfig{true, true, 2, false, false}() + result = EnzymeCore.EnzymeRules.forward( + config_primal, Const(fww), EnzymeCore.BatchDuplicated{Float64, 2}, + BatchDuplicated(3.0, (1.0, 2.0)) + ) + @test result isa BatchDuplicated + @test result.val ≈ 9.0 + @test result.dval isa NTuple{2, Float64} + @test !(result.dval isa NamedTuple) + @test result.dval == (6.0, 12.0) + + # Confirm the conversion generalises to W > 2. + config_w3 = EnzymeCore.EnzymeRules.FwdConfig{false, true, 3, false, false}() + shadow3 = EnzymeCore.EnzymeRules.forward( + config_w3, Const(fww), EnzymeCore.BatchDuplicated{Float64, 3}, + BatchDuplicated(3.0, (1.0, 2.0, 4.0)) + ) + @test shadow3 isa NTuple{3, Float64} + @test shadow3 == (6.0, 12.0, 24.0) +end + +@testset "Enzyme forward mode, Const return (IIP, no return-shadow)" begin + # Covers EnzymeRules.FwdConfig{false, false, W, ...} — Enzyme dispatches on + # this combo for IIP functions with a Const return type where the caller + # needs primal + shadow propagation via Duplicated args only (no return + # value to shadow). Reproduces the SciML/OrdinaryDiffEq.jl v7 Downstream + # regression where this call previously produced: + # - without any rule: MethodError: no method matching forward(…) + # - with a primal-only rule: trivial (zero) arg shadows, wrong Jacobians + # (Rodas4/5/Veldd4 errors 4–9 orders of magnitude above tolerance). + # The rule must delegate to `Enzyme.autodiff` on the unwrapped function + # so Duplicated arg shadows propagate correctly. + f!(du, u) = (du[1] = -u[1]^2; nothing) + fww = FunctionWrappersWrapper( + f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,) + ) + + du = [0.0] + u = [3.0] + du_shadow = [0.0] + u_shadow = [1.0] # seed: ∂/∂u[1] = 1 + + config = EnzymeCore.EnzymeRules.FwdConfig{false, false, 1, false, false}() + ret = EnzymeCore.EnzymeRules.forward( + config, Const(fww), EnzymeCore.Const{Nothing}, + Duplicated(du, du_shadow), Duplicated(u, u_shadow) + ) + @test ret === nothing + # Primal side-effect: du[1] = -u[1]^2 = -9 + @test du[1] ≈ -9.0 + # Shadow propagation: ∂du[1]/∂u[1] * u_shadow[1] = -2*u[1]*1 = -6 + @test du_shadow[1] ≈ -6.0 +end + +@testset "Enzyme reverse mode, Const return — augmented_primal runs primal" begin + # Mirrors the forward {false, false} case on the reverse side. Augmented + # primal runs the wrapped function for its side effects and returns + # AugmentedReturn(nothing, nothing, nothing). Reverse returns `nothing` + # per arg since there is no return derivative to propagate. + counter = Ref(0) + g(x, y) = (counter[] += 1; x + y) # returns Float64 (ignored via Const RT) + fww = FunctionWrappersWrapper(g, (Tuple{Float64, Float64},), (Float64,)) + + # Construct a concrete RevConfig. Fields: + # (NeedsPrimal, NeedsShadow, Width, Overwritten, RuntimeActivity, StrongZero) + rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false), false, false}() + + counter[] = 0 + aug = EnzymeRules.augmented_primal( + rconfig, Const(fww), EnzymeCore.Const{Float64}, + Active(3.0), Active(4.0) + ) + @test counter[] == 1 # primal ran exactly once + @test aug.primal === nothing # NeedsPrimal == false + @test aug.shadow === nothing + @test aug.tape === nothing + + # Reverse step — dret is Const (passed as TYPE not instance in reverse + # rules). Enzyme's rule protocol requires concrete gradients for Active + # args; under a Const return they're zero (no gradient source). + grads = EnzymeRules.reverse( + rconfig, Const(fww), EnzymeCore.Const{Float64}, + aug.tape, Active(3.0), Active(4.0) + ) + @test grads == (0.0, 0.0) +end + +@testset "Enzyme reverse mode, Duplicated return — augmented_primal initializes shadow" begin + # Covers augmented_primal with RT <: Duplicated{T}. + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + rconfig = EnzymeRules.RevConfig{true, true, 1, (false,), false, false}() + + aug = EnzymeRules.augmented_primal( + rconfig, Const(fww), EnzymeCore.Duplicated{Float64}, + Active(3.0) + ) + @test aug.primal ≈ 9.0 # f(3) = 9 + @test aug.shadow ≈ 0.0 # zero-initialized shadow + @test aug.tape === nothing +end + +@testset "Enzyme reverse mode, BatchDuplicated return — augmented_primal initializes shadows" begin + # Covers augmented_primal with RT <: BatchDuplicated{T, W}. + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + rconfig = EnzymeRules.RevConfig{true, true, 2, (false,), false, false}() + + aug = EnzymeRules.augmented_primal( + rconfig, Const(fww), EnzymeCore.BatchDuplicated{Float64, 2}, + Active(3.0) + ) + @test aug.primal ≈ 9.0 + @test aug.shadow isa NTuple{2, Float64} + @test aug.shadow == (0.0, 0.0) + @test aug.tape === nothing +end + +# ============================================================================= +# End-to-end reverse-mode derivative tests — exercise Enzyme.autodiff(Reverse, +# …) through the FWW and assert the resulting gradients are numerically correct. +# The prior reverse-mode testsets only checked dispatch / shape of +# AugmentedReturn; they did NOT verify the gradients are right. +# ============================================================================= + +@testset "Enzyme Reverse: Const return, Active args — no-flow gradients" begin + # For a function whose return is annotated Const in Reverse mode, there is + # no gradient source from the return, so Active arg gradients must be 0. + # (Enzyme's rule-return protocol requires concrete gradients for Active + # args — `nothing` is not allowed — so the rule returns zeros.) + g(x, y) = x * y + x^2 + fww = FunctionWrappersWrapper(g, (Tuple{Float64, Float64},), (Float64,)) + + # Const return (instead of Active) → no gradient flows back + result = Enzyme.autodiff(Reverse, Const(fww), Const, Active(3.0), Active(4.0)) + @test result[1] === (0.0, 0.0) +end + +@testset "Enzyme Reverse: IIP with Duplicated args, Const return" begin + # SciML's standard pattern: IIP RHS `f!(du, u)` with Const return, both du + # and u are Duplicated. Reverse mode should accumulate + # u_shadow[i] += du_shadow[j] * ∂(du[j])/∂(u[i]) + # into u_shadow. For f!(du, u) = (du[1] = u[1]^2; nothing) with + # du_shadow = [1.0] (incoming adjoint), + # u[1] = 3.0, + # ∂du[1]/∂u[1] = 2*u[1] = 6, + # the expected result is u_shadow[1] = 6.0 after the call. + f!(du, u) = (du[1] = u[1]^2; nothing) + fww = FunctionWrappersWrapper( + f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,) + ) + + du = [0.0] + u = [3.0] + du_shadow = [1.0] + u_shadow = [0.0] + + Enzyme.autodiff( + Reverse, Const(fww), Const, + Duplicated(du, du_shadow), Duplicated(u, u_shadow) + ) + @test du[1] ≈ 9.0 # primal effect: du[1] = u[1]^2 = 9 + @test u_shadow[1] ≈ 6.0 # reverse accumulation: 2 * u[1] * du_shadow[1] +end + +@testset "Enzyme Reverse: IIP multi-component IIP with Duplicated args" begin + # Cross-coupled IIP RHS: each output depends on multiple inputs. + # du[1] = u[1] * u[2] + # du[2] = u[1]^2 + u[2]^3 + # Jacobian at u = (x, y): + # J = [ y x ; + # 2x 3y^2 ] + # In reverse mode with du_shadow = [a, b], transpose of J applied to + # du_shadow gives the accumulation into u_shadow: + # u_shadow[1] += a*y + b*2x + # u_shadow[2] += a*x + b*3y^2 + f!(du, u) = (du[1] = u[1] * u[2]; du[2] = u[1]^2 + u[2]^3; nothing) + fww = FunctionWrappersWrapper( + f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,) + ) + + x, y = 2.0, 5.0 + a, b = 1.0, 0.5 + du = zeros(2) + u = [x, y] + du_shadow = [a, b] + u_shadow = zeros(2) + + Enzyme.autodiff( + Reverse, Const(fww), Const, + Duplicated(du, du_shadow), Duplicated(u, u_shadow) + ) + @test du ≈ [x * y, x^2 + y^3] + @test u_shadow[1] ≈ a * y + b * 2 * x # 5 + 2 = 7 + @test u_shadow[2] ≈ a * x + b * 3 * y^2 # 2 + 37.5 = 39.5 +end + +@testset "Enzyme ReverseWithPrimal: IIP with Duplicated args" begin + # Same IIP pattern but with ReverseWithPrimal so we also check the primal + # is available when the rule is asked to include it. + f!(du, u) = (du[1] = u[1]^3; nothing) + fww = FunctionWrappersWrapper( + f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,) + ) + + du = [0.0] + u = [2.0] + du_shadow = [1.0] + u_shadow = [0.0] + + # Capture the expected gradient BEFORE the call — Enzyme may zero + # `du_shadow` after consuming it during the reverse pass. + expected_u_grad = 3 * u[1]^2 * du_shadow[1] # = 12.0 + + Enzyme.autodiff( + ReverseWithPrimal, Const(fww), Const, + Duplicated(du, du_shadow), Duplicated(u, u_shadow) + ) + @test du[1] ≈ 8.0 + @test u_shadow[1] ≈ expected_u_grad +end + +# ============================================================================= +# Runtime-activity propagation through the FWW forward rules. +# +# Prior to this fix the rules hard-coded plain `Forward` when delegating to +# `Enzyme.autodiff`, silently dropping the caller's +# `set_runtime_activity(Forward)` flag. Enzyme's static IR-level activity +# analysis can't see through `FunctionWrappersWrapper`'s opaque cfunction +# indirection, so the inner call raised `EnzymeRuntimeActivityError` inside +# `@.` broadcast's `broadcast_unalias` → `mightalias` — despite +# `set_runtime_activity` being set on the outer `autodiff` call. +# +# Upstream motivation: OrdinaryDiffEq.jl PR #3518 — +# Rosenbrock23(autodiff = AutoEnzyme(set_runtime_activity(Enzyme.Forward))) +# on any time-dependent in-place RHS routed through DiffEqBase's +# `wrapfun_iip`. Here we reproduce the shape (`f!(du, u, p, t) = @. du = …`) +# in a 4-arg `FunctionWrappersWrapper` matching DiffEqBase's +# `wrapfun_iip` output, and assert both that (a) the call completes without +# an `EnzymeRuntimeActivityError` and (b) the resulting tangent is +# numerically correct. +# ============================================================================= + +# ============================================================================= +# Duplicated function annotation on the FWW itself. +# +# Reproduces SciML/FunctionWrappersWrappers.jl#48: when Enzyme differentiates +# through a closure that captures an FWW (e.g. NonlinearSolve + +# SciMLSensitivity), the rule is invoked with +# `Duplicated{<:FunctionWrappersWrapper}` for the function argument, not +# `Const{<:FunctionWrappersWrapper}`. The FWW struct itself only carries +# `FunctionWrapper`s + cache storage, so its "shadow" is ignored — we route +# through `unwrap(func.val)` exactly as with `Const`. +# ============================================================================= + +@testset "Enzyme forward, Duplicated FWW annotation — IIP Const return" begin + f!(residual, u, p) = (residual[1] = u[1]^2 - p[1]; nothing) + fww = FunctionWrappersWrapper( + f!, + (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}},), + (Nothing,) + ) + + residual = [0.0]; dresidual = [0.0] + u = [2.0]; du = [1.0] + p = [1.0]; dp = [0.0] + + config = EnzymeCore.EnzymeRules.FwdConfig{false, false, 1, false, false}() + ret = EnzymeCore.EnzymeRules.forward( + config, + Duplicated(fww, fww), # <-- the failing dispatch in #48 + EnzymeCore.Const{Nothing}, + Duplicated(residual, dresidual), + Duplicated(u, du), + Duplicated(p, dp), + ) + @test ret === nothing + @test residual[1] ≈ 3.0 # u[1]^2 - p[1] = 4 - 1 + @test dresidual[1] ≈ 4.0 # 2*u[1]*du[1] - 1*dp[1] = 4 +end + +@testset "Enzyme forward, Duplicated FWW annotation — shadow-only return" begin + # Drive the {false, true, W, …} rule (shadow only, no primal) with a + # Duplicated FWW. + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + + config = EnzymeCore.EnzymeRules.FwdConfig{false, true, 1, false, false}() + shadow = EnzymeCore.EnzymeRules.forward( + config, + Duplicated(fww, fww), + EnzymeCore.Duplicated{Float64}, + Duplicated(3.0, 1.0), + ) + @test shadow ≈ 6.0 # f'(3) = 2*3 = 6 +end + +@testset "Enzyme forward, Duplicated FWW annotation — primal + shadow return" begin + # Drive the {true, true, W, …} rule (ForwardWithPrimal) with a Duplicated + # FWW. + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + + config = EnzymeCore.EnzymeRules.FwdConfig{true, true, 1, false, false}() + result = EnzymeCore.EnzymeRules.forward( + config, + Duplicated(fww, fww), + EnzymeCore.Duplicated{Float64}, + Duplicated(3.0, 1.0), + ) + @test result isa Duplicated + @test result.val ≈ 9.0 # primal + @test result.dval ≈ 6.0 # shadow +end + +@testset "Enzyme reverse, Duplicated FWW annotation — Const return IIP" begin + # Mirror the forward IIP case on the reverse side. Duplicated FWW must + # still drive the rule, gradients must accumulate into u_shadow. + f!(du, u) = (du[1] = u[1]^2; nothing) + fww = FunctionWrappersWrapper( + f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,) + ) + + du = [0.0]; du_shadow = [1.0] + u = [3.0]; u_shadow = [0.0] + + rconfig = EnzymeRules.RevConfig{false, false, 1, (false, false), false, false}() + aug = EnzymeRules.augmented_primal( + rconfig, + Duplicated(fww, fww), # <-- Duplicated FWW + EnzymeCore.Const{Nothing}, + Duplicated(du, du_shadow), + Duplicated(u, u_shadow), + ) + @test aug.primal === nothing + @test aug.shadow === nothing + + EnzymeRules.reverse( + rconfig, + Duplicated(fww, fww), + EnzymeCore.Const{Nothing}, + aug.tape, + Duplicated(du, du_shadow), + Duplicated(u, u_shadow), + ) + @test du[1] ≈ 9.0 # primal effect from augmented_primal + @test u_shadow[1] ≈ 6.0 # reverse accumulation: 2*u[1]*du_shadow[1] +end + +@testset "Enzyme Forward: set_runtime_activity propagates through FWW (IIP, time-dependent)" begin + # DiffEqBase's `wrapfun_iip(ff, (u, u, p, t))` shape. + const_INPUTS = Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64} + + # 1) Time-independent RHS — ∂du/∂t = 0. + f!(du, u, p, t) = (@. du = p * u; nothing) + fww = FunctionWrappersWrapper(f!, (const_INPUTS,), (Nothing,)) + + u = [1.0, 2.0, 3.0] + p = [0.5, 0.5, 0.5] + t = 1.7 + du = zero(u); ddu = zero(u); dt = 1.0 + + Enzyme.autodiff( + Enzyme.set_runtime_activity(Forward), + Const(fww), Const, + Duplicated(du, ddu), + Const(u), Const(p), + Duplicated(t, dt), + ) + @test du ≈ p .* u + @test all(iszero, ddu) + + # 2) Non-trivial time dependence: g!(du,u,p,t) = @. sin(t)*u. + # Expected ∂du/∂t = cos(t) .* u. + g!(du, u, p, t) = (@. du = sin(t) * u; nothing) + gww = FunctionWrappersWrapper(g!, (const_INPUTS,), (Nothing,)) + + du2 = zero(u); ddu2 = zero(u) + Enzyme.autodiff( + Enzyme.set_runtime_activity(Forward), + Const(gww), Const, + Duplicated(du2, ddu2), + Const(u), Const(p), + Duplicated(t, 1.0), + ) + @test du2 ≈ sin(t) .* u + @test ddu2 ≈ cos(t) .* u + + # 3) Confirm the rule also propagates `set_strong_zero(Forward)` (the + # other ForwardMode flag carried in FwdConfig) — another RHS that + # doesn't need runtime activity but exercises a distinct flag. + h!(du, u, p, t) = (du[1] = u[1] * t; nothing) + hww = FunctionWrappersWrapper( + h!, (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},), + (Nothing,) + ) + du_h = [0.0]; ddu_h = [0.0] + Enzyme.autodiff( + Enzyme.set_strong_zero(Forward), + Const(hww), Const, + Duplicated(du_h, ddu_h), + Const([2.0]), Const([0.0]), + Duplicated(3.5, 1.0), + ) + @test du_h[1] ≈ 2.0 * 3.5 # primal: u[1] * t = 7.0 + @test ddu_h[1] ≈ 2.0 # ∂(u[1]*t)/∂t = u[1] +end diff --git a/test/Mooncake/Project.toml b/test/Mooncake/Project.toml new file mode 100644 index 0000000..ac5b7be --- /dev/null +++ b/test/Mooncake/Project.toml @@ -0,0 +1,13 @@ +[deps] +FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +FunctionWrappersWrappers = {path = "../.."} + +[compat] +Mooncake = "0.5.36" +SafeTestsets = "0.0.1, 0.1, 1" +Test = "1" diff --git a/test/Mooncake/mooncake_tests.jl b/test/Mooncake/mooncake_tests.jl new file mode 100644 index 0000000..8ca86fd --- /dev/null +++ b/test/Mooncake/mooncake_tests.jl @@ -0,0 +1,118 @@ +using FunctionWrappersWrappers +using Mooncake +using Test + +@testset "Mooncake reverse mode - single arg" begin + f(x) = x^2 + fww = FunctionWrappersWrapper(f, (Tuple{Float64},), (Float64,)) + + rule = Mooncake.build_rrule(fww, 3.0) + val, (_, dx) = Mooncake.value_and_gradient!!(rule, fww, 3.0) + @test val ≈ 9.0 + @test dx ≈ 6.0 +end + +@testset "Mooncake reverse mode - multi arg" begin + g(x, y) = x * y + x^2 + fww = FunctionWrappersWrapper(g, (Tuple{Float64, Float64},), (Float64,)) + + # g(x,y) = x*y + x^2 → ∂g/∂x = y + 2x, ∂g/∂y = x + rule = Mooncake.build_rrule(fww, 3.0, 4.0) + val, (_, dx, dy) = Mooncake.value_and_gradient!!(rule, fww, 3.0, 4.0) + @test val ≈ 21.0 + @test dx ≈ 10.0 # ∂g/∂x at (3,4) = 4 + 6 + @test dy ≈ 3.0 # ∂g/∂y at (3,4) = 3 +end + +@testset "Mooncake with trig functions" begin + fww_sin = FunctionWrappersWrapper(sin, (Tuple{Float64},), (Float64,)) + + rule = Mooncake.build_rrule(fww_sin, 1.0) + val, (_, dx) = Mooncake.value_and_gradient!!(rule, fww_sin, 1.0) + @test val ≈ sin(1.0) + @test dx ≈ cos(1.0) +end + +@testset "Mooncake through loss function" begin + # Test that Mooncake can differentiate a loss function that calls FunctionWrappersWrapper + f(x) = x[1]^2 + x[2]^2 + fww = FunctionWrappersWrapper(f, (Tuple{Vector{Float64}},), (Float64,)) + + loss(x) = fww(x) + rule = Mooncake.build_rrule(loss, [3.0, 4.0]) + val, (_, dx) = Mooncake.value_and_gradient!!(rule, loss, [3.0, 4.0]) + @test val ≈ 25.0 + @test dx ≈ [6.0, 8.0] +end + +@testset "Mooncake with wrapped callable struct" begin + # SciML wraps functions in Void{F} or similar callable structs before + # putting them in FunctionWrappersWrapper. The unwrapped function is + # then a non-primitive callable struct, not a plain function. + struct VoidWrapper{F} + f::F + end + function (v::VoidWrapper)(args...) + v.f(args...) + return nothing + end + + function f!(du, u, p) + du[1] = p[1] * u[1] + du[2] = p[2] * u[2] + return nothing + end + + wrapped = VoidWrapper(f!) + fww = FunctionWrappersWrapper( + wrapped, + (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}},), + (Nothing,), + ) + + function loss(p) + u = [3.0, 4.0] + du = similar(u) + fww(du, u, p) + return sum(abs2, du) + end + + rule = Mooncake.build_rrule(loss, [2.0, 3.0]) + val, (_, dp) = Mooncake.value_and_gradient!!(rule, loss, [2.0, 3.0]) + # du = [2*3, 3*4] = [6, 12], loss = 36 + 144 = 180 + @test val ≈ 180.0 + # ∂loss/∂p1 = 2*du[1]*u[1] = 2*6*3 = 36 + # ∂loss/∂p2 = 2*du[2]*u[2] = 2*12*4 = 96 + @test dp ≈ [36.0, 96.0] +end + +@testset "Mooncake in-place function" begin + # In-place functions are common in SciML (f!(du, u, p, t)) + function f!(du, u, p) + du[1] = p[1] * u[1] + p[2] * u[2] + du[2] = p[3] * u[1] - u[2] + return nothing + end + fww = FunctionWrappersWrapper( + f!, + (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}},), + (Nothing,), + ) + + function loss(p) + u = [1.0, 2.0] + du = similar(u) + fww(du, u, p) + return sum(abs2, du) + end + + rule = Mooncake.build_rrule(loss, [1.0, 2.0, 3.0]) + val, (_, dp) = Mooncake.value_and_gradient!!(rule, loss, [1.0, 2.0, 3.0]) + # f!(du, [1,2], [1,2,3]) → du = [1*1+2*2, 3*1-2] = [5, 1] + # loss = 25 + 1 = 26 + @test val ≈ 26.0 + # ∂loss/∂p1 = 2*du[1]*u[1] = 2*5*1 = 10 + # ∂loss/∂p2 = 2*du[1]*u[2] = 2*5*2 = 20 + # ∂loss/∂p3 = 2*du[2]*u[1] = 2*1*1 = 2 + @test dp ≈ [10.0, 20.0, 2.0] +end diff --git a/test/basictests.jl b/test/basictests.jl new file mode 100644 index 0000000..a469614 --- /dev/null +++ b/test/basictests.jl @@ -0,0 +1,264 @@ +using FunctionWrappersWrappers +using Test + +@testset "FunctionWrappersWrappers.jl" begin + fwplus = FunctionWrappersWrapper( + +, (Tuple{Float64, Float64}, Tuple{Int, Int}), ( + Float64, Int, + ) + ) + @test fwplus(4.0, 8.0) === 12.0 + @test fwplus(4, 8) === 12 + + fwexp2 = FunctionWrappersWrapper( + exp2, (Tuple{Float64}, Tuple{Float32}, Tuple{Int}), (Float64, Float32, Float64) + ) + @test fwexp2(4.0) === 16.0 + @test fwexp2(4.0f0) === 16.0f0 + @test fwexp2(4) === 16.0 +end + +@testset "Type inference" begin + fwplus = FunctionWrappersWrapper( + +, (Tuple{Float64, Float64}, Tuple{Int, Int}), ( + Float64, Int, + ) + ) + @test @inferred(fwplus(4.0, 8.0)) === 12.0 + @test @inferred(fwplus(4, 8)) === 12 + + fwexp2 = FunctionWrappersWrapper( + exp2, (Tuple{Float64}, Tuple{Float32}, Tuple{Int}), (Float64, Float32, Float64) + ) + @test @inferred(fwexp2(4.0)) === 16.0 + @test @inferred(fwexp2(4.0f0)) === 16.0f0 + @test @inferred(fwexp2(4)) === 16.0 +end + +@testset "Introspection functions" begin + fwsin = FunctionWrappersWrapper(sin, (Tuple{Float64},), (Float64,)) + + @testset "unwrap" begin + f = unwrap(fwsin) + @test f === sin + @test f(0.5) == sin(0.5) + end + + @testset "wrapped_signatures" begin + sigs = wrapped_signatures(fwsin) + @test sigs == (Tuple{Float64},) + end + + @testset "wrapped_return_types" begin + rets = wrapped_return_types(fwsin) + @test rets == (Float64,) + end + + fwplus = FunctionWrappersWrapper( + +, (Tuple{Float64, Float64}, Tuple{Int, Int}), ( + Float64, Int, + ) + ) + + @testset "unwrap with multiple signatures" begin + f = unwrap(fwplus) + @test f === + + @test f(1, 2) == 3 + end + + @testset "wrapped_signatures with multiple signatures" begin + sigs = wrapped_signatures(fwplus) + @test sigs == (Tuple{Float64, Float64}, Tuple{Int, Int}) + end + + @testset "wrapped_return_types with multiple signatures" begin + rets = wrapped_return_types(fwplus) + @test rets == (Float64, Int) + end + + my_func(x) = x^2 + fwcustom = FunctionWrappersWrapper( + my_func, (Tuple{Float64}, Tuple{Int}), ( + Float64, Int, + ) + ) + + @testset "unwrap with custom function" begin + f = unwrap(fwcustom) + @test f === my_func + @test f(3) == 9 + @test f(2.5) == 6.25 + end +end + +@testset "Fallback policies" begin + @testset "Strict" begin + fww = FunctionWrappersWrapper( + +, (Tuple{Float64, Float64},), (Float64,); + cache = NoCache(), policy = Strict() + ) + @test fww(4.0, 8.0) === 12.0 + @test_throws FunctionWrappersWrappers.NoFunctionWrapperFoundError fww(4, 8) + @test_throws FunctionWrappersWrappers.NoFunctionWrapperFoundError fww( + BigFloat(4), BigFloat(8) + ) + end + + @testset "AllowAll" begin + fww = FunctionWrappersWrapper( + +, (Tuple{Float64, Float64},), (Float64,); + cache = NoCache(), policy = AllowAll() + ) + @test fww(4.0, 8.0) === 12.0 + @test fww(4, 8) === 12 + @test fww(4.0f0, 8.0f0) == 12.0f0 + @test fww(BigFloat(4), BigFloat(8)) == BigFloat(12) + end + + @testset "AllowNonIsBits" begin + fww = FunctionWrappersWrapper( + +, (Tuple{Float64, Float64},), (Float64,); + cache = NoCache(), policy = AllowNonIsBits() + ) + @test fww(4.0, 8.0) === 12.0 + # Float32 is isbits but doesn't match Float64 wrapper → error + @test_throws FunctionWrappersWrappers.NoFunctionWrapperFoundError fww(4.0f0, 8.0f0) + # Int is isbits but doesn't match Float64 wrapper → error + @test_throws FunctionWrappersWrappers.NoFunctionWrapperFoundError fww(4, 8) + # BigFloat is non-isbits → allowed + @test fww(BigFloat(4), BigFloat(8)) == BigFloat(12) + end + + @testset "AllowNonIsBits with arrays" begin + f!(du, u) = (du[1] = u[1]^2; nothing) + fww = FunctionWrappersWrapper( + f!, (Tuple{Vector{Float64}, Vector{Float64}},), (Nothing,); + cache = NoCache(), policy = AllowNonIsBits() + ) + du_f = [0.0]; u_f = [3.0] + fww(du_f, u_f) + @test du_f[1] === 9.0 + + # Float32 arrays: eltype is isbits but doesn't match → error + @test_throws FunctionWrappersWrappers.NoFunctionWrapperFoundError fww( + Float32[0.0], Float32[3.0] + ) + + # BigFloat arrays: eltype is non-isbits → allowed + du_bf = BigFloat[0]; u_bf = BigFloat[3] + fww(du_bf, u_bf) + @test du_bf[1] == BigFloat(9) + end +end + +@testset "Cache modes" begin + f!(du, u, p, t) = (du[1] = p[1] * u[1]; nothing) + + @testset "NoCache" begin + fww = FunctionWrappersWrapper( + f!, + (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},), + (Nothing,); + cache = NoCache(), policy = AllowAll() + ) + # Float64 match + du = [0.0]; u = [2.0]; p = [3.0] + fww(du, u, p, 0.0) + @test du[1] === 6.0 + + # BigFloat fallback (NoCache: 1 alloc per call) + du_bf = BigFloat[0]; u_bf = BigFloat[2]; p_bf = BigFloat[3]; t_bf = BigFloat(0) + fww(du_bf, u_bf, p_bf, t_bf) + @test du_bf[1] == BigFloat(6) + end + + @testset "SingleCache" begin + fww = FunctionWrappersWrapper( + f!, + (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},), + (Nothing,); + cache = SingleCache(), policy = AllowAll() + ) + du_bf = BigFloat[0]; u_bf = BigFloat[2]; p_bf = BigFloat[3]; t_bf = BigFloat(0) + # First call caches + fww(du_bf, u_bf, p_bf, t_bf) + @test du_bf[1] == BigFloat(6) + # Second call uses cache (0 alloc) + du_bf[1] = BigFloat(0) + fww(du_bf, u_bf, p_bf, t_bf) + @test du_bf[1] == BigFloat(6) + end + + @testset "DictCache" begin + fww = FunctionWrappersWrapper( + f!, + (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},), + (Nothing,); + cache = DictCache(), policy = AllowAll() + ) + du_bf = BigFloat[0]; u_bf = BigFloat[2]; p_bf = BigFloat[3]; t_bf = BigFloat(0) + fww(du_bf, u_bf, p_bf, t_bf) + @test du_bf[1] == BigFloat(6) + + # Different type also works and caches separately + du_f32 = Float32[0]; u_f32 = Float32[2]; p_f32 = Float32[3]; t_f32 = Float32(0) + fww(du_f32, u_f32, p_f32, t_f32) + @test du_f32[1] === Float32(6) + + # BigFloat still cached + du_bf[1] = BigFloat(0) + fww(du_bf, u_bf, p_bf, t_bf) + @test du_bf[1] == BigFloat(6) + end + + @testset "SingleCache thrashing recovers" begin + fww = FunctionWrappersWrapper( + f!, + (Tuple{Vector{Float64}, Vector{Float64}, Vector{Float64}, Float64},), + (Nothing,); + cache = SingleCache(), policy = AllowAll() + ) + du_bf = BigFloat[0]; u_bf = BigFloat[2]; p_bf = BigFloat[3]; t_bf = BigFloat(0) + du_f32 = Float32[0]; u_f32 = Float32[2]; p_f32 = Float32[3]; t_f32 = Float32(0) + + # Alternate types — each call replaces the cache but still works + fww(du_bf, u_bf, p_bf, t_bf) + @test du_bf[1] == BigFloat(6) + fww(du_f32, u_f32, p_f32, t_f32) + @test du_f32[1] === Float32(6) + du_bf[1] = BigFloat(0) + fww(du_bf, u_bf, p_bf, t_bf) + @test du_bf[1] == BigFloat(6) + end +end + +@testset "Default constructor uses SingleCache + AllowNonIsBits" begin + fww = FunctionWrappersWrapper( + +, (Tuple{Float64, Float64},), (Float64,) + ) + # Float64 matches wrapper + @test fww(4.0, 8.0) === 12.0 + # BigFloat is non-isbits → falls back + @test fww(BigFloat(4), BigFloat(8)) == BigFloat(12) + # Float32 is isbits mismatch → errors + @test_throws FunctionWrappersWrappers.NoFunctionWrapperFoundError fww(4.0f0, 8.0f0) +end + +@testset "Conversion" begin + fww_exp = FunctionWrappersWrapper( + exp, + (Tuple{Float64}, Tuple{Float32}), + (Float64, Float32) + ) + FWW = typeof(fww_exp) + + fww_cos = FWW(cos) + @test typeof(fww_cos) == FWW + @test fww_cos(0.5) == cos(0.5) + + fww_vector = FWW[sin, tan, log] + @test eltype(fww_vector) == FWW + fww_vector[1](0.5) == sin(0.5) + fww_vector[2](0.5) == tan(0.5) + fww_vector[3](0.5) == log(0.5) +end diff --git a/test/nopre/Project.toml b/test/nopre/Project.toml new file mode 100644 index 0000000..2021cec --- /dev/null +++ b/test/nopre/Project.toml @@ -0,0 +1,14 @@ +[deps] +FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf" +JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" +SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +FunctionWrappersWrappers = {path = "../.."} + +[compat] +SafeTestsets = "0.0.1, 0.1" +SciMLTesting = "1" +Test = "1" diff --git a/test/nopre/jet_tests.jl b/test/nopre/jet_tests.jl new file mode 100644 index 0000000..42a3ee7 --- /dev/null +++ b/test/nopre/jet_tests.jl @@ -0,0 +1,30 @@ +using FunctionWrappersWrappers +using JET: JET, @test_opt +using Test + +@testset "JET static analysis" begin + # Test that the main call path is type-stable (no fallback) + fwplus = FunctionWrappersWrapper( + +, (Tuple{Float64, Float64}, Tuple{Int, Int}), ( + Float64, Int, + ) + ) + + # Core functionality should be type-stable + @test_opt target_modules = (FunctionWrappersWrappers,) fwplus(4.0, 8.0) + @test_opt target_modules = (FunctionWrappersWrappers,) fwplus(4, 8) + + # Test single-argument wrapper + fwexp2 = FunctionWrappersWrapper( + exp2, (Tuple{Float64}, Tuple{Float32}, Tuple{Int}), (Float64, Float32, Float64) + ) + @test_opt target_modules = (FunctionWrappersWrappers,) fwexp2(4.0) + @test_opt target_modules = (FunctionWrappersWrappers,) fwexp2(4.0f0) + @test_opt target_modules = (FunctionWrappersWrappers,) fwexp2(4) + + # Verify no errors detected by JET.report_call for core paths + rep = JET.report_call(fwplus, (Float64, Float64)) + @test length(JET.get_reports(rep)) == 0 + rep = JET.report_call(fwplus, (Int, Int)) + @test length(JET.get_reports(rep)) == 0 +end diff --git a/test/qa/Project.toml b/test/qa/Project.toml new file mode 100644 index 0000000..754f8b8 --- /dev/null +++ b/test/qa/Project.toml @@ -0,0 +1,18 @@ +[deps] +Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +FunctionWrappersWrappers = "77dc65aa-8811-40c2-897b-53d922fa7daf" +JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" +SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" +SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[sources] +FunctionWrappersWrappers = {path = "../.."} + +[compat] +Aqua = "0.8" +JET = "0.9, 0.10, 0.11" +SafeTestsets = "0.0.1, 0.1, 1" +SciMLTesting = "1.6" +Test = "1" +julia = "1.10" diff --git a/test/qa/qa.jl b/test/qa/qa.jl new file mode 100644 index 0000000..0cbadbc --- /dev/null +++ b/test/qa/qa.jl @@ -0,0 +1,17 @@ +using SciMLTesting, FunctionWrappersWrappers, Test +using JET + +run_qa( + FunctionWrappersWrappers; + explicit_imports = true, + ei_kwargs = (; + # Qualified accesses to names that are not (yet) declared `public` in their + # owning module: `Base.tail`, `FunctionWrappers.FunctionWrapper`, and + # `TruncatedStacktraces.@truncate_stacktrace`. These are core, long-standing + # API of Base / FunctionWrappers / TruncatedStacktraces that predate the + # `public` keyword; ignore until those packages mark them public. + all_qualified_accesses_are_public = (; + ignore = (:tail, :FunctionWrapper, Symbol("@truncate_stacktrace")), + ), + ), +) diff --git a/test/runtests.jl b/test/runtests.jl index a66c601..e4da1f8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,15 +1,36 @@ -using FunctionWrappersWrappers -using Test +using SafeTestsets +using SciMLTesting -@testset "FunctionWrappersWrappers.jl" begin - fwplus = FunctionWrappersWrapper(+, (Tuple{Float64, Float64}, Tuple{Int, Int}), ( - Float64, Int)); - @test fwplus(4.0, 8.0) === 12.0 - @test fwplus(4, 8) === 12 - - fwexp2 = FunctionWrappersWrapper( - exp2, (Tuple{Float64}, Tuple{Float32}, Tuple{Int}), (Float64, Float32, Float64)); - @test fwexp2(4.0) === 16.0 - @test fwexp2(4.0f0) === 16.0f0 - @test fwexp2(4) === 16.0 -end +run_tests(; + core = function () + @safetestset "FunctionWrappersWrappers.jl" begin + include(joinpath(@__DIR__, "basictests.jl")) + end + return @safetestset "BigFloat + UnionAll" begin + include(joinpath(@__DIR__, "shared", "bigfloat_unionall_tests.jl")) + end + end, + qa = (; env = joinpath(@__DIR__, "qa"), body = joinpath(@__DIR__, "qa", "qa.jl")), + groups = Dict( + "nopre" => (; + env = joinpath(@__DIR__, "nopre"), + body = function () + @safetestset "JET" begin + include(joinpath(@__DIR__, "nopre", "jet_tests.jl")) + end + return @safetestset "BigFloat + UnionAll" begin + include(joinpath(@__DIR__, "shared", "bigfloat_unionall_tests.jl")) + end + end, + ), + "Enzyme" => (; + env = joinpath(@__DIR__, "Enzyme"), + body = joinpath(@__DIR__, "Enzyme", "enzyme_tests.jl"), + ), + "Mooncake" => (; + env = joinpath(@__DIR__, "Mooncake"), + body = joinpath(@__DIR__, "Mooncake", "mooncake_tests.jl"), + ), + ), + all = ["Core", "Enzyme", "Mooncake"], +) diff --git a/test/shared/bigfloat_unionall_tests.jl b/test/shared/bigfloat_unionall_tests.jl new file mode 100644 index 0000000..de091c5 --- /dev/null +++ b/test/shared/bigfloat_unionall_tests.jl @@ -0,0 +1,41 @@ +using FunctionWrappersWrappers +using Test + +@testset "BigFloat support" begin + fwplus_big = FunctionWrappersWrapper( + +, + (Tuple{BigFloat, BigFloat}, Tuple{Float64, Float64}), + (BigFloat, Float64) + ) + a = BigFloat("3.14159265358979323846264338327950288") + b = BigFloat("2.71828182845904523536028747135266250") + @test fwplus_big(a, b) isa BigFloat + @test fwplus_big(a, b) == a + b + @test fwplus_big(1.0, 2.0) === 3.0 + + fwsin_big = FunctionWrappersWrapper( + sin, + (Tuple{BigFloat}, Tuple{Float64}), + (BigFloat, Float64) + ) + @test fwsin_big(BigFloat("1.0")) isa BigFloat + @test fwsin_big(1.0) === sin(1.0) +end + +@testset "UnionAll return types" begin + # Test that UnionAll types (like AbstractArray{Float64}) work as return types + function double_array(x::AbstractArray{Float64}) + return x .* 2 + end + + fwdouble = FunctionWrappersWrapper( + double_array, + (Tuple{AbstractArray{Float64}},), + (AbstractArray{Float64},) + ) + + v = [1.0, 2.0, 3.0] + result = fwdouble(v) + @test result isa Vector{Float64} + @test result == [2.0, 4.0, 6.0] +end diff --git a/test/test_groups.toml b/test/test_groups.toml new file mode 100644 index 0000000..ab28124 --- /dev/null +++ b/test/test_groups.toml @@ -0,0 +1,14 @@ +[Core] +versions = ["lts", "1", "pre"] + +[QA] +versions = ["lts", "1"] + +[nopre] +versions = ["lts", "1"] + +[Enzyme] +versions = ["lts", "1"] + +[Mooncake] +versions = ["lts", "1"]