remove Zygote, add DynamicPPL.DebugUtils in performance tips

Author: penelopeysm

usage/performance-tips/index.qmd

@@ -25,7 +25,7 @@ The following example:
 using Turing
 @model function gmodel(x)
     m ~ Normal()
-    for i in 1:length(x)
+    for i in eachindex(x)
         x[i] ~ Normal(m, 0.2)
     end
 end
@@ -44,34 +44,23 @@ end
 
 ## Choose your AD backend
 
-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
-supports several AD backends, including [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) (the default),
-[Mooncake](https://github.com/compintell/Mooncake.jl),
-[Zygote](https://github.com/FluxML/Zygote.jl), and
-[ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl).
+Automatic differentiation (AD) makes it possible to use modern, efficient gradient-based samplers like NUTS and HMC.
+This, however, also means that using a performant AD system is incredibly important.
+Turing currently supports several AD backends, including [ForwardDiff](https://github.com/JuliaDiff/ForwardDiff.jl) (the default), [Mooncake](https://github.com/chalk-lab/Mooncake.jl), and [ReverseDiff](https://github.com/JuliaDiff/ReverseDiff.jl).
 
-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
-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.
-`NUTS(adtype = AutoZygote())`. See [Automatic Differentiation]({{<meta usage-automatic-differentiation>}}) for details. Generally, `adtype = AutoForwardDiff()` is likely to be the fastest and most reliable for models with
-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
-operations. If in doubt, it's easy to try a few different backends to see how they compare.
+For many common types of models, the default ForwardDiff backend performs well, and there is no need to worry about changing it.
+However, if you need more speed, you can try 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. `NUTS(; adtype = AutoMooncake())`.
 
-### Special care for Zygote
-
-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
-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,
-vectorized operations are still likely to perform best.
-
-Avoiding loops can be done using `filldist(dist, N)` and `arraydist(dists)`. `filldist(dist, N)` creates a multivariate distribution that is composed of `N` identical and independent
-copies of the univariate distribution `dist` if `dist` is univariate, or it creates a matrix-variate distribution composed of `N` identical and independent copies of the multivariate
-distribution `dist` if `dist` is multivariate. `filldist(dist, N, M)` can also be used to create a matrix-variate distribution from a univariate distribution `dist`.  `arraydist(dists)`
-is similar to `filldist` but it takes an array of distributions `dists` as input. Writing a [custom distribution](advanced) with a custom adjoint is another option to avoid loops.
+Generally, `adtype = AutoForwardDiff()` is likely to be the fastest and most reliable for models with few parameters (say, less than 20 or so), while reverse-mode backends such as `AutoMooncake()` or `AutoReverseDiff()` will perform better for models with many parameters or linear algebra operations.
+If in doubt, you can benchmark your model with different backends to see which one performs best.
+See the [Automatic Differentiation]({{<meta usage-automatic-differentiation>}}) page for details.
 
 ### Special care for ReverseDiff with a compiled tape
 
-For large models, the fastest option is often ReverseDiff with a compiled tape, specified as `adtype=AutoReverseDiff(true)`. However, it is important to note that if your model contains any
-branching code, such as `if`-`else` statements, **the gradients from a compiled tape may be inaccurate, leading to erroneous results**. If you use this option for the (considerable) speedup it
-can provide, make sure to check your code. It's also a good idea to verify your gradients with another backend.
+For large models, the fastest option is often ReverseDiff with a compiled tape, specified as `adtype=AutoReverseDiff(; compile=true)`.
+However, it is important to note that if your model contains any branching code, such as `if`-`else` statements, **the gradients from a compiled tape may be inaccurate, leading to erroneous results**.
+If you use this option for the (considerable) speedup it can provide, make sure to check your code for branching and ensure that it does not affect the gradients.
+It is also a good idea to verify your gradients with another backend.
 
 ## Ensure that types in your model can be inferred
 
@@ -117,10 +106,10 @@ Alternatively, you could use `filldist` in this example:
 end
 ```
 
-Note that you can use `@code_warntype` to find types in your model definition that the compiler cannot infer.
-They are marked in red in the Julia REPL.
+You can use DynamicPPL's debugging utilities to find types in your model definition that the compiler cannot infer.
+These will be marked in red in the Julia REPL (much like when using the `@code_warntype` macro).
 
-For example, consider the following simple program:
+For example, consider the following model:
 
 ```{julia}
 @model function tmodel(x)
@@ -131,22 +120,22 @@ For example, consider the following simple program:
 end
 ```
 
-We can use
+Because the element type of `p` is an abstract type (`Real`), the compiler cannot infer a concrete type for `p[1]`.
+To detect this, we can use
 
 ```{julia}
 #| eval: false
-using Random
-
 model = tmodel(1.0)
 
-@code_warntype model.f(
-    model,
-    Turing.VarInfo(model),
-    Turing.SamplingContext(
-        Random.default_rng(), Turing.SampleFromPrior(), Turing.DefaultContext()
-    ),
-    model.args...,
-)
+using DynamicPPL
+DynamicPPL.DebugUtils.model_warntype(model)
 ```
 
-to inspect type inference in the model.
+In this particular model, the following call to `getindex` should be highlighted in red (the exact numbers may vary):
+
+```
+[...]
+│    %120 = p::AbstractVector
+│    %121 = Base.getindex(%120, 1)::Any
+[...]
+```