API reference

DifferentiableExpectations

A Julia package for differentiating through expectations with Monte-Carlo estimates.

Exports

• DifferentiableExpectation
• FixedAtomsProbabilityDistribution
• Reinforce
• Reparametrization
• empirical_distribution

DifferentiableExpectation{t}

Abstract supertype for differentiable parametric expectations E : θ -> 𝔼[f(X)] where X ∼ p(θ), whose value and derivative are approximated with Monte-Carlo averages.

Subtypes

• Reinforce
• Reparametrization

Calling behavior

(E::DifferentiableExpectation)(θ...; kwargs...)

Return a Monte-Carlo average (1/S) ∑f(xᵢ) where the xᵢ ∼ p(θ) are iid samples.

Type parameters

• threaded::Bool: specifies whether the sampling should be performed in parallel

Required fields

• f: The function applied inside the expectation.
• dist_constructor: The constructor of the probability distribution.
• rng::AbstractRNG: The random number generator.
• nb_samples::Integer: The number of Monte-Carlo samples.
• seed::Union{Nothing,Integer}: The seed for the random number generator, reset before each call. Set to nothing for no seeding.

The field dist_constructor must be a callable such that dist_constructor(θ...) generates an object dist that corresponds to p(θ). The resulting object dist needs to satisfy:

• the Random API for sampling with rand(rng, dist)
• the DensityInterface.jl API for loglikelihoods with logdensityof(dist, x)

FixedAtomsProbabilityDistribution{threaded}

A probability distribution with finite support and fixed atoms.

Whenever its expectation is differentiated, only the weights are considered active, whereas the atoms are considered constant.

Example

julia> using DifferentiableExpectations, Statistics, Zygote

julia> using DifferentiableExpectations: atoms, weights

julia> dist = FixedAtomsProbabilityDistribution([2, 3], [0.4, 0.6]);

julia> atoms(map(abs2, dist))
2-element Vector{Int64}:
 4
 9

julia> weights(map(abs2, dist))
2-element Vector{Float64}:
 0.4
 0.6

julia> mean(abs2, dist)
7.0

julia> gradient(mean, abs2, dist)[2]
(atoms = nothing, weights = [4.0, 9.0])

Constructor

FixedAtomsProbabilityDistribution(
    atoms::AbstractVector,
    weights::AbstractVector=uniform_weights(atoms);
    threaded=false
)

Fields

• atoms::AbstractVector

• weights::AbstractVector{<:Real}

Reinforce{threaded} <: DifferentiableExpectation{threaded}

Differentiable parametric expectation F : θ -> 𝔼[f(X)] where X ∼ p(θ) using the REINFORCE (or score function) gradient estimator:

∂F(θ) = 𝔼[f(X) ∇₂logp(X,θ)ᵀ]

Example

using DifferentiableExpectations, Distributions, Zygote

E = Reinforce(exp, Normal; nb_samples=10^5)
E_true(μ, σ) = mean(LogNormal(μ, σ))

μ, σ = 0.5, 1,0
∇E, ∇E_true = gradient(E, μ, σ), gradient(E_true, μ, σ)
isapprox(collect(∇E), collect(∇E_true); rtol=1e-1)

# output

true

Constructor

Reinforce(
    f,
    dist_constructor,
    dist_logdensity_grad=nothing;
    rng=Random.default_rng(),
    nb_samples=1,
    threaded=false,
    seed=nothing
)

Fields

• f::Any: function applied inside the expectation

• dist_constructor::Any: constructor of the probability distribution (θ...) -> p(θ)

• dist_logdensity_grad::Any: either nothing or a parameter gradient callable (x, θ...) -> ∇₂logp(x, θ)

• rng::Random.AbstractRNG: random number generator

• nb_samples::Int64: number of Monte-Carlo samples

• seed::Union{Nothing, Int64}: seed for the random number generator, reset before each call. Set to nothing for no seeding.

See also

• DifferentiableExpectation

Reparametrization{threaded} <: DifferentiableExpectation{threaded}

Differentiable parametric expectation F : θ -> 𝔼[f(X)] where X ∼ p(θ) using the reparametrization (or pathwise) gradient estimator: if X = g(Z,θ) where Z ∼ q then

∂F(θ) = 𝔼_q[∂f(g(Z,θ)) ∂₂g(Z,θ)ᵀ]

Example

using DifferentiableExpectations, Distributions, Zygote

E = Reparametrization(exp, Normal; nb_samples=10^4)
E_true(μ, σ) = mean(LogNormal(μ, σ))

μ, σ = 0.5, 1,0
∇E, ∇E_true = gradient(E, μ, σ), gradient(E_true, μ, σ)
isapprox(collect(∇E), collect(∇E_true); rtol=1e-1)

# output

true

Constructor

Reparametrization(
    f,
    dist_constructor,
    rng=Random.default_rng(),
    nb_samples=1,
    threaded=false,
    seed=nothing
)

Fields

• f::Any: function applied inside the expectation

• dist_constructor::Any: constructor of the probability distribution (θ...) -> p(θ)

• rng::Random.AbstractRNG: random number generator

• nb_samples::Int64: number of Monte-Carlo samples

• seed::Union{Nothing, Int64}: seed for the random number generator, reset before each call. Set to nothing for no seeding.

See also

• DifferentiableExpectation

empirical_distribution(E::DifferentiableExpectation, θ...; kwargs...)

Return a uniform FixedAtomsProbabilityDistribution over {f(x₁), ..., f(xₛ)}, where the xᵢ ∼ p(θ) are iid samples.

FixKwargs(f, kwargs)

Callable struct that fixes the keyword arguments of f to kwargs..., and only accepts positional arguments.

TransformedDistribution

Represent the probability distribution p of a random variable X ∼ p with a transformation X = T(Z) where Z ∼ q.

Fields

• base_dist::Any: the distribution q that gets transformed into p

• transformation::Any: the transformation function T

map(f, dist::FixedAtomsProbabilityDistribution)

Apply f to the atoms of dist, leave the weights unchanged.

rand(rng, dist::TransformedDistribution)

Sample from dist by applying dist.transformation to dist.base_dist.

rand(rng, dist::FixedAtomsProbabilityDistribution)

Sample from the atoms of dist with probability proportional to their weights.

atoms(dist::FixedAtomsProbabilityDistribution)

Get the vector of atoms of a distribution.

empirical_predistribution(E::DifferentiableExpectation, θ...)

Return a uniform FixedAtomsProbabilityDistribution over {x₁, ..., xₛ}, where the xᵢ ∼ p(θ) are iid samples.

maybe_eachcol(x::AbstractMatrix)

Return eachcol(x).

maybe_eachcol(x::AbstractVector)

Return x.

mymap(::Val{threaded}, args...)

Apply either tmap(args...) or map(args...) depending on the value of threaded.

mymapreduce(::Val{threaded}, args...)

Apply either tmapreduce(args...) or mapreduce(args...) depending on the value of threaded.

reparametrize(dist)

Turn a probability distribution p into a TransformedDistribution (q, T) such that the new distribution q does not depend on the parameters of p. These parameters are encoded (closed over) in the transformation function T.

weights(dist::FixedAtomsProbabilityDistribution)

Get the vector of weights of a distribution.

mean(f, dist::FixedAtomsProbabilityDistribution)

Shortcut for mean(map(f, dist)).

mean(dist::FixedAtomsProbabilityDistribution)

Compute the expectation of dist, i.e. the sum of all atoms multiplied by their respective weights.