pushpull: backpropagate through the sampling grid (closes #16)

[balbasty]

Closes#16.

Top of a seven-repo stack — this is pure Python and merges cleanly, but it only works once the chain below it lands (see the bottom of this description).

The gap

pull and push treated grid as a constant: _no_grid_grad raised NotImplementedError if it required grad, so differentiating through sample positions — the thing you need for a learned deformation field or a registration model — was simply unavailable.

The C++ adjoints existed in the impl layer the whole time but were never exported. The stack below now exposes them, so _Pull/_Push's backward() can return a real grid gradient and the guard is gone.

Implementation

Both branches are kept, because they cost different amounts:

• grid needs grad → one fused pull_backward / push_backward call returns both adjoints.
• only the field needs grad → the old, cheaper push / pull path (the plain adjoint never touches the field to compute a spatial gradient).

inp is likewise only saved for backward when grid.requires_grad, so the field-only path keeps its previous memory profile.

backward is now marked @once_differentiable. It calls straight into C++, which is opaque to autograd, so double-backward was never actually supported — it would silently have produced a wrong second derivative under create_graph. Better to say so than to keep returning quiet garbage.

count and grad stay non-differentiable here. Their adjoints are now exported by fastfields.dlpack (count_backward / grad_backward, both covered by the C++ oracle), they are just not wired into an autograd.Function — that is beyond what #16 asked for.

Tests

gradcheck is the actual proof, so the removed "not supported" test is replaced by 26 of them: pull and push, wrt grid alone and wrt (inp, grid) jointly, in 1D/2D/3D, orders 1/2/3, five boundary conditions, and with a leading batch dimension.

Plus two behavioural tests that finite differences cannot cover:

• the field gradient must be identical whichever of the two branches computes it;
• out-of-FOV samples (extrapolate=0) must get an exactly zero grid gradient — the FOV test is a step, not a smooth factor.

Grid coordinates in these tests are kept off the spline knots. A spline of order k is only C^(k−1), so at order 1 the derivative wrt position genuinely jumps at a knot and a central difference straddling one is meaningless. That is a property of the interpolant, not a limitation of the implementation.

Full suite: 104 passed.

What this turned up underneath

Exposing the adjoints put them under a finite-difference oracle for the first time, and the pre-existing math was not correct — four bugs, all inherited from jitfields:

1. push_backward indexed ginp with the wrong stride array — silent wrong gradients whenever the pushed volume and the grid have different spatial strides (801/6381 oracle checks)
2. 1D linear grad_backward returned a zero gradient wrt inp
3. 2D linear grad_backward returned a zero gradient wrt grid (dropped mixed second derivative)
4. 3D linear grad_backward — same, all three mixed terms

The stack

Merge bottom-up; each pin needs re-bumping to its parent's merged SHA as you go.

#	repo	PR
1	fastfields-kernels	#45
2	fastfields-cpu-impl	#39
3	fastfields-cuda-impl	#22
4	fastfields-cpu-lib	#49
5	fastfields-cuda-lib	#32
6	fastfields-lib	#46
7	fastfields-dlpack	#19
8	fastfields-torch	this PR

PRs 1–6 are based on claude/cuda-pushpull-dynamic-spline (kernels#44 / cpu-lib#48), not main — that work is unmerged and this builds on its FF_SPLINE_*/SplineVec machinery. They need retargeting to main once it lands.

Workstream: claude-jitfields-to-fastfields

🤖 Generated with Claude Code

---

Generated by Claude Code

[balbasty]

We want to support as many pytorch versions as possible: check that we are not using torch functions that are too narrow (@oncedifferentiable?) and if yes, provide compatibility fallbacks.

[balbasty]

check that we are not using torch functions that are too narrow (@once_differentiable?) and if yes, provide compatibility fallbacks

Checked properly rather than guessing — no shim needed, nothing changed. Details, so you can disagree with the evidence rather than with me:

What the floor actually is

pyproject.toml declares dependencies = ["fastfields-dlpack", "torch", "numpy"] — torch is unpinned — and requires-python = ">=3.9". So the effective floor is "the oldest torch that has a py3.9 wheel", i.e. torch 1.7.1. CI (fastfields/.github/.github/workflows/python-test.yml) is py 3.9/3.11/3.12 × ubuntu/macos with pip install ".[test]", so what actually gets exercised is the latest 2.x. Range to defend: 1.7.1 → current (I ran against 2.13).

once_differentiable

Present in torch/autograd/function.py at every tag I fetched: v0.4.1, v1.0.0, v1.5.0, v1.7.1, v1.13.1, v2.0.0, v2.4.0, v2.8.0, main. The body is the same logic character-for-character from 1.7.1 through main — the only diffs are black reformatting and a typed signature (Concatenate/ParamSpec) on main. It is in __all__ on main, and torch's own Function docstring uses it in five separate examples; the only @deprecated markers in that module are mark_shared_storage and legacy non-static forward.

Import path has never moved. torch/autograd/__init__.py does from .function import Function, NestedIOFunction as far back as 1.7.1, so torch.autograd.function.once_differentiable is reachable from a bare import torch — no import torch.autograd.function needed, then or now. A try/except ImportError fallback would be unreachable code.

Semantics are also stable: no_grad() around the call, then a DelayedError alias so a second backward raises. Confirmed on 2.13 that create_graph=True → RuntimeError rather than silent garbage, and that nothing emits a warning under -W error::DeprecationWarning -W error::FutureWarning.

The other new APIs in this diff

• ctx.save_for_backward(inp if grid.requires_grad else None, grid) — saving None is the one genuinely non-obvious thing here, so I checked the C++: _save_variables in torch/csrc/autograd/python_function.cpp has the explicit if (obj == Py_None) { saved_variables.emplace_back(); continue; } branch at v1.7.1, v1.13.1, v2.4.0 and main, and the unpack side maps the undefined slot back to None. Verified at runtime on 2.13 that saved_tensors[0] is None.
• reading grid.requires_grad inside forward — unpack_input forwards the original Python tensor objects (Py_INCREF(arg), no detach) and derives needs_input_grad[i] from that very same requires_grad() flag; identical at v1.7.1 and today. So the forward-time branch and ctx.needs_input_grad[1] in backward cannot disagree, on any version.
• classic forward(ctx, …) signature — still supported on main; setup_context is opt-in, not a replacement.
• Everything else is pre-1.0: .detach(), .contiguous(), new_zeros, ctx.needs_input_grad, ctx.saved_tensors, and torch.autograd.gradcheck(eps=, atol=) in the tests.

Caveat + the one thing that would be actionable

The old-version evidence is source-level at the tags (1.7.1 has no wheel for the py3.11 in this environment, so I couldn't execute it); the runtime checks were on 2.13.

The real gap isn't @once_differentiable, it's that nothing pins or tests a floor: torch is unpinned and CI only ever resolves the newest 2.x, so "we support old torch" is currently unverified for the whole package, not just this PR. If you want that nailed down, the orthogonal fix is a torch>=X pin in pyproject.toml plus an oldest-supported leg in the matrix. Say the word and I'll do it in a separate PR — it shouldn't ride along with the grid-backward change.

---

Generated by Claude Code