Skip to content

Reusing a GradientConfig/JacobianConfig with a differently structured input silently computes wrong derivatives #842

Description

@devmotion

Two independent bugs in the structure-aware seeding introduced in #739. Both reproduce on master
(v1.4.5, 569af35) and on the branch of #840, which addresses neither. I intend to fix both in #840.

1. Reusing a config with a differently structured input is silently wrong

seed! and seed_zero_partials! take the positions to seed from the config's work buffer
(structural_eachindex(duals, x)), while extraction takes them from somewhere else — from result
on master, from x after #840. When a config is reused with an input of the same length but a
different structure, the two sets disagree and nothing complains:

using ForwardDiff, LinearAlgebra
using ForwardDiff: GradientConfig, Chunk, gradient

f(z) = sum(abs2, z) / 2                                    # ∇f(z) == z
U = UpperTriangular([1.0 2.0 3.0; 0.0 4.0 5.0; 0.0 0.0 6.0])
A = reshape(1.0:9.0, 3, 3)

cfgU = GradientConfig(f, U, Chunk{2}())
cfgA = GradientConfig(f, A, Chunk{2}())

vec(gradient(f, U, cfgU))   # [1, 0, 0, 2, 4, 0, 3, 5, 6]   ✅ correct
vec(gradient(f, A, cfgA))   # [1, 2, 3, 4, 5, 6, 7, 8, 9]   ✅ correct

vec(gradient(f, U, cfgA))   # [1, 0, 0, 0, 0, 0, 2, 4, 0]   ❌ should be [1, 0, 0, 2, 4, 0, 3, 5, 6]
vec(gradient(f, A, cfgU))   # [1, 4, 5, 7, 8, 9, 0, 0, 0]   ❌ should be [1, 2, 3, 4, 5, 6, 7, 8, 9]

Affected:

  • both modes — with this x, Chunk{2} is chunk mode and Chunk{6} is vector mode; both
    produce exactly the wrong values above;

  • gradient, jacobian and hessianhessian through both of its sub-configs, e.g.
    diag(hessian(f, A, HessianConfig(f, U, Chunk{2}()))) is [1, 0, 0, 0, 0, 0, 0, 0, 0] where it
    should be all ones;

  • all three structured wrappers (LowerTriangular, UpperTriangular, Diagonal), in either
    direction, including structured→structured:

    L = LowerTriangular([1.0 0.0 0.0; 2.0 4.0 0.0; 3.0 5.0 6.0])
    vec(gradient(f, U, GradientConfig(f, L, Chunk{2}())))
    # [1, 0, 0, 0, 0, 0, 4, 0, 6]   ❌ should be [1, 0, 0, 2, 4, 0, 3, 5, 6]

    Note that these two agree on size and on structural_length (both 6), so no count- or
    size-based check can catch this pair.

…and the work buffer is left partially uninitialized

The reuse does not merely mislabel the output. The buffer entries that the mismatched structure never
visits are never written at all, and the target function reads them. Above, the sweep for an
UpperTriangular input seeds and clears 6 positions of a work buffer that has 9. With Float64 that
is silent garbage; a non-bits element type surfaces it:

Ub = UpperTriangular(big.([1.0 2.0 3.0; 0.0 4.0 5.0; 0.0 0.0 6.0]))
Ab = big.(collect(reshape(1.0:9.0, 3, 3)))

gradient(f, Ub, GradientConfig(f, Ab, Chunk{2}()))
# ERROR: UndefRefError: access to undefined reference

This is also why the fix cannot simply be "take the seeding positions from x as well": the buffer
entries outside x's structure would then never be initialized. The config has to be checked against
the input instead.

2. Independent: Base._unsetindex! is linear-index-only, so the unassigned-entry branch of seed! mostly does not work

seed! and seed_zero_partials! have a branch for non-isbitstype value types that propagates
unassigned entries of x into the buffer via Base._unsetindex!(duals, idx). Base has no
_unsetindex!(::AbstractArray, ::CartesianIndex) method at all — not even for Array — and its
AbstractArray fallback for a linear index calls itself forever, since to_index(::Int) is the
identity (_unsetindex!(A, i::Integer) = _unsetindex!(A, to_index(i)), abstractarray.jl:1482). So
that branch only ever worked for a dense Array. With one unassigned entry at a structural position
and Chunk{2}():

x, element type BigFloat ForwardDiff.gradient(f, x)
Matrix UndefRefError — raised by f reading the hole, i.e. seeding worked
adjoint(Matrix) MethodError: no method matching _unsetindex!(::Matrix{Dual{…}}, ::CartesianIndex{2})
PermutedDimsArray same MethodError
UpperTriangular, LowerTriangular MethodError: no method matching _unsetindex!(::UpperTriangular{Dual{…}}, ::CartesianIndex{2})
Diagonal StackOverflowError

Between them the two failure modes cover the whole of Base's _unsetindex! surface, which has
exactly two concrete methods, for Array and for Memory:

  • No CartesianIndex method exists at all. Rows 2 and 3 are the informative ones here: the buffer
    is a plain dense Matrix, because similar does not preserve those wrappers, and it still fails.
    Cartesian indices reach the seeding loop whenever either array of eachindex(duals, x) is
    IndexCartesian, so this is not confined to LinearAlgebra's wrapper types.
  • The AbstractArray fallback for a linear index recurses forever. Diagonal is the one
    structured case whose positions are already linear (structural_eachindex(::Diagonal, _) returns
    diagind(x)), so it gets past the first problem and straight into the second.

Proposed fix (in #840)

  1. Store the structural positions in the config, as an indexable vector of linear indices, built
    from the work buffer the config owns. This also removes the Iterators.drop walk that currently
    re-traverses the lazy triangular position iterators from the front to reach each chunk — three
    times per middle chunk in the gradient sweep, which measures at 35–49% of gradient! for an
    UpperTriangular input.
  2. Validate the config against the input at every API entry point, next to checktag. Comparing the
    structural kind is O(1) and a compile-time constant, so it can run on every call; a size or
    count comparison is not sufficient, per the LowerTriangular/UpperTriangular case above.
  3. Unset through the array that actually stores the entry, with a linear index.

Making the stored positions linear is what makes (3) expressible, and it fixes the non-bits path for
adjoint, transpose, PermutedDimsArray and non-strided views as well as for the three wrappers
similar preserves.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions