### @opt_out usage example

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Example showing how to opt-out of a specific rule for a specialized type.

```julia
function rrule(::typeof(foo), x::AbstractArray)
    foo_pullback(ȳ) = ...
    return foo(x), foo_pullback
end

@opt_out rrule(::typeof(foo), ::FancyArray)
```

--------------------------------

### @scalar_rule simplified syntax

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Example of the simplified @scalar_rule syntax without explicit setup code.

```julia
@scalar_rule(f(x₁, x₂, ...),
             (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
             (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
             ...)
```

```julia
@scalar_rule(f(x₁, x₂, ...),
             @setup(nothing),
             (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
             (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
             ...)
```

--------------------------------

### Original sin pullback using inputs only

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the original formulation for sin that stores only x and uses cos(pb.x) in the pullback. This is the starting point for the refactoring example.

```julia
function augmented_primal(::typeof(sin), x)
    y = sin(x)
    return y, PullbackMemory(sin; x=x)
end
(pb::PullbackMemory)(ȳ) = ȳ * cos(pb.x)
```

--------------------------------

### @scalar_rule(f(x₁, x₂, ...), @setup(statement₁, statement₂, ...), (∂f₁_∂x₁, ∂f₁_∂x₂, ...), (∂f₂_∂x₁, ∂f₂_∂x₂, ...), ...)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Generates simple scalar forward or reverse rules using provided partial derivatives. Creates corresponding methods for frule and rrule. Supports optional type constraints and setup code.

```APIDOC
## @scalar_rule(f(x₁, x₂, ...), @setup(statement₁, statement₂, ...), (∂f₁_∂x₁, ∂f₁_∂x₂, ...), (∂f₂_∂x₁, ∂f₂_∂x₂, ...), ...)

### Description
A convenience macro that generates simple scalar forward or reverse rules using the provided partial derivatives. It generates corresponding methods for `frule` and `rrule`.

### Usage
```julia
@scalar_rule(f(x₁, x₂, ...),
             @setup(statement₁, statement₂, ...),
             (∂f₁_∂x₁, ∂f₁_∂x₂, ...),
             (∂f₂_∂x₁, ∂f₂_∂x₂, ...),
             ...)
```

### Parameters
- `f(x₁, x₂, ...)` - The function signature. Type constraints may be provided; otherwise each parameter is constrained to `Number`.
- `@setup(statement₁, statement₂, ...)` - Optional setup code. Can be elided if no setup is needed.
- `(∂f₁_∂x₁, ∂f₁_∂x₂, ...)` - Partial derivatives for the first output.
- `(∂f₂_∂x₁, ∂f₂_∂x₂, ...)` - Partial derivatives for the second output.
- Additional derivative tuples for more outputs.

### Generated Methods
```julia
function ChainRulesCore.frule((NoTangent(), Δx₁, Δx₂, ...), ::typeof(f), x₁::Number, x₂::Number, ...)
    Ω = f(x₁, x₂, ...)
    $(statement₁, statement₂, ...)
    return Ω, (
            (∂f₁_∂x₁ * Δx₁ + ∂f₁_∂x₂ * Δx₂ + ...),
            (∂f₂_∂x₁ * Δx₁ + ∂f₂_∂x₂ * Δx₂ + ...),
            ...
        )
end

function ChainRulesCore.rrule(::typeof(f), x₁::Number, x₂::Number, ...)
    Ω = f(x₁, x₂, ...)
    $(statement₁, statement₂, ...)
    return Ω, ((ΔΩ₁, ΔΩ₂, ...)) -> (
            NoTangent(),
            ∂f₁_∂x₁ * ΔΩ₁ + ∂f₂_∂x₁ * ΔΩ₂ + ...),
            ∂f₁_∂x₂ * ΔΩ₁ + ∂f₂_∂x₂ * ΔΩ₂ + ...),
            ...
        )
end
```

### Notes
- If no type constraints are provided, each parameter is constrained to `Number`. Explicit constraints override this.
- Does not support closures/functors. The derivative with respect to the function itself is always `NoTangent()`.
- The result of `f(x₁, x₂, ...)` is bound to `Ω` for use in derivative/setup expressions.
- Assumes complex functions are holomorphic. For non-holomorphic functions, define `frule` and `rrule` manually.
- If the derivative is one, `true` can be used as the most general multiplicative identity.
- The `@setup` argument can be elided if no setup code is needed.
```

--------------------------------

### augmented_primal and pullback_at for sin using PullbackMemory

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Example for sin: augmented_primal computes sincos and stores cx in PullbackMemory; pullback_at uses pb.cx to compute the gradient.

```julia
function augmented_primal(::typeof(sin), x)
    y, cx = sincos(x)
    return y, PullbackMemory(sin; cx=cx)
end

pullback_at(pb::PullbackMemory{typeof(sin)}, ȳ) = ȳ * pb.cx
```

--------------------------------

### Define an frule for sin(x)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/index.html

Example implementation of an frule for the sine function.

```julia
function frule((_, Δx), ::typeof(sin), x)
    return sin(x), cos(x) * Δx
end
```

--------------------------------

### @non_differentiable usage example

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Example showing how to declare a function as non-differentiable and verify the pullback returns NoTangent().

```julia
julia> @non_differentiable Base.:(==)(a, b)

julia> _, pullback = rrule(==, 2.0, 3.0);

julia> pullback(1.0)
(NoTangent(), NoTangent(), NoTangent())
```

--------------------------------

### Callable PullbackMemory for sin

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Example for sin: augmented_primal stores cx, and the callable PullbackMemory computes the gradient using pb.cx.

```julia
function augmented_primal(::typeof(sin), x)
    y, cx = sincos(x)
    return y, PullbackMemory(sin; cx=cx)
end
(pb::PullbackMemory{typeof(sin)})(ȳ) = ȳ * pb.cx
```

--------------------------------

### augmented_primal and pullback_at for sigmoid using PullbackMemory

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Example for the logistic sigmoid σ: augmented_primal computes exp and stores y and ex; pullback_at uses pb.y and pb.ex to compute the gradient.

```julia
function augmented_primal(::typeof(σ), x)
    ex = exp(x)
    y = ex / (1 + ex)
    return y, PullbackMemory(σ; y=y, ex=ex)
end

pullback_at(pb::PullbackMemory{typeof(σ)}, ȳ) = ȳ * pb.y / (1 + pb.ex)
```

--------------------------------

### @non_differentiable with type constraints

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Example showing how to apply type constraints to the signature when using @non_differentiable.

```julia
julia> @non_differentiable Base.length(xs::Union{Number, Array})

julia> frule((ZeroTangent(), 1), length, [2.0, 3.0])
(2, NoTangent())
```

--------------------------------

### Define an rrule for sin(x)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/index.html

Example implementation of an rrule for the sine function, returning a closure as the pullback.

```julia
function rrule(::typeof(sin), x)
    sin_pullback(Δy) = (NoTangent(), cos(x)' * Δy)
    return sin(x), sin_pullback
end
```

--------------------------------

### Example of code causing AD mutation errors

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/tips_for_packages.html

An example of code that uses mutation, which is not supported by Zygote and will cause errors during differentiation.

```julia
somes = []
things = []

function loss(x, y)
    some = f(x, y)
    thing = g(x)
    
    # log
    push!(somes, some)
    push!(things, thing)

    return some + thing
end
```

--------------------------------

### ignore_derivatives incorrect usage example

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Demonstrates how incorrect usage can lead to zero gradients with respect to function arguments.

```julia
function wrong_grads(x)
    y = ones(3)
    ignore_derivatives() do
        push!(y, x)
    end
    return sum(y)
end
```

--------------------------------

### frule for mutating function foo!

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/mutation_support.html

Example frule for foo! that mutates the tangent partners in sync with the primal values. It asserts that the tangents are MutableTangent, updates ȧ.x and ḃ.x, and returns the primal and its tangent.

```julia
function ChainRulesCore.frule((_, ȧ, ḃ), ::typeof(foo!), a::Base.RefValue, b::Base.RefValue)
    @assert ȧ isa MutableTangent{typeof(a)}
    @assert ḃ isa MutableTangent{typeof(b)}

    a[] *= 2
    ȧ.x *= 2  # `.x` is the field that lives behind RefValues

    b[] = 5.0
    ḃ.x = zero_tangent(5.0)  # or since we know that the zero for a Float64 is zero could write `ḃ.x = 0.0`

    return a, ȧ
end
```

--------------------------------

### Callable PullbackMemory for sigmoid

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Example for the logistic sigmoid σ: augmented_primal stores y and ex, and the callable PullbackMemory computes the gradient using pb.y and pb.ex.

```julia
function augmented_primal(::typeof(σ), x)
    ex = exp(x)
    y = ex / (1 + ex)
    return y, PullbackMemory(σ; y=y, ex=ex)
end

(pb::PullbackMemory{typeof(σ)})(ȳ) = ȳ * pb.y / (1 + pb.ex)
```

--------------------------------

### Primal mutating function foo!

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/mutation_support.html

Example primal function that takes two Refs, doubles the first in place, overwrites the second with 5.0, and returns the first.

```julia
function foo!(a::Base.RefValue, b::Base.RefValue)
    a[] *= 2
    b[] = 5.0
    return a
end
```

--------------------------------

### Plot piecewise slope change

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/nondiff_points.html

Piecewise slope change example. Options for the derivative at 0 include 1 (backwards finite differencing), 5 (forwards finite differencing), or 3 (mean, central finite differencing).

```julia
plot(x-> x < 0 ? x : 5x)
```

--------------------------------

### Plot abs for local minimum

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/nondiff_points.html

Classic example of a function where the derivative is not defined at a point (local minimum). The AD system would choose either 1 or -1, but returning 0 is nicer.

```julia
plot(abs)
```

--------------------------------

### Julia function with multiple uses of X requiring gradient accumulation

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/gradient_accumulation.html

Example function where X is used twice (X[1] and X[2]), so the AD system must accumulate the gradients ā and b̄ to compute X̄. Demonstrates the need for + on tangent types.

```julia
function sum_first_and_second(X::Array{Float64})
    a = X[1]
    b = X[2]
    y = a + b
    return y
end
```

--------------------------------

### Define the function `f` with a custom struct `Foo`

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/converting_zygoterules.html

Example function definition that will be used for both the ZygoteRules and ChainRules implementations. It defines a struct Foo with fields a and b, and a function f(x, y::Foo, z) = 2*x + y.a.

```julia
struct Foo
    a::Float64,
    b::Float64
end

f(x, y::Foo, z) = 2*x + y.a
```

--------------------------------

### rrule for map using frule_via_ad with HasForwardsMode

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/ruleconfig.html

Defines a reverse-mode rrule for map that calls forward-mode AD via frule_via_ad when the RuleConfig supports HasForwardsMode. The example asserts that f is not a functor/closure and returns a pullback that multiplies the output sensitivity by the forward-mode derivatives.

```julia
function rrule(config::RuleConfig{>:HasForwardsMode}, ::typeof(map), f::Function, x::Array{<:Real})
    # real code would support functors/closures, but in interest of keeping example short we exclude it:
    @assert (fieldcount(typeof(f)) == 0) "Functors/Closures are not supported"

    y_and_ẏ = map(x) do xi
        frule_via_ad(config, (NoTangent(), one(xi)), f, xi)
    end
    y = first.(y_and_ẏ)
    ẏ = last.(y_and_ẏ)

    pullback_map(ȳ) = NoTangent(), NoTangent(), ȳ .* ẏ
    return y, pullback_map
end
```

--------------------------------

### Function that mutates arrays internally (addone)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html

Example of a function that mutates an internal array but not its input, causing Zygote's 'Mutating arrays is not supported' error. This pattern is common and requires an rrule to work with Zygote.

```julia
function addone(a::AbstractArray)
    b = similar(a)
    b .= a .+ 1
    return sum(b)
end
```

--------------------------------

### Compare with other autodiff packages

Source: https://juliadiff.org/ChainRulesCore.jl/stable/index.html

Demonstrates how to calculate the derivative of the same function using FiniteDifferences.jl, ForwardDiff.jl, and Zygote.jl.

```julia
#### Find dfoo/dx via FiniteDifferences.jl
using FiniteDifferences
central_fdm(5, 1)(foo, x)
# output
-1.0531613736418257

#### Find dfoo/dx via ForwardDiff.jl
using ForwardDiff
ForwardDiff.derivative(foo, x)
# output
-1.0531613736418153

#### Find dfoo/dx via Zygote.jl
using Zygote
Zygote.gradient(foo, x)
# output
(-1.0531613736418153,)
```

--------------------------------

### pullback_at concept for primal program

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the conceptual primal program and pullback_at call that records the primal function, inputs, and outputs on the tape, then pulls back the output sensitivity to the input sensitivity.

```julia
y = f(x)  # primal program
x̄ = pullback_at(f, x, y, ȳ)
```

--------------------------------

### Benchmark logistic sigmoid separately vs reusing exp(x)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Benchmarks computing σ(x) and σ(-x) separately (11.658 ns total) versus reusing exp(x) (7.878 ns total), demonstrating a 30-40% speedup.

```julia
julia> @btime 1/(1+exp(x)) setup=(x=rand());
  5.622 ns (0 allocations: 0 bytes)

julia> @btime 1/(1+exp(-x)) setup=(x=rand());
  6.036 ns (0 allocations: 0 bytes)

julia> 5.622 + 6.036
11.658
```

```julia
julia> @btime exp(x) setup=(x=rand());
  5.367 ns (0 allocations: 0 bytes)

julia> @btime ex/(1+ex) setup=(ex=exp(rand()));
  1.255 ns (0 allocations: 0 bytes)

julia> @btime 1/(1+ex) setup=(ex=exp(rand()));
  1.256 ns (0 allocations: 0 bytes)

julia> 5.367 + 1.255 + 1.256
7.878
```

--------------------------------

### New API using augmented_primal and pullback_at with PullbackMemory

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the updated API where augmented_primal returns a PullbackMemory and pullback_at takes only the memory and the output sensitivity. This is the cleaner interface after introducing PullbackMemory.

```julia
y = f(x)  # primal program
y, pb = augmented_primal(f, x)
x̄ = pullback_at(pb, ȳ)
```

--------------------------------

### Thunk Type and Usage

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

A Thunk represents a deferred computation. The example demonstrates creating a thunk with @thunk and evaluating it with unthunk.

```julia
Thunk(()->v)
```

```julia
julia> t = @thunk(3);

julia> unthunk(t)
3
```

--------------------------------

### Verify rrule for exception handling

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html

Successful gradient calculation after defining the rrule.

```julia
julia> gradient(exception, 3.0)
(6.0,)
```

--------------------------------

### Implement frule for sum(abs2, X; dims)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Generic forward-mode rule for sum(abs2, X; dims). The pushforward is computed as sum(2 .* real.(conj.(X) .* ΔX); dims = dims).

```julia
function frule(
    (_, _, ΔX),
    ::typeof(sum),
    ::typeof(abs2),
    X::Array{<:RealOrComplex};
    dims = :,
)
    Ω = sum(abs2, X; dims = dims)
    ∂Ω = sum(2 .* real.(conj.(X) .* ΔX); dims = dims)
    return (Ω, ∂Ω)
end
```

--------------------------------

### frule for matrix inversion (inv)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Implements the forward-mode rule for matrix inversion. The pushforward is -Ω * ΔA * Ω, derived from differentiating the constraint ΩA = I. Use for Matrix inputs.

```julia
function frule((_, ΔA), ::typeof(inv), A::Matrix{<:RealOrComplex})
    Ω = inv(A)
    ∂Ω = -Ω * ΔA * Ω
    return (Ω, ∂Ω)
end
```

--------------------------------

### Dump TimeSample instance

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/many_tangents.html

Shows the internal structure of a TimeSample instance, revealing the nested fields of DateTime.

```julia
julia> dump(TimeSample(now(), 2.6e9))
TimeSample
  time: DateTime
    instant: Dates.UTInstant{Millisecond}
      periods: Millisecond
        value: Int64 63720043490844
  value: Float64 2.6e9
```

--------------------------------

### Generated frule and rrule for @non_differentiable

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/rule_definition_tools.html

These are the frule and rrule definitions automatically generated by the @non_differentiable macro for the string concatenation example. The types of arguments are propagated to the definitions, which is needed for functions that are differentiable for some argument types but not others.

```julia
function ChainRulesCore.frule(var"##_#1600", ::Core.Typeof(*), String::Any...; kwargs...)
    return (*(String...; kwargs...), NoTangent())
end
function ChainRulesCore.rrule(::Core.Typeof(*), String::Any...; kwargs...)
    return (*(String...; kwargs...), function var"*_pullback"(_)
        (ZeroTangent(), ntuple((_->NoTangent()), 0 + length(String))...)
    end)
end
```

--------------------------------

### Define the test function using ZygoteRules.@adjoint

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/converting_zygoterules.html

Shows how the rule is written with the ZygoteRules macro, where the pullback returns partial derivatives and the macro automatically adds a nothing for the function itself.

```julia
@adjoint function f(x, y::Foo, z)
    f_pullback(Ω̄) = (2Ω̄, NamedTuple(;a=Ω̄, b=nothing), nothing)
    return f(x, y, z), f_pullback
end
```

--------------------------------

### pullback_at with intermediates for sigmoid

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the pullback_at implementation for the logistic sigmoid σ that uses the stored intermediate ex to compute the gradient. This is the first step in reducing the primal and gradient pass time.

```julia
pullback_at(::typeof(σ), x, y, ȳ, intermediates) = ȳ * y / (1 + intermediates.ex)
```

--------------------------------

### Original sigmoid pullback using inputs only

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the original formulation for the logistic sigmoid σ that stores y and x and uses σ(-pb.x) in the pullback.

```julia
function augmented_primal(::typeof(σ), x)
    y = σ(x)
    return y, PullbackMemory(σ; y=y, x=x)
end
(pb::PullbackMemory{typeof(σ)})(ȳ) = ȳ * pb.y * σ(-pb.x)
```

--------------------------------

### augmented_primal concept with intermediates

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the augmented_primal call that returns the primal output and intermediates, which are then passed to pullback_at for the gradient pass.

```julia
y = f(x)  # primal program
y, intermediates = augmented_primal(f, x)
x̄ = pullback_at(f, x, y, ȳ, intermediates)
```

--------------------------------

### Using @ignore_derivatives for one-liners

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/tips_for_packages.html

Using the @ignore_derivatives macro for concise one-line expressions.

```julia
@ignore_derivatives push!(things, thing)
```

--------------------------------

### Benchmark optimized sum3

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html

Performance measurement of the sum3 function after applying the rrule.

```julia
julia> @btime gradient(sum3, rand(30))
  192.818 ns (3 allocations: 784 bytes)
```

--------------------------------

### Write rules conditionally on AD features

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/ruleconfig.html

Shows how to write rrule and frule definitions that only apply when the AD has specific features, using the RuleConfig type parameter with >: syntax.

```julia
rrule(::RuleConfig{>:Feature1}, f, args...) = # rrule that should only be define for ADs with `Feature1`

frule(::RuleConfig{>:Union{Feature1,Feature2}}, f, args...) = # frule that should only be define for ADs with both `Feature1` and `Feature2`
```

--------------------------------

### Simplifying pullbacks with duck-typing

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/writing_good_rules.html

Alternatively, simplify the pullback by using duck-typing to handle both AbstractArray and AbstractThunk inputs in a single method.

```julia
function ChainRulesCore.rrule(::typeof(double_it), x)
    double_it_pullback(ȳ) = (NoTangent(), 2 .* ȳ)
    return double_it(x), double_it_pullback
end
```

--------------------------------

### Maybe-mutating accumulation with add!!

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/gradient_accumulation.html

Shows the pattern of mutating one gradient in-place to reduce allocations, and mentions that ChainRulesCore provides add!! for this purpose. The code shows ā .+= b̄ and X̄ = ā, which cuts allocations from three to two.

```julia
ā .+= b̄
X̄ = ā
```

--------------------------------

### frule for matrix multiplication (*)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Implements the forward-mode rule for matrix multiplication. The pushforward is ΔA * B + A * ΔB, following the matrix product rule. Use for Matrix inputs.

```julia
function frule(
    (_, ΔA, ΔB),
    ::typeof(*),
    A::Matrix{<:RealOrComplex},
    B::Matrix{<:RealOrComplex},
)
    Ω = A * B
    ∂Ω = ΔA * B + A * ΔB
    return (Ω, ∂Ω)
end
```

--------------------------------

### Matrix inversion rrule for inv

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Implements the rrule for matrix inversion (inv) with real or complex matrices. The pullback computes ∂A = -Ω' * ΔΩ * Ω' and returns NoTangent() for the primal output.

```julia
function rrule(::typeof(inv), A::Matrix{<:RealOrComplex})
    Ω = inv(A)
    function inv_pullback(ΔΩ)
        ∂A = -Ω' * ΔΩ * Ω'
        return (NoTangent(), ∂A)
    end
    return Ω, inv_pullback
end
```

--------------------------------

### Define AD-specific rule for add!! using ZygoteConfig

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/ruleconfig.html

Defines a ZygoteConfig subtype of RuleConfig and implements an rrule for add!! that performs addition instead of mutation.

```julia
struct ZygoteConfig <: RuleConfig{Union{}} end

rrule(::ZygoteConfig, typeof(ChainRulesCore.add!!), a, b) = a+b, Δ->(NoTangent(), Δ, Δ)
```

--------------------------------

### Compute partial derivative b̄ for X[2]

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/gradient_accumulation.html

Shows how to compute the partial derivative b̄ for X[2] as a zero matrix with the scalar sensitivity ȳ placed at index 2. This is the value representation of the inplaceable thunk.

```julia
b̄ = zeros(size(X))
b̄[2] = ȳ  # the scalar sensitivity of the `sum_first_and_second` output
```

--------------------------------

### Benchmark loop-based MSE

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html

Performance measurement of the loop-based MSE function.

```julia
julia> y = rand(30);

julia> ŷ = rand(30);

julia> @btime gradient(mse, $y, $ŷ)
  38.180 μs (993 allocations: 65.00 KiB)
```

--------------------------------

### Benchmark sin and cos separately vs sincos

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Benchmarks computing sin and cos separately (8.633 ns total) versus sincos together (6.028 ns), showing the speed advantage of sincos.

```julia
julia> using BenchmarkTools

julia> @btime sin(x) setup=(x=rand());
  3.838 ns (0 allocations: 0 bytes)

julia> @btime cos(x) setup=(x=rand());
  4.795 ns (0 allocations: 0 bytes)

julia> 3.838 + 4.795
8.633
```

```julia
julia> @btime sincos(x) setup=(x=rand());
  6.028 ns (0 allocations: 0 bytes)
```

--------------------------------

### Test the rrule with ChainRulesTestUtils

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/example.html

Verifies the implementation against finite-differences using test_rrule.

```julia
julia> using ChainRulesTestUtils

julia> test_rrule(foo_mul, Foo(rand(3, 3), 3.0), rand(3, 3))
Test Summary:                                       | Pass  Total
test_rrule: foo_mul on Foo{Float64},Matrix{Float64} |   10     10
Test.DefaultTestSet("test_rrule: foo_mul on Foo{Float64},Matrix{Float64}", Any[], 10, false, false)
```

--------------------------------

### Benchmark optimized MSE

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html

Performance measurement of the MSE function after applying the rrule.

```julia
julia> @btime gradient(mse, $y, $ŷ)
  143.697 ns (2 allocations: 672 bytes)
```

--------------------------------

### Benchmark sum3 function

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html

Performance measurement of the sum3 function.

```julia
julia> @btime gradient(sum3, rand(30))
  424.510 ns (9 allocations: 2.06 KiB)
```

--------------------------------

### rrule for logabsdet

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Implements the reverse-mode rule for logabsdet. It reuses the LU decomposition to compute inv(A). The pullback computes the gradient with respect to A, using the adjoint of the sign and the real part of the adjoint of the log determinant.

```julia
function rrule(::typeof(logabsdet), A::Matrix{<:RealOrComplex})
    # The primal function uses the lu decomposition to compute logabsdet
    # we reuse this decomposition to compute inv(A)
    F = lu(A, check = false)
    Ω = logabsdet(F)  # == logabsdet(A)
    s = last(Ω)
    function logabsdet_pullback(ΔΩ)
        (Δl, Δs) = ΔΩ
        f = conj(s) * Δs
        imagf = f - real(f)  # 0 for real A and Δs, im * imag(f) for complex A and/or Δs
        g = real(Δl) + imagf
        ∂A = g * inv(F)'  # == g * inv(A)'
        return (NoTangent(), ∂A)
    end
    return (Ω, logabsdet_pullback)
end
```

--------------------------------

### Dump DateTime structure

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/many_tangents.html

Shows the internal structure of a DateTime object using dump(now()). This illustrates the nested fields that a structural tangent must mirror.

```julia
julia> dump(now())
DateTime
  instant: UTInstant{Millisecond}
    periods: Millisecond
      value: Int64 63719890305605

```

--------------------------------

### Use named local functions for rrule pullbacks

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/writing_good_rules.html

Using named local functions for pullbacks improves stacktrace readability during debugging compared to anonymous functions.

```julia
# good:
function rrule(::typeof(foo), x)
    Y = foo(x)
    function foo_pullback(Ȳ)
        return NoTangent(), bar(Ȳ)
    end
    return Y, foo_pullback
end
#== output
julia> rrule(foo, 2)
(4, var"#foo_pullback#11"())
==#

# bad:
function rrule(::typeof(foo), x)
    return foo(x), x̄ -> (NoTangent(), bar(x̄))
end
#== output:
julia> rrule(foo, 2)
(4, var"##9#10"())
==#
```

--------------------------------

### Define function and rules for ChainRulesCore

Source: https://juliadiff.org/ChainRulesCore.jl/stable/index.html

Defines a sample function foo and its associated scalar rules using @scalar_rule.

```julia
using ChainRulesCore

function foo(x)
    a = sin(x)
    b = 0.2 + a
    c = asin(b)
    return c
end

# Define rules (alternatively get them for free via `using ChainRules`)
@scalar_rule(sin(x), cos(x))
@scalar_rule(+(x, y), (1.0, 1.0))
@scalar_rule(asin(x), inv(sqrt(1 - x^2)))
```

--------------------------------

### Refactored sin pullback using sincos

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the refactored version for sin that uses sincos and stores cx instead of x, reducing redundant computation.

```julia
function augmented_primal(::typeof(sin), x)
    y, cx = sincos(x)
    return y, PullbackMemory(sin; cx=cx)
end
(pb::PullbackMemory)(ȳ) = ȳ * pb.cx
```

--------------------------------

### Corrected rrule using ProjectTo

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/projectto.html

An improved rrule implementation that uses ProjectTo to ensure tangents are projected onto the correct subspace of the input arguments.

```julia
function rrule(::typeof(*), A::AbstractMatrix, B::AbstractMatrix)
    project_A = ProjectTo(A)
    project_B = ProjectTo(B)
    function times_pullback(ȳ)
        dA = ȳ * B'
        dB = A' * ȳ
        return NoTangent(), project_A(dA), project_B(dB)
    end
    return A * B, times_pullback
end
```

--------------------------------

### rrule

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Computes the output of a function and a pullback for reverse mode automatic differentiation. Returns nothing if no rule is defined.

```APIDOC
## rrule

### Description
Computes the output of a function and a pullback for reverse mode automatic differentiation. Returns nothing if no rule is defined.

### Method
rrule

### Signature
rrule([::RuleConfig,] f, x...)

### Parameters
- RuleConfig (optional) - AD system configuration for feature support.
- f - Function to differentiate.
- x... - Input arguments.

### Returns
- Ω - Output of f(x...).
- Pullback - A function (Ω̄₁, Ω̄₂, ...) -> (s̄elf, x̄₁, x̄₂, ...) that takes cotangents of the outputs and returns cotangents for the function itself and inputs.

### Examples
Unary input, unary output scalar function:
```julia
julia> x = rand();

julia> sinx, sin_pullback = rrule(sin, x);

julia> sinx == sin(x)
true

julia> sin_pullback(1) == (NoTangent(), cos(x))
true
```

Binary input, unary output scalar function:
```julia
julia> x, y = rand(2);

julia> hypotxy, hypot_pullback = rrule(hypot, x, y);

julia> hypotxy == hypot(x, y)
true

julia> hypot_pullback(1) == (NoTangent(), (x / hypot(x, y)), (y / hypot(x, y)))
true
```

### Notes
The optional RuleConfig option allows specifying rrules only for AD systems that support given features. If not needed, it can be omitted and the rrule without it will be hit as a fallback.

### See Also
frule, @scalar_rule, RuleConfig
```

--------------------------------

### Callable PullbackMemory API

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the API where PullbackMemory is callable, so the pullback is invoked as pb(ȳ) instead of pullback_at(pb, ȳ).

```julia
y = f(x)  # primal program
y, pb = augmented_primal(f, x)
x̄ = pb(ȳ)
```

--------------------------------

### Compute Jacobian via rrule for complex function

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/complex.html

Uses `rrule` and its pullback with inputs 1 and im to compute the Jacobian of a complex function as a real 2x2 matrix.

```julia
function jacobian_via_rrule(f,z)
    _, pullback = rrule(f,z)
    du_dx, du_dy = reim(pullback( 1)[2])
    dv_dx, dv_dy = reim(pullback(im)[2])
    return [
        du_dx  du_dy
        dv_dx  dv_dy
    ]
end
```

--------------------------------

### frule for logabsdet

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Implements the forward-mode rule for logabsdet. It reuses the LU decomposition from the primal computation to compute tr(inv(A) * ΔA). For real A, the tangent of the sign is ZeroTangent() to maintain type stability.

```julia
function frule((_, ΔA), ::typeof(logabsdet), A::Matrix{<:RealOrComplex})
    # The primal function uses the lu decomposition to compute logabsdet
    # we reuse this decomposition to compute inv(A) * ΔA
    F = lu(A, check = false)
    Ω = logabsdet(F)  # == logabsdet(A)
    b = tr(F \ ΔA)  # == tr(inv(A) * ΔA)
    s = last(Ω)
    ∂l = real(b)
    # for real A, ∂s will always be zero (because imag(b) = 0)
    # this is type-stable because the eltype is known
    ∂s = eltype(A) <: Real ? ZeroTangent() : im * imag(b) * s
    # tangents of tuples are of type Tangent{<:Tuple}
    ∂Ω = Tangent{typeof(Ω)}(∂l, ∂s)
    return (Ω, ∂Ω)
end
```

--------------------------------

### Gradient after defining rrule

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html

Shows the successful gradient evaluation after defining the rrule. The output is a tuple containing the gradient vector.

```julia
julia> gradient(addone, a)
([1.0, 1.0, 1.0],)
```

--------------------------------

### Action of b̄ on ā (b̄_add!)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/gradient_accumulation.html

Defines the action of b̄ upon ā to produce the same result as add!!(ā, b̄) or ā + b̄. It mutates ā by adding ȳ at index 2 and returns ā.

```julia
function b̄_add!(ā)
    ā[2] += ȳ
    return ā
end
```

--------------------------------

### Naive rrule for matrix multiplication

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/projectto.html

A naive implementation of matrix multiplication that fails to account for specific tangent types when inputs are abstractly typed.

```julia
function rrule(::typeof(*), A::AbstractMatrix, B::AbstractMatrix)
    function times_pullback(ȳ)
        dA = ȳ * B'
        dB = A' * ȳ
        return NoTangent(), dA, dB
    end
    return A * B, times_pullback
end
```

--------------------------------

### augmented_primal and pullback_at for sin

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Defines augmented_primal for sin using sincos to compute both y and cos(x) as intermediates, and pullback_at that uses the stored cos(x).

```julia
function augmented_primal(::typeof(sin), x)
    y, cx = sincos(x)
    return y, (; cx=cx)  # use a NamedTuple for the intermediates
end

pullback_at(::typeof(sin), x, y, ȳ, intermediates) = ȳ * intermediates.cx
```

--------------------------------

### Matrix multiplication rrule for *

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Implements the rrule for matrix multiplication (*) with real or complex matrices. The pullback returns thunks for the partial derivatives ∂A = ΔΩ * B' and ∂B = A' * ΔΩ, along with NoTangent() for the primal output.

```julia
function rrule(::typeof(*), A::Matrix{<:RealOrComplex}, B::Matrix{<:RealOrComplex})
    function times_pullback(ΔΩ)
        ∂A = @thunk(ΔΩ * B')
        ∂B = @thunk(A' * ΔΩ)
        return (NoTangent(), ∂A, ∂B)
    end
    return A * B, times_pullback
end
```

--------------------------------

### Chain of expressions: a = f(x), b = g(a), c = h(b)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/propagators.html

Shows a simple chain of expressions used to illustrate the chain rule. The pullback of g applies the chain rule from ∂c/∂b to ∂c/∂a, while the pushforward applies it from ∂a/∂x to ∂b/∂x.

```plaintext
a = f(x)
b = g(a)
c = h(b)
```

--------------------------------

### Structure of the rrule dispatch

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/example.html

Shows the dispatch pattern for the rrule function.

```julia
function ChainRulesCore.rrule(::typeof(foo_mul), foo::Foo, b::AbstractArray)
    ...
    return y, foo_mul_pullback
end
```

--------------------------------

### Compute the primal result

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/example.html

The primal computation step within the rrule.

```julia
y = foo_mul(foo, b)
```

--------------------------------

### Plot cbrt for nonfinite derivative

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/nondiff_points.html

Derivative nonfinite and same on both sides. The derivative at 0 is Inf; alternatives like nextfloat(0.0) or prevfloat(0.0) are essentially the same.

```julia
plot(cbrt)
```

--------------------------------

### Implement rrule for sum(abs2, X; dims)

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Generic reverse-mode rule for sum(abs2, X; dims). The pullback returns NoTangent for the function and abs2, and a thunked gradient for X computed as 2 .* real.(ΔΩ) .* X.

```julia
function rrule(::typeof(sum), ::typeof(abs2), X::Array{<:RealOrComplex}; dims = :)
    function sum_abs2_pullback(ΔΩ)
        ∂abs2 = NoTangent()
        ∂X = @thunk(2 .* real.(ΔΩ) .* X)
        return (NoTangent(), ∂abs2, ∂X)
    end
    return sum(abs2, X; dims = dims), sum_abs2_pullback
end
```

--------------------------------

### Optimized rrule with @thunk

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/projectto.html

An optimized version of the pullback that uses @thunk to defer projection and product evaluation until the tangent is actually used.

```julia
times_pullback(ȳ) = NoTangent(), @thunk(project_A(ȳ * B')), @thunk(project_B(A' * ȳ))
```

--------------------------------

### Pushforward function structure

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/propagators.html

Shows the general structure of a pushforward, which is part of frule. It takes Δself and Δargs as input and returns ∂y.

```julia
function pushforward(Δself, Δargs...)
    ...
    return ∂y
end
```

--------------------------------

### Define rrule for a constructor

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/writing_good_rules.html

Define an rrule for a constructor by dispatching on Type{Bar}.

```julia
function ChainRulesCore.rrule(::Type{Bar}, a)
    Bar_pullback(Δbar) = NoTangent(), Δbar.a
    return Bar(a), Bar_pullback
end
```

--------------------------------

### Calculate gradient with rrule

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/propagators.html

Computes the gradient of a function f using rrule and the returned pullback.

```julia
y, f_pullback = rrule(f, a, b, c)
∇f = f_pullback(1)  # for appropriate `1`-like seed.
s̄elf, ā, b̄, c̄ = ∇f
```

--------------------------------

### Naive sum function using getindex

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/gradient_accumulation.html

This function sums the elements of an array using getindex. It is used to illustrate the allocation problem when using value representation of derivatives, where a derivative array would be allocated for each element.

```julia
function mysum(X::Array{Float64})
    total = 0.0
    for i in eachindex(X)
        total += X[i]
    end
    return total
end
```

--------------------------------

### Define an AD's RuleConfig with feature union

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/superpowers/ruleconfig.html

Shows how an AD system declares its special features by defining a RuleConfig with a union of feature types as its type parameter.

```julia
struct MyADRuleConfig <: RuleConfig{Union{Feature1, Feature2}} end
```

--------------------------------

### Plot x^3 for standard derivative

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/nondiff_points.html

Standard case where the derivative is defined in the usual sense. At x=0 the derivative is 3*0^2=0.

```julia
plot(x->x^3)
```

--------------------------------

### Accumulation with add!!

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

The add!! function returns x+y, potentially mutating x in-place to avoid allocations.

```julia
add!!(x, y)
```

--------------------------------

### Refactored sigmoid pullback using exp

Source: https://juliadiff.org/ChainRulesCore.jl/stable/design/changing_the_primal.html

Shows the refactored version for the logistic sigmoid σ that computes exp and stores y and ex, avoiding recomputation of σ(-x).

```julia
function augmented_primal(::typeof(σ), x)
    ex = exp(x)
    y = ex/(1 + ex)
    return y, PullbackMemory(σ; y=y, ex=ex)
end
(pb::PullbackMemory{typeof(σ)})(ȳ) = ȳ * pb.y/(1 + pb.ex)
```

--------------------------------

### Compute holomorphic derivative via rrule

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/complex.html

For holomorphic f, the pullback returns (f′(z))∗Δf; conjugating the result yields the derivative.

```julia
function holomorphic_derivative_via_rrule(f,z)
    fz, pullback = rrule(f,z)
    dself, conj_df_dz = pullback(1)
    return conj(conj_df_dz)
end
```

--------------------------------

### Define rrule for sum3

Source: https://juliadiff.org/ChainRulesCore.jl/stable/rule_author/which_functions_need_rules.html

An rrule to optimize the gradient calculation for sum3.

```julia
function ChainRulesCore.rrule(::typeof(sum3), a)
    y = sum3(a)
    function sum3_pullback(ȳ)
        grad = zeros(length(a))
        grad[1:3] .+= ȳ
        return NoTangent(), grad
    end
    return y, sum3_pullback
end
```

--------------------------------

### Define a reverse-only AD RuleConfig and rrule_via_ad

Source: https://juliadiff.org/ChainRulesCore.jl/stable/ad_author/call_back_into_ad.html

Declares a custom RuleConfig subtype that supports reverse mode but not forward mode, and defines the required rrule_via_ad method. The function body is omitted in the source, indicated by ellipsis.

```julia
struct MyReverseOnlyADRuleConfig <: RuleConfig{Union{HasReverseMode, NoForwardsMode}} end

function ChainRulesCore.rrule_via_ad(::MyReverseOnlyADRuleConfig, f, args...)
    ...
    return y, pullback
end
```

--------------------------------

### frule for sylvester

Source: https://juliadiff.org/ChainRulesCore.jl/stable/maths/arrays.html

Forward-mode rule for sylvester. Computes X = sylvester(A, B, C) and returns the pushforward as sylvester(A, B, ΔC + ΔA*X + X*ΔB).

```julia
function frule((_, ΔA, ΔB, ΔC), ::typeof(sylvester), A, B, C)
    X = sylvester(A, B, C)
    return X, sylvester(A, B, ΔC + ΔA * X + X * ΔB)
end
```

--------------------------------

### RuleConfig configuration

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Defines the configuration for rules, where T is a Union of special traits.

```julia
RuleConfig{T}
```

```julia
# only define rrule for `pop!` on AD systems where mutation is supported.
rrule(::RuleConfig{>:SupportsMutation}, typeof(pop!), ::Vector) = ...

# this definition of map is for any AD that defines a forwards mode
rrule(conf::RuleConfig{>:HasForwardsMode}, typeof(map), ::Vector) = ...

# this definition of map is for any AD that only defines a reverse mode.
# It is not as good as the rrule that can be used if the AD defines a forward-mode as well.
rrule(conf::RuleConfig{>:Union{NoForwardsMode, HasReverseMode}}, typeof(map), ::Vector) = ...
```

--------------------------------

### ProjectTo tangent projection

Source: https://juliadiff.org/ChainRulesCore.jl/stable/api.html

Projects a tangent dx onto a specific tangent space T, typically enforcing p(dx)::T.

```julia
(p::ProjectTo{T})(dx)
```