From 10cbfec1993970257d173ffd308c10e87a69bc7e Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 17:43:52 +0100 Subject: [PATCH 1/9] Add numerical VEJP tests --- julia/test/runtests.jl | 4 +++ julia/test/test_vejp.jl | 66 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 70 insertions(+) create mode 100644 julia/test/test_vejp.jl 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..6c04d90 --- /dev/null +++ b/julia/test/test_vejp.jl @@ -0,0 +1,66 @@ +using SmoothedDifferentiation +using SmoothedDifferentiation: ReluAccumulator, MaxPoolAccumulator, reset_counts! +using Test +using Flux: Flux, MaxPool +using NNlib: relu +using Zygote: pullback + +@testset "ReLU VEJP" begin + # 4 stats inputs, then 1 smooth backward pass. + # 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] + # Expected VEJP: count / n = [4/4, 3/4, 2/4, 1/4, 0/4] + # = [1.0, 0.75, 0.5, 0.25, 0.0] + + 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 + + # Compute VeJP with grad_output = ones + 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.75, 0.5, 0.25, 0.0] +end + +@testset "MaxPool VEJP" begin + # 4 stats inputs (W=2, H=2, C=1, N=1), kernel_size=2, stride=2. + # The max is at the top-left position in 3 of 4 inputs + # and at the top-right in 1 of 4. + # Expected VEJP: [0.75 0.0; 0.25 0.0] (in Julia's column-major WHCN layout) + + pool = MaxPool((2, 2); stride = (2, 2)) + count = zeros(Int, 2, 2, 1, 1) + layer = MaxPoolAccumulator(; layer = pool, count = count) + + # In Julia WHCN layout: array[w, h, c, n] + # "max at (w=1, h=1)" = top-left + stats_inputs = [ + Float32[1 0; 0 0;;;], # max at (1,1) + Float32[1 0; 0 0;;;], # max at (1,1) + Float32[1 0; 0 0;;;], # max at (1,1) + Float32[0 0; 1 0;;;], # max at (2,1) + ] + for x in stats_inputs + layer(x) + end + + # Compute VeJP with grad_output = ones + test_input = Float32[1 1; 1 1;;;] + _, vejp_fn = pullback(layer, test_input) + grad = only(vejp_fn(ones(Float32, 1, 1, 1, 1))) + + @test grad ≈ Float32[0.75 0.0; 0.25 0.0;;;] +end From 2f041e65090245627c5878c2842970e496b3c37e Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 17:51:03 +0100 Subject: [PATCH 2/9] Add Zygote to test env --- julia/test/Project.toml | 1 + 1 file changed, 1 insertion(+) 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" From 285af2b86b35d7ed3e6c3beb11795c38cee324c1 Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 19:49:19 +0100 Subject: [PATCH 3/9] Update PyTorch test --- python/tests/test_maxpool_vejp.py | 31 ++++++++++++++++++------------- python/tests/test_relu_vejp.py | 12 ++++++------ 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/python/tests/test_maxpool_vejp.py b/python/tests/test_maxpool_vejp.py index 6e220ac..c6c42e7 100644 --- a/python/tests/test_maxpool_vejp.py +++ b/python/tests/test_maxpool_vejp.py @@ -6,25 +6,30 @@ def test_maxpool2d_vejp(): - """Test the MaxPool2d VEJP with 4 manual 2x2 inputs. + """Test the MaxPool2d VEJP with 10 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]] + All 10 inputs count toward the statistics (n=10): 9 stats-only inputs + plus the 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) + # 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,0) torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]), - # max at (0,1) + # 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, 1.0], [0.0, 0.0]]]]), + # 2x max at (1,0) + torch.tensor([[[[0.0, 0.0], [1.0, 0.0]]]]), + torch.tensor([[[[0.0, 0.0], [1.0, 0.0]]]]), + # 1x max at (1,1) + torch.tensor([[[[0.0, 0.0], [0.0, 1.0]]]]), ] - 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() @@ -35,13 +40,13 @@ def test_maxpool2d_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([[[[0.75, 0.25], [0.0, 0.0]]]]) + 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}" ) 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}" ) From 3b454c28e4cf1d4c1e067305a45a347de14647d3 Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 19:49:41 +0100 Subject: [PATCH 4/9] Update Julia test --- julia/test/test_vejp.jl | 40 +++++++++++++++++++++++----------------- 1 file changed, 23 insertions(+), 17 deletions(-) diff --git a/julia/test/test_vejp.jl b/julia/test/test_vejp.jl index 6c04d90..d9f51be 100644 --- a/julia/test/test_vejp.jl +++ b/julia/test/test_vejp.jl @@ -6,14 +6,15 @@ using NNlib: relu using Zygote: pullback @testset "ReLU VEJP" begin - # 4 stats inputs, then 1 smooth backward pass. + # 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] - # Expected VEJP: count / n = [4/4, 3/4, 2/4, 1/4, 0/4] - # = [1.0, 0.75, 0.5, 0.25, 0.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)) @@ -27,40 +28,45 @@ using Zygote: pullback layer(x) end - # Compute VeJP with grad_output = ones + # 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.75, 0.5, 0.25, 0.0] + @test grad ≈ Float32[1.0, 0.8, 0.6, 0.4, 0.2] end @testset "MaxPool VEJP" begin - # 4 stats inputs (W=2, H=2, C=1, N=1), kernel_size=2, stride=2. - # The max is at the top-left position in 3 of 4 inputs - # and at the top-right in 1 of 4. - # Expected VEJP: [0.75 0.0; 0.25 0.0] (in Julia's column-major WHCN layout) + # 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) # In Julia WHCN layout: array[w, h, c, n] - # "max at (w=1, h=1)" = top-left stats_inputs = [ - Float32[1 0; 0 0;;;], # max at (1,1) - Float32[1 0; 0 0;;;], # max at (1,1) - Float32[1 0; 0 0;;;], # max at (1,1) - Float32[0 0; 1 0;;;], # max at (2,1) + 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 - # Compute VeJP with grad_output = ones - test_input = Float32[1 1; 1 1;;;] + # 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.75 0.0; 0.25 0.0;;;] + @test grad ≈ Float32[0.4 0.2; 0.3 0.1;;;;] end From df8b31ca76f97c3ee68340bcdb00b530be0f2fe0 Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 19:49:48 +0100 Subject: [PATCH 5/9] Bump Zygote compat --- julia/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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" From 71a5e6f9c9dfaefe1cf444031bc57659ef31750e Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 20:34:01 +0100 Subject: [PATCH 6/9] Update Zygote --- julia/src/SmoothedDifferentiation.jl | 2 +- julia/src/vejp/maxpool.jl | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) 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..7e6bb14 100644 --- a/julia/src/vejp/maxpool.jl +++ b/julia/src/vejp/maxpool.jl @@ -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 From 5b0b62c2af72b7dd51969adf5dc604eb4f51aab8 Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 20:34:08 +0100 Subject: [PATCH 7/9] Appease JET --- julia/src/vejp/maxpool.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/julia/src/vejp/maxpool.jl b/julia/src/vejp/maxpool.jl index 7e6bb14..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) From aa033ff374004bcf1cb64d3eeacd3999e881140e Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 20:54:53 +0100 Subject: [PATCH 8/9] Add more tests --- julia/test/test_vejp.jl | 38 ++++++++++++++++- python/tests/test_maxpool_vejp.py | 70 +++++++++++++++++++++++-------- 2 files changed, 88 insertions(+), 20 deletions(-) diff --git a/julia/test/test_vejp.jl b/julia/test/test_vejp.jl index d9f51be..6bfb4bf 100644 --- a/julia/test/test_vejp.jl +++ b/julia/test/test_vejp.jl @@ -36,7 +36,7 @@ using Zygote: pullback @test grad ≈ Float32[1.0, 0.8, 0.6, 0.4, 0.2] end -@testset "MaxPool VEJP" begin +@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. @@ -47,7 +47,6 @@ end count = zeros(Int, 2, 2, 1, 1) layer = MaxPoolAccumulator(; layer = pool, count = count) - # In Julia WHCN layout: array[w, h, c, n] stats_inputs = [ Float32[1 0; 0 0;;;;], # 3x max at (1,1) Float32[1 0; 0 0;;;;], @@ -70,3 +69,38 @@ end @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 c6c42e7..393835e 100644 --- a/python/tests/test_maxpool_vejp.py +++ b/python/tests/test_maxpool_vejp.py @@ -5,42 +5,34 @@ from smoothdiff_torch.smoothdiff import _SmoothMaxPool2d -def test_maxpool2d_vejp(): - """Test the MaxPool2d VEJP with 10 manual 2x2 inputs. - - Uses kernel_size=2, stride=2 on 1x1x2x2 inputs, producing a scalar output. - All 10 inputs count toward the statistics (n=10): 9 stats-only inputs - plus the 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]] +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 = [ - # 3x max at (0,0) - torch.tensor([[[[1.0, 0.0], [0.0, 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]]]]), torch.tensor([[[[1.0, 0.0], [0.0, 0.0]]]]), - # 3x 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, 1.0], [0.0, 0.0]]]]), - # 2x max at (1,0) - torch.tensor([[[[0.0, 0.0], [1.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]]]]), - # 1x max at (1,1) - torch.tensor([[[[0.0, 0.0], [0.0, 1.0]]]]), + torch.tensor([[[[0.0, 0.0], [0.0, 1.0]]]]), # 1x max at (1,1) ] 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 (still collecting stats so test input counts) layer.collect_stats = True layer.smooth_backward = True output = layer(test_input) @@ -50,3 +42,45 @@ def test_maxpool2d_vejp(): 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.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}" + ) From 4c40d814545ddc47a8c7874989b3ad9b3d6665db Mon Sep 17 00:00:00 2001 From: adrhill Date: Mon, 16 Feb 2026 20:58:28 +0100 Subject: [PATCH 9/9] Run ruff --- python/tests/test_maxpool_vejp.py | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/python/tests/test_maxpool_vejp.py b/python/tests/test_maxpool_vejp.py index 393835e..f4b04de 100644 --- a/python/tests/test_maxpool_vejp.py +++ b/python/tests/test_maxpool_vejp.py @@ -53,15 +53,23 @@ def test_maxpool2d_k3s3_vejp(): → 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]]]] + ), # 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]]]] + ), # 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]]]] + ), # 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) + 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