JuliaDiff / JuliaDiff/ForwardDiff.jl

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

Open
#842 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Julia
Stars
1k
Forks
160
PR merge metrics
No merged PRs in 30d

Description

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.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the seed!, seed_zero_partials!, GradientConfig, JacobianConfig, and checktag entry points, then run the supplied structured-input and non-bits examples on master. Done means reused configs reject differently structured inputs and unassigned entries are handled across the listed array types without wrong derivatives, undefined references, method errors, or recursion.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.