Derivatives

A model may be differentiated, to fit hyperparameters or to run a gradient-based sampler. Its derivatives then pass through its rules. @test_rule_derivatives checks that they do. It builds the rule's inputs from a parameter θ, reduces the rule's output to a number with summary, and compares ForwardDiff's derivative with a central finite difference.

The rule under test is the belief propagation rule towards out of the Gaussian node of the first tutorial:

using MessagePassingRulesBase, MessagePassingRulesTestUtils
using BayesBase, ExponentialFamily, Distributions, Test

struct Gaussian end
Gaussian(μ, v) = NormalMeanVariance(μ, v)

@define_factor_node(node = Gaussian, type = Stochastic, interfaces = [:out, :μ, :v])

@define_message_update_rule(
    node = Gaussian, target = :out,
    args = (m[:μ]::NormalMeanVariance, m[:v]::PointMass),
    logscale = 0,
    body = (args) -> NormalMeanVariance(mean(args.m[:μ]), var(args.m[:μ]) + mean(args.m[:v])),
)

A scalar θ sets the mean of the message on μ and the variance v. A vector θ sets the mean and the variance of the message on μ, and the check compares the gradient:

@testset "Gaussian: derivatives" begin
    @test_rule_derivatives(
        node = Gaussian, target = :out,
        inputs = θ -> (m = (μ = NormalMeanVariance(θ, 1.0), v = PointMass(θ^2)),),
        at = 1.5, summary = d -> mean(d) + var(d),
    )
    @test_rule_derivatives(
        node = Gaussian, target = :out,
        inputs = θ -> (m = (μ = NormalMeanVariance(θ[1], θ[2]), v = PointMass(0.5)),),
        at = [1.0, 2.0], summary = var,
    )
end

Test Summary:         | Pass  Total  Time
Gaussian: derivatives |    2      2  1.5s

A rule that converts its inputs to Float64, or strips the dual part of a number, gives a wrong or zero derivative and fails. An in-place rule is checked twice, through rule and through rule!. Reading a failure shows what a failing check reports.

API

MessagePassingRulesTestUtils.@test_rule_derivatives — Macro

@test_rule_derivatives(node = ..., target = ..., inputs = θ -> (m = (...),), at = ..., ...)

Check that derivatives propagate through the message rule of a node towards one of its interfaces, by ForwardDiff against a central finite difference. It takes the arguments and keywords of test_rule_derivatives, where they are described, all by name.

The macro takes keyword arguments only, written name = value, and passes them to test_rule_derivatives; node and target are required, and a positional argument or a missing required keyword is an error when the macro expands. Every check is reported against the line of the macro call, so a failure points at the call that made it.

Examples

@test_rule_derivatives(
    node = NormalMeanVariance, target = :out,
    inputs = θ -> (m = (μ = NormalMeanVariance(θ, 1.0), v = PointMass(θ^2)),),
    at = 1.5, summary = d -> mean(d) + var(d),
)

source

MessagePassingRulesTestUtils.test_rule_derivatives — Function

test_rule_derivatives(node, target; inputs, at, summary = mean, algorithm = default_algorithm(node),
    step = 1e-6, rtol = 1e-5, atol = 1e-8) -> nothing

Check that derivatives propagate through the message rule of node towards target: the derivative of θ -> summary(rule(inputs(θ))) by ForwardDiff must agree with a central finite difference at θ = at. A rule that drops the dual part of its inputs fails the comparison. One that cannot take dual numbers, or throws at a perturbed θ, fails as rule_runs with the exception; no rule for the inputs, at at or at a dual or perturbed θ, fails as rule_found with the near misses. Nothing is thrown for a rule, and the checks after it still run. @test_rule_derivatives is the same check written with keywords, and reports failures at its own line.

The rule is resolved for the inputs at each θ, and the one found at at is recorded as selected for check_rule_coverage. An in-place rule is checked a second time through rule!, into a buffer its preallocate makes from the dual inputs, so both paths are covered. The rule runs with an empty RuleContext, and its log scale is not checked.

Arguments

node: the node, as declared with @define_factor_node.
target: the interface the message goes to: :out, or (:m, k) for member k of the group m.

Keywords

inputs: a function of θ returning a NamedTuple with the rule's inbound messages m and marginals q: θ -> (m = (out = Normal(θ, 2.0), σ = PointMass(θ^2)),). Required. Joint marginals of clusters are not supported.
at: the point to differentiate at, a number or a vector. Required.
summary: a function from the rule's output to a number, whose derivative is compared. Default: mean.
algorithm: the algorithm value to run under. Default: default_algorithm(node).
step: the finite-difference step. Default: 1e-6.
rtol, atol: the tolerances of the comparison, element by element for a vector at. Defaults: 1e-5 and 1e-8.
source: the line checks are reported against. Default: LineNumberNode(0, :unknown); the macro form passes its own line.

source