Skip to content

Commit d42a408

Browse files
authored
Update to Turing v0.35 (#559)
* Add readme note on Julia version * Bump Turing to 0.35 * Update minimum supported Julia version * Remove unnecessary version qualifier * Remove Tracker and replace with Mooncake, except in BNN doc * Use Mooncake in BNN doc (#521) * Fix BNN doc to work with Mooncake
1 parent ba5e408 commit d42a408

File tree

9 files changed

+73
-18
lines changed

9 files changed

+73
-18
lines changed

Manifest.toml

+38-1
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
julia_version = "1.10.6"
44
manifest_format = "2.0"
5-
project_hash = "52677c2802b9702303d49a2f40fdefa34939b429"
5+
project_hash = "53e97ef537fe2e98eecbe367cd36c4bc6a201e3e"
66

77
[[deps.ADTypes]]
88
git-tree-sha1 = "72af59f5b8f09faee36b4ec48e014a79210f2f4f"
@@ -708,6 +708,12 @@ git-tree-sha1 = "23163d55f885173722d1e4cf0f6110cdbaf7e272"
708708
uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
709709
version = "1.15.1"
710710

711+
[[deps.DiffTests]]
712+
deps = ["LinearAlgebra", "SparseArrays", "Statistics"]
713+
git-tree-sha1 = "b92beb1933df01bf4915d3a05e54c2a0aad312c7"
714+
uuid = "de460e47-3fe3-5279-bb4a-814414816d5d"
715+
version = "0.1.2"
716+
711717
[[deps.DifferentialEquations]]
712718
deps = ["BoundaryValueDiffEq", "DelayDiffEq", "DiffEqBase", "DiffEqCallbacks", "DiffEqNoiseProcess", "JumpProcesses", "LinearAlgebra", "LinearSolve", "NonlinearSolve", "OrdinaryDiffEq", "Random", "RecursiveArrayTools", "Reexport", "SciMLBase", "SteadyStateDiffEq", "StochasticDiffEq", "Sundials"]
713719
git-tree-sha1 = "d55af9d6b51c54f81ae30d1a463206d32cc4c24a"
@@ -2109,6 +2115,11 @@ git-tree-sha1 = "ec4f7fbeab05d7747bdf98eb74d130a2a2ed298d"
21092115
uuid = "e1d29d7a-bbdc-5cf2-9ac0-f12de2c33e28"
21102116
version = "1.2.0"
21112117

2118+
[[deps.MistyClosures]]
2119+
git-tree-sha1 = "1142aefd845c608f3c70e4c202c4aae725cab67b"
2120+
uuid = "dbe65cb8-6be2-42dd-bbc5-4196aaced4f4"
2121+
version = "2.0.0"
2122+
21122123
[[deps.Mmap]]
21132124
uuid = "a63ad114-7e13-5084-954f-fe012c677804"
21142125

@@ -2118,6 +2129,32 @@ git-tree-sha1 = "2c140d60d7cb82badf06d8783800d0bcd1a7daa2"
21182129
uuid = "78c3b35d-d492-501b-9361-3d52fe80e533"
21192130
version = "0.8.1"
21202131

2132+
[[deps.Mooncake]]
2133+
deps = ["ADTypes", "ChainRules", "ChainRulesCore", "DiffRules", "DiffTests", "ExprTools", "FunctionWrappers", "Graphs", "InteractiveUtils", "LinearAlgebra", "MistyClosures", "Random", "Setfield", "Test"]
2134+
git-tree-sha1 = "9fa3c8c11fe5286badb9c24e8d2420b3a07f9a8c"
2135+
uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
2136+
version = "0.4.50"
2137+
2138+
[deps.Mooncake.extensions]
2139+
MooncakeAllocCheckExt = "AllocCheck"
2140+
MooncakeCUDAExt = "CUDA"
2141+
MooncakeDynamicPPLExt = "DynamicPPL"
2142+
MooncakeJETExt = "JET"
2143+
MooncakeLuxLibExt = "LuxLib"
2144+
MooncakeLuxLibSLEEFPiratesExtension = ["LuxLib", "SLEEFPirates"]
2145+
MooncakeNNlibExt = "NNlib"
2146+
MooncakeSpecialFunctionsExt = "SpecialFunctions"
2147+
2148+
[deps.Mooncake.weakdeps]
2149+
AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
2150+
CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
2151+
DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
2152+
JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
2153+
LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
2154+
NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
2155+
SLEEFPirates = "476501e8-09a2-5ece-8869-fb82de89a1fa"
2156+
SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
2157+
21212158
[[deps.MozillaCACerts_jll]]
21222159
uuid = "14a3606d-f60d-562e-9121-12d972cd8159"
21232160
version = "2023.1.10"

Project.toml

+1-1
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
3333
Measures = "442fdcdd-2543-5da2-b0f3-8c86c306513e"
3434
Memoization = "6fafb56a-5788-4b4e-91ca-c0cea6611c73"
3535
MicroCanonicalHMC = "234d2aa0-2291-45f7-9047-6fa6f316b0a8"
36+
Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
3637
NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
3738
Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
3839
OptimizationNLopt = "4e6fcdb7-1186-4e1f-a706-475e75c168bb"
@@ -48,7 +49,6 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
4849
StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
4950
StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
5051
StatsPlots = "f3b207a7-027a-5e70-b257-86293d7955fd"
51-
Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
5252
Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
5353
TuringBenchmarking = "0db1332d-5c25-4deb-809f-459bc696f94f"
5454
UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"

README.md

+17
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,17 @@ This repository is part of [Turing.jl's](https://turinglang.org/) website (i.e.
44
- The `master` branch contains the quarto source
55
- The `gh-pages` branch contains the `html` version of these documents compiled from the `master` branch.
66

7+
> [!NOTE]
8+
> Due to [an incompatibility between Bijectors.jl and
9+
> Enzyme.jl](https://github.com/TuringLang/Bijectors.jl/pull/341), the docs
10+
> currently must be built with Julia 1.10. You can do this either by making
11+
> Julia 1.10 your default Julia version (`juliaup default 1.10`), or by
12+
> explicitly specifying Julia 1.10 when running Quarto:
13+
>
14+
> ```bash
15+
> QUARTO_JULIA=$(julia +1.10 -e "println(Sys.BINDIR)")/julia quarto render
16+
> ```
17+
718
## Local development
819
920
To get started with the docs website locally, you'll need to have [Quarto](https://quarto.org/docs/download/) installed.
@@ -79,6 +90,12 @@ If you find that Quarto's execution is failing with errors that aren't reproduci
7990
quarto render /path/to/index.qmd --execute-daemon-restart
8091
```
8192
93+
And also, kill any stray Quarto processes that are still running (sometimes it keeps running in the background):
94+
95+
```bash
96+
pkill -9 -f quarto
97+
```
98+
8299
## License
83100

84101
This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

tutorials/03-bayesian-neural-network/index.qmd

+3-3
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ using Turing
1919
using FillArrays
2020
using Lux
2121
using Plots
22-
using Tracker
22+
import Mooncake
2323
using Functors
2424
2525
using LinearAlgebra
@@ -189,7 +189,7 @@ const nn = StatefulLuxLayer{true}(nn_initial, nothing, st)
189189
parameters ~ MvNormal(zeros(nparameters), Diagonal(abs2.(sigma .* ones(nparameters))))
190190
191191
# Forward NN to make predictions
192-
preds = Lux.apply(nn, xs, vector_to_parameters(parameters, ps))
192+
preds = Lux.apply(nn, xs, f32(vector_to_parameters(parameters, ps)))
193193
194194
# Observe each prediction.
195195
for i in eachindex(ts)
@@ -208,7 +208,7 @@ setprogress!(false)
208208
```{julia}
209209
# Perform inference.
210210
N = 2_000
211-
ch = sample(bayes_nn(reduce(hcat, xs), ts), NUTS(; adtype=AutoTracker()), N);
211+
ch = sample(bayes_nn(reduce(hcat, xs), ts), NUTS(; adtype=AutoMooncake(; config=nothing)), N);
212212
```
213213

214214
Now we extract the parameter samples from the sampled chain as `θ` (this is of size `5000 x 20` where `5000` is the number of iterations and `20` is the number of parameters).

tutorials/10-bayesian-differential-equations/index.qmd

+1-1
Original file line numberDiff line numberDiff line change
@@ -320,7 +320,7 @@ More theoretical details on these methods can be found at: https://docs.sciml.ai
320320
While these sensitivity analysis methods may seem complicated, using them is dead simple.
321321
Here is a version of the Lotka-Volterra model using adjoint sensitivities.
322322

323-
All we have to do is switch the AD backend to one of the adjoint-compatible backends (ReverseDiff, Tracker, or Zygote)!
323+
All we have to do is switch the AD backend to one of the adjoint-compatible backends (ReverseDiff or Zygote)!
324324
Notice that on this model adjoints are slower.
325325
This is because adjoints have a higher overhead on small parameter models and therefore we suggest using these methods only for models with around 100 parameters or more.
326326
For more details, see https://arxiv.org/abs/1812.01892.

tutorials/docs-00-getting-started/index.qmd

+1-1
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@ Pkg.instantiate();
1616

1717
To use Turing, you need to install Julia first and then install Turing.
1818

19-
You will need to install Julia 1.7 or greater, which you can get from [the official Julia website](http://julialang.org/downloads/).
19+
You will need to install Julia 1.10 or greater, which you can get from [the official Julia website](http://julialang.org/downloads/).
2020

2121
Turing is officially registered in the [Julia General package registry](https://github.com/JuliaRegistries/General), which means that you can install a stable version of Turing by running the following in the Julia REPL:
2222

tutorials/docs-10-using-turing-autodiff/index.qmd

+4-4
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,8 @@ Pkg.instantiate();
1212

1313
## Switching AD Modes
1414

15-
Turing currently supports four automatic differentiation (AD) backends for sampling: [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) for forward-mode AD; and [ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl), [Zygote](https://github.com/FluxML/Zygote.jl), and [Tracker](https://github.com/FluxML/Tracker.jl) for reverse-mode AD.
16-
While `Tracker` is still available, its use is discouraged due to a lack of active maintenance.
17-
`ForwardDiff` is automatically imported by Turing. To utilize `Zygote` or `ReverseDiff` for AD, users must explicitly import them with `using Zygote` or `using ReverseDiff`, alongside `using Turing`.
15+
Turing currently supports four automatic differentiation (AD) backends for sampling: [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) for forward-mode AD; and [Mooncake](https://github.com/compintell/Mooncake.jl), [ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl), and [Zygote](https://github.com/FluxML/Zygote.jl) for reverse-mode AD.
16+
`ForwardDiff` is automatically imported by Turing. To utilize `Mooncake`, `Zygote`, or `ReverseDiff` for AD, users must explicitly import them with `import Mooncake`, `import Zygote` or `import ReverseDiff`, alongside `using Turing`.
1817

1918
As of Turing version v0.30, the global configuration flag for the AD backend has been removed in favour of [`AdTypes.jl`](https://github.com/SciML/ADTypes.jl), allowing users to specify the AD backend for individual samplers independently.
2019
Users can pass the `adtype` keyword argument to the sampler constructor to select the desired AD backend, with the default being `AutoForwardDiff(; chunksize=0)`.
@@ -69,7 +68,8 @@ Generally, reverse-mode AD, for instance `ReverseDiff`, is faster when sampling
6968
If the differentiation method is not specified in this way, Turing will default to using whatever the global AD backend is.
7069
Currently, this defaults to `ForwardDiff`.
7170

72-
The most reliable way to ensure you are using the fastest AD that works for your problem is to benchmark them using `TuringBenchmarking`:
71+
The most reliable way to ensure you are using the fastest AD that works for your problem is to benchmark them using [`TuringBenchmarking`](https://github.com/TuringLang/TuringBenchmarking.jl):
72+
7373
```{julia}
7474
using TuringBenchmarking
7575
benchmark_model(gdemo(1.5, 2), adbackends=[AutoForwardDiff(), AutoReverseDiff()])

tutorials/docs-12-using-turing-guide/index.qmd

+2-2
Original file line numberDiff line numberDiff line change
@@ -166,7 +166,7 @@ The `chains` variable now contains a `Chains` object which can be indexed by cha
166166

167167
#### Multithreaded sampling
168168

169-
If you wish to perform multithreaded sampling and are running Julia 1.3 or greater, you can call `sample` with the following signature:
169+
If you wish to perform multithreaded sampling, you can call `sample` with the following signature:
170170

171171
```{julia}
172172
#| eval: false
@@ -514,7 +514,7 @@ ForwardDiff (Turing's default AD backend) uses forward-mode chunk-wise AD. The c
514514

515515
#### AD Backend
516516

517-
Turing supports four automatic differentiation (AD) packages in the back end during sampling. The default AD backend is [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) for forward-mode AD. Three reverse-mode AD backends are also supported, namely [Tracker](https://github.com/FluxML/Tracker.jl), [Zygote](https://github.com/FluxML/Zygote.jl) and [ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl). `Zygote` and `ReverseDiff` are supported optionally if explicitly loaded by the user with `using Zygote` or `using ReverseDiff` next to `using Turing`.
517+
Turing supports four automatic differentiation (AD) packages in the back end during sampling. The default AD backend is [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) for forward-mode AD. Three reverse-mode AD backends are also supported, namely [Mooncake](https://github.com/compintell/Mooncake.jl), [Zygote](https://github.com/FluxML/Zygote.jl) and [ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl). `Mooncake`, `Zygote`, and `ReverseDiff` also require the user to explicitly load them using `import Mooncake`, `import Zygote`, or `import ReverseDiff` next to `using Turing`.
518518

519519
For more information on Turing's automatic differentiation backend, please see the [Automatic Differentiation]({{<meta using-turing-autodiff>}}) article.
520520

tutorials/docs-13-using-turing-performance-tips/index.qmd

+6-5
Original file line numberDiff line numberDiff line change
@@ -43,19 +43,20 @@ end
4343
## Choose your AD backend
4444

4545
Automatic differentiation (AD) makes it possible to use modern, efficient gradient-based samplers like NUTS and HMC, and that means a good AD system is incredibly important. Turing currently
46-
supports several AD backends, including [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) (the default), [Zygote](https://github.com/FluxML/Zygote.jl),
47-
[ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl), and [Tracker](https://github.com/FluxML/Tracker.jl). Experimental support is also available for
48-
[Tapir](https://github.com/withbayes/Tapir.jl).
46+
supports several AD backends, including [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) (the default),
47+
[Mooncake](https://github.com/compintell/Mooncake.jl),
48+
[Zygote](https://github.com/FluxML/Zygote.jl), and
49+
[ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl).
4950

5051
For many common types of models, the default ForwardDiff backend performs great, and there is no need to worry about changing it. However, if you need more speed, you can try
5152
different backends via the standard [ADTypes](https://github.com/SciML/ADTypes.jl) interface by passing an `AbstractADType` to the sampler with the optional `adtype` argument, e.g.
5253
`NUTS(adtype = AutoZygote())`. See [Automatic Differentiation]({{<meta using-turing-autodiff>}}) for details. Generally, `adtype = AutoForwardDiff()` is likely to be the fastest and most reliable for models with
5354
few parameters (say, less than 20 or so), while reverse-mode backends such as `AutoZygote()` or `AutoReverseDiff()` will perform better for models with many parameters or linear algebra
5455
operations. If in doubt, it's easy to try a few different backends to see how they compare.
5556

56-
### Special care for Zygote and Tracker
57+
### Special care for Zygote
5758

58-
Note that Zygote and Tracker will not perform well if your model contains `for`-loops, due to the way reverse-mode AD is implemented in these packages. Zygote also cannot differentiate code
59+
Note that Zygote will not perform well if your model contains `for`-loops, due to the way reverse-mode AD is implemented in these packages. Zygote also cannot differentiate code
5960
that contains mutating operations. If you can't implement your model without `for`-loops or mutation, `ReverseDiff` will be a better, more performant option. In general, though,
6061
vectorized operations are still likely to perform best.
6162

0 commit comments

Comments
 (0)