diff --git a/julia/Project.toml b/julia/Project.toml index 510cef8..3e7ee5b 100644 --- a/julia/Project.toml +++ b/julia/Project.toml @@ -33,5 +33,5 @@ Random = "1" Reexport = "1" Statistics = "1" XAIBase = "4" -Zygote = "0.6" +Zygote = "0.6, 0.7" julia = "1.10" diff --git a/julia/src/SmoothedDifferentiation.jl b/julia/src/SmoothedDifferentiation.jl index ced0066..ef2d282 100644 --- a/julia/src/SmoothedDifferentiation.jl +++ b/julia/src/SmoothedDifferentiation.jl @@ -11,7 +11,7 @@ using ProgressMeter: Progress, next! using NNlib: relu, ∇maxpool, maxpool, upsample_nearest, σ, softplus using Zygote: pullback -import ChainRulesCore: rrule, NoTangent +import ChainRulesCore: rrule, NoTangent, unthunk using Flux: Flux diff --git a/julia/src/vejp/maxpool.jl b/julia/src/vejp/maxpool.jl index ed7a3b1..8cc9a2c 100644 --- a/julia/src/vejp/maxpool.jl +++ b/julia/src/vejp/maxpool.jl @@ -10,7 +10,7 @@ reset_counts!(m::MaxPoolAccumulator) = fill!(m.count, 0) samplingmode!(m::MaxPoolAccumulator, mode::Bool) = (m.sampling = mode; m) # Forward pass -function (m::MaxPoolAccumulator)(x) +function (m::MaxPoolAccumulator)(x::AbstractArray) pool = m.layer pdims = Flux.PoolDims(x, pool.k; padding = pool.pad, stride = pool.stride) y = maxpool(x, pdims) @@ -29,10 +29,10 @@ Flux.@layer MaxPoolAccumulator trainable = () # Custom VJP computing VeJP function rrule(m::MaxPoolAccumulator, x) y = m(x) - function modified_maxpool_pullback(ȳ) - ȳ_expanded = upsample_nearest(ȳ, m.layer.stride) + function modified_maxpool_pullback(ȳ) + ȳ_expanded = upsample_nearest(unthunk(ȳ), m.layer.stride) J = convert.(Float32, m.count) / m.n - x̄ = J .* ȳ_expanded + x̄ = J .* ȳ_expanded return (NoTangent(), x̄) end return y, modified_maxpool_pullback diff --git a/julia/test/Project.toml b/julia/test/Project.toml index a1d4bbc..e2fcafa 100644 --- a/julia/test/Project.toml +++ b/julia/test/Project.toml @@ -10,3 +10,4 @@ NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" XAIBase = "9b48221d-a747-4c1b-9860-46a1d8ba24a7" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/julia/test/runtests.jl b/julia/test/runtests.jl index ffbd255..585a0d8 100644 --- a/julia/test/runtests.jl +++ b/julia/test/runtests.jl @@ -11,6 +11,10 @@ using JET JET.test_package(SmoothedDifferentiation; target_defined_modules = true) end + @testset "VEJP numerical tests" begin + include("test_vejp.jl") + end + @testset "Flux" begin @testset "VGG preparation tests" begin include("test_preparation_flux.jl") diff --git a/julia/test/test_vejp.jl b/julia/test/test_vejp.jl new file mode 100644 index 0000000..6bfb4bf --- /dev/null +++ b/julia/test/test_vejp.jl @@ -0,0 +1,106 @@ +using SmoothedDifferentiation +using SmoothedDifferentiation: ReluAccumulator, MaxPoolAccumulator, reset_counts! +using Test +using Flux: Flux, MaxPool +using NNlib: relu +using Zygote: pullback + +@testset "ReLU VEJP" begin + # 5 inputs total: 4 stats inputs + 1 pullback call (which also counts). + # Position: 1 2 3 4 5 + # Input 1: [1, 0, 0, 0, 0] → count += [1, 0, 0, 0, 0] + # Input 2: [1, 1, 0, 0, 0] → count += [1, 1, 0, 0, 0] + # Input 3: [1, 1, 1, 0, 0] → count += [1, 1, 1, 0, 0] + # Input 4: [1, 1, 1, 1, 0] → count += [1, 1, 1, 1, 0] + # Pullback input: [1, 1, 1, 1, 1] → count += [1, 1, 1, 1, 1] + # Total count: [5, 4, 3, 2, 1], n = 5 + # Expected VEJP: count / n = [1.0, 0.8, 0.6, 0.4, 0.2] + + layer = ReluAccumulator(; count = zeros(Int, 5)) + + stats_inputs = [ + Float32[1, 0, 0, 0, 0], + Float32[1, 1, 0, 0, 0], + Float32[1, 1, 1, 0, 0], + Float32[1, 1, 1, 1, 0], + ] + for x in stats_inputs + layer(x) + end + + # pullback also triggers a forward pass that increments count and n + test_input = Float32[1, 1, 1, 1, 1] + _, vejp_fn = pullback(layer, test_input) + grad = only(vejp_fn(ones(Float32, 5))) + + @test grad ≈ Float32[1.0, 0.8, 0.6, 0.4, 0.2] +end + +@testset "MaxPool k=2 s=2 VEJP" begin + # 10 inputs total (WHCN layout): 9 stats inputs + 1 pullback call. + # kernel_size=2, stride=2 on 2x2x1x1 inputs. + # The pullback input has a clear max at (1,1) to avoid tie-breaking differences. + # Total count: [4 2; 3 1;;;;], n = 10 + # Expected VEJP: [0.4 0.2; 0.3 0.1] + + pool = MaxPool((2, 2); stride = (2, 2)) + count = zeros(Int, 2, 2, 1, 1) + layer = MaxPoolAccumulator(; layer = pool, count = count) + + stats_inputs = [ + Float32[1 0; 0 0;;;;], # 3x max at (1,1) + Float32[1 0; 0 0;;;;], + Float32[1 0; 0 0;;;;], + Float32[0 0; 1 0;;;;], # 3x max at (2,1) + Float32[0 0; 1 0;;;;], + Float32[0 0; 1 0;;;;], + Float32[0 1; 0 0;;;;], # 2x max at (1,2) + Float32[0 1; 0 0;;;;], + Float32[0 0; 0 1;;;;], # 1x max at (2,2) + ] + for x in stats_inputs + layer(x) + end + + # pullback also triggers a forward pass + test_input = Float32[4 1; 1 1;;;;] + _, vejp_fn = pullback(layer, test_input) + grad = only(vejp_fn(ones(Float32, 1, 1, 1, 1))) + + @test grad ≈ Float32[0.4 0.2; 0.3 0.1;;;;] +end + +@testset "MaxPool k=3 s=3 VEJP" begin + # 10 inputs total (WHCN layout): 9 stats inputs + 1 pullback call. + # kernel_size=3, stride=3 on 3x3x1x1 inputs (single pool). + # The pullback input has a clear max at (1,1). + # Total count: [4 0 1; 0 3 0; 2 0 0;;;;], n = 10 + # Expected VEJP: [0.4 0.0 0.1; 0.0 0.3 0.0; 0.2 0.0 0.0] + + pool = MaxPool((3, 3); stride = (3, 3)) + count = zeros(Int, 3, 3, 1, 1) + layer = MaxPoolAccumulator(; layer = pool, count = count) + + # WHCN layout: [w, h, c, n]. Position (w,h) maps to Python (h-1, w-1). + stats_inputs = [ + Float32[1 0 0; 0 0 0; 0 0 0;;;;], # 3x max at (1,1) → Python (0,0) + Float32[1 0 0; 0 0 0; 0 0 0;;;;], + Float32[1 0 0; 0 0 0; 0 0 0;;;;], + Float32[0 0 0; 0 0 0; 1 0 0;;;;], # 2x max at (3,1) → Python (0,2) + Float32[0 0 0; 0 0 0; 1 0 0;;;;], + Float32[0 0 0; 0 1 0; 0 0 0;;;;], # 3x max at (2,2) → Python (1,1) + Float32[0 0 0; 0 1 0; 0 0 0;;;;], + Float32[0 0 0; 0 1 0; 0 0 0;;;;], + Float32[0 0 1; 0 0 0; 0 0 0;;;;], # 1x max at (1,3) → Python (2,0) + ] + for x in stats_inputs + layer(x) + end + + # pullback also triggers a forward pass (adds 1 to (1,1)) + test_input = Float32[9 1 1; 1 1 1; 1 1 1;;;;] + _, vejp_fn = pullback(layer, test_input) + grad = only(vejp_fn(ones(Float32, 1, 1, 1, 1))) + + @test grad ≈ Float32[0.4 0.0 0.1; 0.0 0.3 0.0; 0.2 0.0 0.0;;;;] +end diff --git a/python/tests/test_maxpool_vejp.py b/python/tests/test_maxpool_vejp.py index 6e220ac..f4b04de 100644 --- a/python/tests/test_maxpool_vejp.py +++ b/python/tests/test_maxpool_vejp.py @@ -5,43 +5,90 @@ from smoothdiff_torch.smoothdiff import _SmoothMaxPool2d -def test_maxpool2d_vejp(): - """Test the MaxPool2d VEJP with 4 manual 2x2 inputs. - - Uses kernel_size=2, stride=2 on 1x1x2x2 inputs, producing a scalar output. - The max is at the top-left position in 3 of 4 stats inputs and at the - top-right in 1 of 4, so the expected smooth gradient is: - [[0.75, 0.25], - [0.00, 0.00]] +def test_maxpool2d_k2s2_vejp(): + """Test the MaxPool2d VEJP with kernel_size=2, stride=2 on 2x2 inputs. + + 10 inputs total (n=10): 9 stats-only inputs + 1 backward-pass input. + The test input has a clear max at (0,0) to avoid tie-breaking differences. + Expected counts: [[4, 3], [2, 1]], n=10 → VEJP = [[0.4, 0.3], [0.2, 0.1]] """ stats_inputs = [ - # max at (0,0) - torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]), - # max at (0,0) + torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]), # 3x max at (0,0) torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]), - # max at (0,0) torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]), - # max at (0,1) + torch.tensor([[[[0.0, 1.0], [0.0, 0.0]]]]), # 3x max at (0,1) torch.tensor([[[[0.0, 1.0], [0.0, 0.0]]]]), + torch.tensor([[[[0.0, 1.0], [0.0, 0.0]]]]), + torch.tensor([[[[0.0, 0.0], [1.0, 0.0]]]]), # 2x max at (1,0) + torch.tensor([[[[0.0, 0.0], [1.0, 0.0]]]]), + torch.tensor([[[[0.0, 0.0], [0.0, 1.0]]]]), # 1x max at (1,1) ] - test_input = torch.tensor([[[[1.0, 1.0], [1.0, 1.0]]]], requires_grad=True) + test_input = torch.tensor([[[[4.0, 1.0], [1.0, 1.0]]]], requires_grad=True) layer = _SmoothMaxPool2d(kernel_size=2, stride=2, padding=0) layer.reset_stats() - # Phase 1: Collect stats layer.collect_stats = True layer.smooth_backward = False for x in stats_inputs: layer(x) - # Phase 2: Smooth backward pass - layer.collect_stats = False + layer.collect_stats = True + layer.smooth_backward = True + output = layer(test_input) + output.sum().backward() + + expected = torch.tensor([[[[0.4, 0.3], [0.2, 0.1]]]]) + assert torch.allclose(test_input.grad, expected), ( # type: ignore[arg-type] + f"Expected {expected}, got {test_input.grad}" + ) + + +def test_maxpool2d_k3s3_vejp(): + """Test the MaxPool2d VEJP with kernel_size=3, stride=3 on 3x3 inputs. + + 10 inputs total (n=10): 9 stats-only inputs + 1 backward-pass input. + The test input has a clear max at (0,0) to avoid tie-breaking differences. + Expected counts: [[4, 0, 2], [0, 3, 0], [1, 0, 0]], n=10 + → VEJP = [[0.4, 0.0, 0.2], [0.0, 0.3, 0.0], [0.1, 0.0, 0.0]] + """ + stats_inputs = [ + torch.tensor( + [[[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]] + ), # 3x (0,0) + torch.tensor([[[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]]), + torch.tensor([[[[1.0, 0.0, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]]), + torch.tensor( + [[[[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]] + ), # 2x (0,2) + torch.tensor([[[[0.0, 0.0, 1.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]]]), + torch.tensor( + [[[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]]] + ), # 3x (1,1) + torch.tensor([[[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]]]), + torch.tensor([[[[0.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 0.0]]]]), + torch.tensor( + [[[[0.0, 0.0, 0.0], [0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]]] + ), # 1x (2,0) + ] + test_input = torch.tensor( + [[[[9.0, 1.0, 1.0], [1.0, 1.0, 1.0], [1.0, 1.0, 1.0]]]], requires_grad=True + ) + + layer = _SmoothMaxPool2d(kernel_size=3, stride=3, padding=0) + layer.reset_stats() + + layer.collect_stats = True + layer.smooth_backward = False + for x in stats_inputs: + layer(x) + + layer.collect_stats = True layer.smooth_backward = True output = layer(test_input) output.sum().backward() - expected = torch.tensor([[[[0.75, 0.25], [0.0, 0.0]]]]) + expected = torch.tensor([[[[0.4, 0.0, 0.2], [0.0, 0.3, 0.0], [0.1, 0.0, 0.0]]]]) assert torch.allclose(test_input.grad, expected), ( # type: ignore[arg-type] f"Expected {expected}, got {test_input.grad}" ) diff --git a/python/tests/test_relu_vejp.py b/python/tests/test_relu_vejp.py index 138b942..bf65e2e 100644 --- a/python/tests/test_relu_vejp.py +++ b/python/tests/test_relu_vejp.py @@ -8,9 +8,9 @@ def test_relu_vejp(): """Test the ReLU VEJP with 5 manual inputs. - The first 4 inputs are used for stats collection, and the 5th input - is used for the smooth backward pass. The expected smooth gradient - is [1, 0.75, 0.5, 0.25, 0]. + All 5 inputs count toward the statistics (n=5): 4 stats-only inputs + plus the backward-pass input. The expected smooth gradient is + count / n = [5, 4, 3, 2, 1] / 5 = [1.0, 0.8, 0.6, 0.4, 0.2]. """ stats_inputs = [ torch.tensor([1.0, 0.0, 0.0, 0.0, 0.0]), @@ -29,13 +29,13 @@ def test_relu_vejp(): for x in stats_inputs: layer(x) - # Phase 2: Smooth backward pass - layer.collect_stats = False + # Phase 2: Smooth backward pass (still collecting stats so test input counts) + layer.collect_stats = True layer.smooth_backward = True output = layer(test_input) output.sum().backward() - expected = torch.tensor([1.0, 0.75, 0.5, 0.25, 0.0]) + expected = torch.tensor([1.0, 0.8, 0.6, 0.4, 0.2]) assert torch.allclose(test_input.grad, expected), ( # type: ignore[arg-type] f"Expected {expected}, got {test_input.grad}" )