ChainRulesCore.jl

https://juliadiff.org/ChainRulesCore.jl/stable/
Admin
ChainRulesCore.jl provides an AD-independent system for defining automatic differentiation rules and...
Tokens:70,868
Snippets:499
Trust Score:8.8
Update:1 month ago
Show doc for...
Tokens:
Raw
### @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(ȳ) = ...
    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 ȧ.x and ḃ.x, and returns the primal and its tangent.

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

    a[] *= 2
    ȧ.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
```

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

### 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
```

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

### 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)
```

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

### 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
```

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

### 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
```

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

### 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
```

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

### 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
```

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

### 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)
```

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

### 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)
```

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

### 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
```

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

### 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)
```

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

### 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_ẏ = 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
```

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

### 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
```

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

### 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(ȳ) = (NoTangent(), 2 .* ȳ)
    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̄ = ā
```

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

### 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)
```
Tokens:
Raw