|
| 1 | +module SeedTest |
| 2 | + |
| 3 | +import ForwardDiff |
| 4 | +using ForwardDiff: Partials |
| 5 | +using LinearAlgebra |
| 6 | +using Test |
| 7 | + |
| 8 | +include("utils.jl") |
| 9 | + |
| 10 | +# The windowed `seed_zero_partials!` is only ever called to clear a chunk that was just seeded, so |
| 11 | +# clearing *too much* is harmless and no test written against the public API can distinguish a |
| 12 | +# correctly bounded implementation from an unbounded one. These tests pin the window down directly: |
| 13 | +# they seed every structural position with a marker whose partials are all nonzero, clear a window, |
| 14 | +# and check exactly which positions lost their marker. |
| 15 | +# |
| 16 | +# The expected structural index sets are written out by hand rather than obtained from |
| 17 | +# `structural_eachindex`, so a bug in that iterator cannot hide inside the assertions depending on |
| 18 | +# it; one test ties the two together. Order is significant: `index` and `count` are positions along |
| 19 | +# the sequence, not array indices. The sets are heterogeneous by design — `Vector` and `Diagonal` |
| 20 | +# enumerate linear indices (the latter via `diagind`), `UpperTriangular` enumerates `CartesianIndex` |
| 21 | +# in column-major order. |
| 22 | +const SEED_CASES = ( |
| 23 | + (rand(10), collect(1:10)), |
| 24 | + (UpperTriangular(rand(5, 5)), [CartesianIndex(i, j) for j in 1:5 for i in 1:j]), |
| 25 | + (Diagonal(rand(6, 6)), collect(1:7:36)), |
| 26 | +) |
| 27 | + |
| 28 | +# Positions within `sidx` whose partials are zero. |
| 29 | +zeroed_positions(duals, sidx) = |
| 30 | + [i for (i, idx) in enumerate(sidx) if iszero(ForwardDiff.partials(duals[idx]))] |
| 31 | + |
| 32 | +# Compares over *every* index of `x`, not just the structural ones, so a bug misplacing values |
| 33 | +# outside the structural set is visible. Off-structure reads are safe: the wrapper types return |
| 34 | +# `zero(Dual)` without touching the (uninitialized) parent storage. |
| 35 | +values_match(duals, x) = all(idx -> ForwardDiff.value(duals[idx]) == x[idx], eachindex(x)) |
| 36 | + |
| 37 | +function fill_marker!(duals, x, sidx, marker) |
| 38 | + D = eltype(duals) |
| 39 | + for idx in sidx |
| 40 | + duals[idx] = D(x[idx], marker) |
| 41 | + end |
| 42 | + return duals |
| 43 | +end |
| 44 | + |
| 45 | +@testset "seed_zero_partials!: $(nameof(typeof(x)))" for (x, sidx) in SEED_CASES |
| 46 | + cfg = ForwardDiff.GradientConfig(nothing, x, ForwardDiff.Chunk{3}()) |
| 47 | + duals, seeds = cfg.duals, cfg.seeds |
| 48 | + N = ForwardDiff.npartials(eltype(duals)) |
| 49 | + marker = Partials(ntuple(i -> Float64(i), N)) |
| 50 | + nstruct = length(sidx) |
| 51 | + |
| 52 | + # everything below counts positions along `sidx`, so pin it to the implementation once |
| 53 | + @test collect(ForwardDiff.structural_eachindex(duals, x)) == sidx |
| 54 | + @test ForwardDiff.structural_length(x) == nstruct |
| 55 | + |
| 56 | + # `count` defaults to N |
| 57 | + fill_marker!(duals, x, sidx, marker) |
| 58 | + ForwardDiff.seed_zero_partials!(duals, x, 4) |
| 59 | + @test zeroed_positions(duals, sidx) == collect(4:(4 + N - 1)) |
| 60 | + @test values_match(duals, x) |
| 61 | + |
| 62 | + # an explicit `count` narrows the window; a `count` overrunning the end is clamped by |
| 63 | + # `Iterators.take` rather than throwing; a zero-width window is a no-op, which is what makes |
| 64 | + # `xlen - N` safe as the `count` of chunk mode's tail clear |
| 65 | + @testset "index=$index count=$count" for (index, count, expected) in |
| 66 | + ((4, 2, 4:5), |
| 67 | + (nstruct - 1, N, (nstruct - 1):nstruct), |
| 68 | + (1, 0, 1:0)) |
| 69 | + fill_marker!(duals, x, sidx, marker) |
| 70 | + ForwardDiff.seed_zero_partials!(duals, x, index, count) |
| 71 | + @test zeroed_positions(duals, sidx) == collect(expected) |
| 72 | + @test values_match(duals, x) |
| 73 | + end |
| 74 | + |
| 75 | + # the 2-arg form clears every structural position |
| 76 | + fill_marker!(duals, x, sidx, marker) |
| 77 | + ForwardDiff.seed_zero_partials!(duals, x) |
| 78 | + @test zeroed_positions(duals, sidx) == collect(1:nstruct) |
| 79 | + @test values_match(duals, x) |
| 80 | + |
| 81 | + # `seed!` and `seed_zero_partials!` must agree on what "the chunk at `index`" is, or chunk mode |
| 82 | + # would leave stale seeds behind. `duals` enters each iteration fully cleared. |
| 83 | + @testset "round-trips seed! at index=$index" for index in unique((1, 4, nstruct - N + 1)) |
| 84 | + ForwardDiff.seed!(duals, x, index, seeds) |
| 85 | + @test zeroed_positions(duals, sidx) == |
| 86 | + [i for i in 1:nstruct if !(index <= i <= index + N - 1)] |
| 87 | + ForwardDiff.seed_zero_partials!(duals, x, index) |
| 88 | + @test zeroed_positions(duals, sidx) == collect(1:nstruct) |
| 89 | + @test values_match(duals, x) |
| 90 | + end |
| 91 | +end |
| 92 | + |
| 93 | +end # module |
0 commit comments