Unify numerical VEJP tests

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"

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
 

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

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"

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")

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

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}"
     )

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}"
     )