Skip to content

feat(DI): add Piste backend #1089

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
EvoArt wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat(DI): add Piste backend #1089

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/Test.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -98,6 +98,7 @@ jobs:
- GTPSA
- HyperHessians
- Mooncake
- Piste
- PolyesterForwardDiff
- ReverseDiff
- SparsityDetector
Expand Down
1 change: 1 addition & 0 deletions DifferentiationInterface/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0

### Added

- Support for [Piste.jl](https://github.com/EvoArt/Piste.jl) with `AutoPiste`, in forward and reverse mode
- New operators `value_gradient_and_hvp` and `value_gradient_and_hvp!`, which also return the primal value ([#1086](https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/1086))

### Deprecated
Expand Down
5 changes: 4 additions & 1 deletion DifferentiationInterface/Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527"
GTPSA = "b27dd330-f138-47c5-815b-40db9dd9b6e8"
Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
Piste = "f69a1a94-130c-4fa4-b8d2-e36d656cbe66"
PolyesterForwardDiff = "98d1487c-24ca-40b6-b7ab-df2af84e126b"
ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
Expand All @@ -44,6 +45,7 @@ DifferentiationInterfaceGPUArraysCoreExt = ["GPUArraysCore", "Adapt"]
DifferentiationInterfaceHyperHessiansExt = "HyperHessians"
DifferentiationInterfaceGTPSAExt = "GTPSA"
DifferentiationInterfaceMooncakeExt = "Mooncake"
DifferentiationInterfacePisteExt = "Piste"
DifferentiationInterfacePolyesterForwardDiffExt = [
"PolyesterForwardDiff",
"ForwardDiff",
Expand All @@ -60,7 +62,7 @@ DifferentiationInterfaceZygoteExt = ["Zygote", "ForwardDiff"]

[compat]
Adapt = "4.5.0"
ADTypes = "1.23.0"
ADTypes = "1.25.0"
ChainRulesCore = "1.23.0"
DiffResults = "1.1.0"
Diffractor = "=0.2.6"
Expand All @@ -75,6 +77,7 @@ GTPSA = "1.4.0"
HyperHessians = "0.3"
LinearAlgebra = "1"
Mooncake = "0.5.25"
Piste = "0.4"
PolyesterForwardDiff = "0.1.2"
ReverseDiff = "1.15.1"
SparseArrays = "1"
Expand Down
1 change: 1 addition & 0 deletions DifferentiationInterface/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ We support the following backends defined by [ADTypes.jl](https://github.com/Sci
- [HyperHessians.jl](https://github.com/KristofferC/HyperHessians.jl)
- [GTPSA.jl](https://github.com/bmad-sim/GTPSA.jl)
- [Mooncake.jl](https://github.com/chalk-lab/Mooncake.jl)
- [Piste.jl](https://github.com/EvoArt/Piste.jl)
- [PolyesterForwardDiff.jl](https://github.com/JuliaDiff/PolyesterForwardDiff.jl)
- [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl)
- [Symbolics.jl](https://github.com/JuliaSymbolics/Symbolics.jl)
Expand Down
12 changes: 12 additions & 0 deletions DifferentiationInterface/docs/src/explanation/backends.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ We support the following dense backend choices from [ADTypes.jl](https://github.
- [`AutoGTPSA`](@extref ADTypes.AutoGTPSA)
- [`AutoHyperHessians`](@extref ADTypes.AutoHyperHessians)
- [`AutoMooncake`](@extref ADTypes.AutoMooncake) and [`AutoMooncakeForward`](@extref ADTypes.AutoMooncake) (the latter is experimental)
- [`AutoPiste`](@extref ADTypes.AutoPiste)
- [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff)
- [`AutoReverseDiff`](@extref ADTypes.AutoReverseDiff)
- [`AutoSymbolics`](@extref ADTypes.AutoSymbolics)
Expand Down Expand Up @@ -52,6 +53,8 @@ In practice, many AD backends have custom implementations for high-level operato
| `AutoHyperHessians` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ |
| `AutoMooncake` | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| `AutoMooncakeForward` | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| `AutoPiste` (forward) | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| `AutoPiste` (reverse) | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
| `AutoPolyesterForwardDiff` | 🔀 | ❌ | 🔀 | ✅ | ✅ | 🔀 | 🔀 | 🔀 |
| `AutoReverseDiff` | ❌ | 🔀 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ |
| `AutoSymbolics` | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
Expand All @@ -74,6 +77,8 @@ Moreover, each context type is supported by a specific subset of backends:
| `AutoHyperHessians` | ✅ | ✅ |
| `AutoMooncake` | ✅ | ✅ |
| `AutoMooncakeForward` | ✅ | ✅ |
| `AutoPiste` (forward) | ✅ | ✅ |
| `AutoPiste` (reverse) | ✅ | ✅ |
| `AutoPolyesterForwardDiff` | ✅ | ✅ |
| `AutoReverseDiff` | ✅ | ❌ |
| `AutoSymbolics` | ✅ | ✅ |
Expand Down Expand Up @@ -197,6 +202,13 @@ For all operators, preparation generates an [executable function](https://docs.s

For `pullback`, preparation [builds the reverse rule](https://chalk-lab.github.io/Mooncake.jl/stable/understanding_mooncake/rule_system/) of the function.

### Piste

In forward mode, `pushforward` propagates `Piste.Dual` numbers carrying a whole batch of tangents, and `gradient` (for a `Vector` input) calls Piste's own chunked driver, with its workspace allocated during preparation.
In reverse mode, `pullback` and `gradient` record the function on a `Piste.Tape` and sweep it backwards, reusing a workspace allocated during preparation.

Piste is designed to survive `juliac --trim`, and prepared `gradient` calls through DifferentiationInterface do too, provided that the `chunksize` of a forward-mode backend is set explicitly.

### Tracker

We implement `pullback` based on `Tracker.back`.
Expand Down
24 changes: 24 additions & 0 deletions ...tiationInterface/ext/DifferentiationInterfacePisteExt/DifferentiationInterfacePisteExt.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
module DifferentiationInterfacePisteExt

using ADTypes: AutoPiste, ForwardMode, ReverseMode
import DifferentiationInterface as DI
using Piste:
Piste,
Dual,
GradientWorkspace,
ReverseWorkspace,
TapedReal,
gradient!,
pickchunk,
rev_value_and_pullback,
rev_value_and_pullback!,
value,
value_and_gradient!

DI.check_available(::AutoPiste) = true
DI.inplace_support(::AutoPiste) = DI.InPlaceSupported()

include("forward.jl")
include("reverse.jl")

end # module
Loading
Loading