From 3986e964131bd6b5315e83c8b2bd1692761b35f2 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Tue, 21 Apr 2026 00:03:12 +0200 Subject: [PATCH 1/2] [AutoDiff] Adstack regression tests --- tests/python/test_adstack.py | 159 +++++++++++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index e50f897663..6969c121fb 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -319,3 +319,162 @@ def compute(): compute.grad() assert math.isnan(x.grad[None]) + + +def _run_basic_gradient(qd_dtype, n_iter, rel_tol, approx=test_utils.approx): + # Builds the kernel, runs forward + backward, and asserts a correct gradient. `approx` defaults to + # `test_utils.approx` which is correct for f32 (its backend-specific floor kicks in); f64 callers must pass + # `pytest.approx` to honor a tight `rel_tol=1e-14` that `test_utils.approx` would otherwise floor to 1e-6. + # The adstack is structurally required here so the backward compiler can reverse the dynamic `range(n_iter)` + # at all; the companion `test_adstack_basic_gradient_negative` pins that disabling the adstack raises + # `QuadrantsCompilationError("non static range")`. Value-correctness of the per-iteration `v` spilled on the + # adstack is NOT exercised by the linear body `v = v * 0.95 + 0.01` - see `test_adstack_basic_gradient`'s + # docstring for the details. + n = 4 + x = qd.field(qd_dtype, shape=n, needs_grad=True) + y = qd.field(qd_dtype, shape=(), needs_grad=True) + + @qd.kernel + def compute(): + for i in x: + v = x[i] + for _ in range(n_iter): + v = v * 0.95 + 0.01 + y[None] += v + + x_vals = [0.1, 0.3, 0.5, 0.8] + for i, v in enumerate(x_vals): + x[i] = v + y[None] = 0.0 + compute() + y.grad[None] = 1.0 + for i in range(n): + x.grad[i] = 0.0 + + compute.grad() + + # `v = v * 0.95 + 0.01` iterated n_iter times gives v_final = 0.95**n_iter * x[i] + const, so + # dv_final/dx[i] == 0.95**n_iter independent of x[i], and dy/dx[i] equals the same quantity. + expected = 0.95**n_iter + for i in range(n): + assert x.grad[i] == approx(expected, rel=rel_tol) + + +@pytest.mark.parametrize("n_iter", [1, 3, 10]) +@test_utils.test(require=qd.extension.adstack) +def test_adstack_basic_gradient(n_iter): + # Smallest possible "does reverse-mode AD through a for-loop work at all" check. The kernel runs `n_iter` + # iterations of `v = v * 0.95 + 0.01` per element and asserts that `dy/dx[i]` matches the analytical gradient + # `0.95 ** n_iter` for every element. + # + # Internal details: the adstack is structurally required here so the backward compiler can reverse the + # dynamic `range(n_iter)` at all - the companion `test_adstack_basic_gradient_negative` pins that disabling + # the adstack raises `QuadrantsCompilationError` in exactly this kernel shape. Value-correctness of the + # stored v, on the other hand, is NOT exercised: the loop body `v = v * 0.95 + 0.01` is linear, so the + # backward chain `adj(v_prev) = 0.95 * adj(v_next)` only uses the compile-time constant 0.95 and never reads + # v from the adstack. A broken push/load/pop that returned garbage for v would still produce the same exact + # gradient. For push/load/pop value-correctness coverage, see `test_adstack_unary_loop_carried` (non-linear + # unary ops in the loop body). `n_iter = 1` exercises the single-push adstack code path; `n_iter = 10` + # exercises repeated push/pop under one forward invocation; multi-element coverage (n = 4) guards against + # per-element accumulation bugs that a single-element variant would miss. + _run_basic_gradient(qd.f32, n_iter=n_iter, rel_tol=1e-6) + + +@pytest.mark.parametrize("n_iter", [1, 3, 10]) +@test_utils.test(require=[qd.extension.adstack, qd.extension.data64], default_fp=qd.f64) +def test_adstack_basic_gradient_f64(n_iter): + # f64 counterpart of `test_adstack_basic_gradient`. Uses `pytest.approx` so the tight `rel_tol=1e-14` is + # honored; `test_utils.approx` would floor it to `get_rel_eps()` (typically 1e-6) and silently pass an + # f32-precision regression. + _run_basic_gradient(qd.f64, n_iter=n_iter, rel_tol=1e-14, approx=pytest.approx) + + +@pytest.mark.parametrize("n_iter", [1, 3, 10]) +@test_utils.test(ad_stack_experimental_enabled=False) +def test_adstack_basic_gradient_negative(n_iter): + # Negative counterpart of `test_adstack_basic_gradient`: with the adstack disabled the backward compiler + # cannot reverse a dynamic `range(n_iter)`, so `compute.grad()` raises `QuadrantsCompilationError("Cannot use + # non static range in Backwards mode")` deterministically for every `n_iter`. Inlined rather than reusing + # `_run_basic_gradient` because the shall-not-pass path never reaches the gradient assertion, so a shared + # helper would carry a dead `rel_tol` argument down this branch. + n = 4 + x = qd.field(qd.f32, shape=n, needs_grad=True) + y = qd.field(qd.f32, shape=(), needs_grad=True) + + @qd.kernel + def compute(): + for i in x: + v = x[i] + for _ in range(n_iter): + v = v * 0.95 + 0.01 + y[None] += v + + x_vals = [0.1, 0.3, 0.5, 0.8] + for i, v in enumerate(x_vals): + x[i] = v + y[None] = 0.0 + compute() + + with pytest.raises(qd.QuadrantsCompilationError, match=r"non static range"): + compute.grad() + + +def _run_sum_linear(qd_dtype, use_static_loop, use_varying_coeff, n_iter, rel_tol, approx=test_utils.approx): + n = 4 + x = qd.field(qd_dtype, shape=n, needs_grad=True) + y = qd.field(qd_dtype, shape=(), needs_grad=True) + + @qd.kernel + def compute(): + for i in x: + v = x[i] + for a in qd.static(range(n_iter)) if qd.static(use_static_loop) else range(n_iter): + if qd.static(use_varying_coeff): + y[None] += v * qd.cast(a + 1, qd_dtype) + else: + y[None] += v + + x_vals = [0.1, 0.3, 0.5, 0.8] + for i, v in enumerate(x_vals): + x[i] = v + y[None] = 0.0 + compute() + y.grad[None] = 1.0 + for i in range(n): + x.grad[i] = 0.0 + compute.grad() + + expected = sum(a + 1 for a in range(n_iter)) if use_varying_coeff else float(n_iter) + for i in range(n): + assert x.grad[i] == approx(expected, rel=rel_tol) + + +@pytest.mark.parametrize("n_iter", [1, 3, 10]) +@pytest.mark.parametrize("use_static_loop", [True, False]) +@pytest.mark.parametrize("use_varying_coeff", [True, False]) +@test_utils.test(require=qd.extension.adstack) +def test_adstack_sum_linear(use_static_loop, use_varying_coeff, n_iter): + # Linear accumulation `y = sum_j v * coeff_j` across all four combinations of (static-unrolled vs dynamic loop) + # x (constant coefficient vs loop-index-varying coefficient), at three loop lengths. Replaces the earlier three + # separate tests (`test_adstack_sum_fixed_coeff`, `test_adstack_sum_constant_coeffs`, + # `test_adstack_sum_static_loop_correct`) with a single parametrized version so every branch of that truth + # table is covered at each trip count. + # + # Internal details: this test deliberately does not mutate `v` inside the loop, so the reverse pass does not + # require adstack replay of `v` to compute the right gradient - `v`'s per-iteration value is the same `x[i]`. + # The point of this test is therefore not to stress the adstack (that is `test_adstack_basic_gradient`'s job) + # but to prove that enabling the adstack extension does not silently regress linear reverse-mode AD for either + # unrolled or dynamic loop shapes. No negative counterpart is included: for `use_static_loop=True` the inner + # loop is unrolled and the backward kernel contains no dynamic range, so the adstack option does not change + # the gradient; for `use_static_loop=False` disabling the adstack would raise `QuadrantsCompilationError` + # (same compile-time rejection covered by `test_adstack_basic_gradient_negative`), which is out of scope here. + _run_sum_linear(qd.f32, use_static_loop, use_varying_coeff, n_iter, rel_tol=1e-6) + + +@pytest.mark.parametrize("n_iter", [1, 3, 10]) +@pytest.mark.parametrize("use_static_loop", [True, False]) +@pytest.mark.parametrize("use_varying_coeff", [True, False]) +@test_utils.test(require=[qd.extension.adstack, qd.extension.data64], default_fp=qd.f64) +def test_adstack_sum_linear_f64(use_static_loop, use_varying_coeff, n_iter): + # f64 counterpart uses `pytest.approx` so the tight rel_tol is not floored by `test_utils.approx`. + _run_sum_linear(qd.f64, use_static_loop, use_varying_coeff, n_iter, rel_tol=1e-14, approx=pytest.approx) From f656de7c6bf2e439b173e88ca0c6221b71265f15 Mon Sep 17 00:00:00 2001 From: Alexis Duburcq Date: Tue, 21 Apr 2026 16:33:09 +0200 Subject: [PATCH 2/2] [Tests] Pass abs=0 to pytest.approx in f64 adstack gradient asserts so the tight rel_tol is honored instead of dominated by the default 1e-12 abs floor --- tests/python/test_adstack.py | 40 +++++++++++++++++++++++++++--------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/tests/python/test_adstack.py b/tests/python/test_adstack.py index 6969c121fb..13c56db3df 100644 --- a/tests/python/test_adstack.py +++ b/tests/python/test_adstack.py @@ -321,10 +321,13 @@ def compute(): assert math.isnan(x.grad[None]) -def _run_basic_gradient(qd_dtype, n_iter, rel_tol, approx=test_utils.approx): +def _run_basic_gradient(qd_dtype, n_iter, rel_tol, approx=test_utils.approx, abs_tol=None): # Builds the kernel, runs forward + backward, and asserts a correct gradient. `approx` defaults to # `test_utils.approx` which is correct for f32 (its backend-specific floor kicks in); f64 callers must pass # `pytest.approx` to honor a tight `rel_tol=1e-14` that `test_utils.approx` would otherwise floor to 1e-6. + # `abs_tol` is forwarded as `abs=` to the approx; f64 callers must pass `abs_tol=0` because pytest.approx's + # default `abs=1e-12` would otherwise dominate for the expected magnitudes here (~0.6-0.95), making the + # effective tolerance ~1e-12 absolute rather than 1e-14 relative and missing f32-narrowing regressions. # The adstack is structurally required here so the backward compiler can reverse the dynamic `range(n_iter)` # at all; the companion `test_adstack_basic_gradient_negative` pins that disabling the adstack raises # `QuadrantsCompilationError("non static range")`. Value-correctness of the per-iteration `v` spilled on the @@ -356,8 +359,11 @@ def compute(): # `v = v * 0.95 + 0.01` iterated n_iter times gives v_final = 0.95**n_iter * x[i] + const, so # dv_final/dx[i] == 0.95**n_iter independent of x[i], and dy/dx[i] equals the same quantity. expected = 0.95**n_iter + approx_kwargs = {"rel": rel_tol} + if abs_tol is not None: + approx_kwargs["abs"] = abs_tol for i in range(n): - assert x.grad[i] == approx(expected, rel=rel_tol) + assert x.grad[i] == approx(expected, **approx_kwargs) @pytest.mark.parametrize("n_iter", [1, 3, 10]) @@ -383,10 +389,12 @@ def test_adstack_basic_gradient(n_iter): @pytest.mark.parametrize("n_iter", [1, 3, 10]) @test_utils.test(require=[qd.extension.adstack, qd.extension.data64], default_fp=qd.f64) def test_adstack_basic_gradient_f64(n_iter): - # f64 counterpart of `test_adstack_basic_gradient`. Uses `pytest.approx` so the tight `rel_tol=1e-14` is - # honored; `test_utils.approx` would floor it to `get_rel_eps()` (typically 1e-6) and silently pass an - # f32-precision regression. - _run_basic_gradient(qd.f64, n_iter=n_iter, rel_tol=1e-14, approx=pytest.approx) + # f64 counterpart of `test_adstack_basic_gradient`. Uses `pytest.approx` so the tight `rel_tol=1e-14` is honored; + # `test_utils.approx` would floor it to `get_rel_eps()` (typically 1e-6) and silently pass an f32-precision + # regression. `abs_tol=0` disables pytest.approx's default `abs=1e-12` floor, which would otherwise dominate for + # the expected magnitudes `0.95**n_iter in [~0.6, 0.95]` and make the effective tolerance ~1e-12 absolute rather + # than 1e-14 relative; that is still ~100x looser than f64 roundoff and would miss an f32-narrowing regression. + _run_basic_gradient(qd.f64, n_iter=n_iter, rel_tol=1e-14, approx=pytest.approx, abs_tol=0) @pytest.mark.parametrize("n_iter", [1, 3, 10]) @@ -419,7 +427,9 @@ def compute(): compute.grad() -def _run_sum_linear(qd_dtype, use_static_loop, use_varying_coeff, n_iter, rel_tol, approx=test_utils.approx): +def _run_sum_linear( + qd_dtype, use_static_loop, use_varying_coeff, n_iter, rel_tol, approx=test_utils.approx, abs_tol=None +): n = 4 x = qd.field(qd_dtype, shape=n, needs_grad=True) y = qd.field(qd_dtype, shape=(), needs_grad=True) @@ -445,8 +455,11 @@ def compute(): compute.grad() expected = sum(a + 1 for a in range(n_iter)) if use_varying_coeff else float(n_iter) + approx_kwargs = {"rel": rel_tol} + if abs_tol is not None: + approx_kwargs["abs"] = abs_tol for i in range(n): - assert x.grad[i] == approx(expected, rel=rel_tol) + assert x.grad[i] == approx(expected, **approx_kwargs) @pytest.mark.parametrize("n_iter", [1, 3, 10]) @@ -476,5 +489,12 @@ def test_adstack_sum_linear(use_static_loop, use_varying_coeff, n_iter): @pytest.mark.parametrize("use_varying_coeff", [True, False]) @test_utils.test(require=[qd.extension.adstack, qd.extension.data64], default_fp=qd.f64) def test_adstack_sum_linear_f64(use_static_loop, use_varying_coeff, n_iter): - # f64 counterpart uses `pytest.approx` so the tight rel_tol is not floored by `test_utils.approx`. - _run_sum_linear(qd.f64, use_static_loop, use_varying_coeff, n_iter, rel_tol=1e-14, approx=pytest.approx) + # f64 counterpart uses `pytest.approx` so the tight rel_tol is not floored by `test_utils.approx`, and + # `abs_tol=0` disables pytest.approx's default `abs=1e-12` floor so the 1e-14 relative tolerance is actually + # honored. Note: the expected gradients here are integers in `{1, 3, 6, 10, 55}` which are exactly representable + # in both f32 and f64, so this test cannot catch an f32-narrowing regression of the backward pass regardless of + # tolerance - the value coverage comes from `test_adstack_basic_gradient_f64` whose non-integer expected values + # genuinely exercise the f64 precision floor. Kept here to preserve the shape coverage (static/dynamic inner + # loop x varying/constant coefficient) at f64 since a future type-narrowing bug that depends on shape might + # still surface through compile-time / structural differences between f32 and f64 codegen paths. + _run_sum_linear(qd.f64, use_static_loop, use_varying_coeff, n_iter, rel_tol=1e-14, approx=pytest.approx, abs_tol=0)