FIX(pushpull): implement the N-dimensional backward passes

[balbasty]

Closes#4

What was wrong

jitfields/csrc/lib/pushpull/nd.h was a stub. As reported in #4, the backward passes were copy-pasted verbatim from the 3D kernels: they still took the x, nx, sx, y, ny, sy, z, nz, sz argument convention, called utils_x/y/z::gindex, and wrote exactly three coordinate gradients (gout[0], gout[osg], gout[osg*2]) regardless of D.

The forward passes were in no better shape — they did not compile at all:

• sd = s + 8*d where s is never declared (the array is called f),
• db[d] used as the loop bound where the array is called l,
• coord referenced in push/count/grad, which don't have such a parameter,
• *out = static_cast<scalar_t>(acc) in grad, where acc is an array,
• an outright syntax error in grad: acc[dd] += val * (g[8*d + ] * weights[D-1];,
• #endif JF_PUSHPULL_ND (trailing tokens after #endif).

None of this was ever caught because nothing included the file: lib/pushpull.h only pulled in 1d.h, 2d.h and 3d.h.

What this PR does

nd.h is rewritten as a genuine N-dimensional implementation.

PushPullND<D, ABS> is now a class template of its own rather than a partial specialization of PushPull. Two reasons:

• PushPull only carries three (spline order, boundary condition) pairs, so for D > 3 they have to be runtime parameters anyway — they are read from inter[] / bnd[] arrays, exactly like resize::Multiscale;
• PushPull<D, Z, B0, Z, B0, Z, B0, ABS> would have been ambiguous with the 3D specialization PushPull<three, Z, BX, Z, BY, Z, BZ, ABS> as soon as the header was actually included.

PushPullND was already the name used to reach the fallback (it was an alias in pushpull/utils.h, unused anywhere), so call sites are unaffected.

All nine kernels — pull, push, count, grad, hess, pull_backward, push_backward, count_backward, grad_backward — now enumerate the support of the separable basis over D dimensions. Spatial derivatives are obtained by swapping the weight of the differentiated dimension for its derivative:

• coordinate gradients are accumulated over all D dimensions;
• grad_backward builds the full D x D symmetric matrix of second derivatives (d2/dx_d dx_e) instead of the six hard-coded 3D components, and contracts it with the incoming gradient;
• hess writes the compact symmetric layout used elsewhere in jitfields — the diagonal first, then the upper triangle in row-major order, i.e. [xx, yy, zz, xy, xz, yz] in 3D.

lib/pushpull.h now includes nd.h, so the file is compiled from here on.

Tests

jitfields/csrc/tests/test_pushpull_nd.cpp (new) checks every kernel against an independent reference: it enumerates the support explicitly and uses its own B-spline weight / first-derivative / second-derivative formulas, so the kernels are not validated against the very functions they call.

Coverage: D = 1..6, orders 0-3, boundary conditions dct1 / dct2 / dst2 / dft / replicate, mixed order+boundary across dimensions, several channels, non-cubic shapes, and the ABS=true variant. On top of that, the D == 3 results are cross-checked against the specialized 3D kernels.

jitfields/tests/test_pushpull_nd.py (new) compiles and runs it as part of the normal pytest suite; it skips itself if no C++ compiler is on the machine. It does not import torch or cppyy, so it runs anywhere.

$ pytest jitfields/tests/test_pushpull_nd.py
141/141 checks passed
PASSED

Verified with both g++ 13 and clang++. A mutation check (making node_grad always differentiate along dimension 0) turns 51 of the 141 checks red, so the suite does bite.

The existing cpp/pushpull.hpp translation unit was compiled before and after the change to confirm that including nd.h breaks nothing on the existing paths.

Two things deliberately left out of this PR

1. The N-D path is still not reachable from python.bindings/{cpp,cuda}/pushpull.py dispatch to jf::pushpull::pullnd, pushnd, countnd, gradnd, hessnd, pullnd_backward, ... when ndim > 3, but those loop wrappers do not exist in csrc/cpp/pushpull.hpp / csrc/cuda/pushpull.cu (and the template argument lists the bindings build for them do not match the ones the 1D/2D/3D wrappers use). Wiring that up — plus the CUDA side — is a bigger, separate piece of work; this PR fixes and tests the kernels themselves, which is what Fix pushpull backward when D > 3 #4 asks for. Happy to open a follow-up issue for the plumbing.

2. Two pre-existing bugs in 3d.h, spotted while cross-checking, left untouched so as not to bundle unrelated changes:

   • in the generic (any-order) hess, accxz is written to both the zz and the xz slot, so acczz is computed and then dropped — this is why hess is not part of the N-D vs 3D cross-check;
   • the lineargrad_backward returns a zero coordinate gradient. The diagonal second derivatives of a multilinear basis are indeed zero, but the mixed ones (d2/dx dy) are not, so the term is not identically zero. The N-D kernel computes them; the cross-check for the linear case therefore skips grad_backward's gout (flagged with a comment in the test).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XQTnUfyLAVMnY1xvRya8U1

---

Generated by Claude Code

[balbasty]

Closes#4

What was wrong

jitfields/csrc/lib/pushpull/nd.h was a stub. As reported in #4, the backward passes were copy-pasted verbatim from the 3D kernels: they still took the x, nx, sx, y, ny, sy, z, nz, sz argument convention, called utils_x/y/z::gindex, and wrote exactly three coordinate gradients (gout[0], gout[osg], gout[osg*2]) regardless of D.

The forward passes were in no better shape — they did not compile at all:

• sd = s + 8*d where s is never declared (the array is called f),
• db[d] used as the loop bound where the array is called l,
• coord referenced in push/count/grad, which don't have such a parameter,
• *out = static_cast<scalar_t>(acc) in grad, where acc is an array,
• an outright syntax error in grad: acc[dd] += val * (g[8*d + ] * weights[D-1];,
• #endif JF_PUSHPULL_ND (trailing tokens after #endif).

None of this was ever caught because nothing included the file: lib/pushpull.h only pulled in 1d.h, 2d.h and 3d.h.

What this PR does

nd.h is rewritten as a genuine N-dimensional implementation.

PushPullND<D, ABS> is now a class template of its own rather than a partial specialization of PushPull. Two reasons:

• PushPull only carries three (spline order, boundary condition) pairs, so for D > 3 they have to be runtime parameters anyway — they are read from inter[] / bnd[] arrays, exactly like resize::Multiscale;
• PushPull<D, Z, B0, Z, B0, Z, B0, ABS> would have been ambiguous with the 3D specialization PushPull<three, Z, BX, Z, BY, Z, BZ, ABS> as soon as the header was actually included.

PushPullND was already the name used to reach the fallback (it was an alias in pushpull/utils.h, unused anywhere), so call sites are unaffected.

All nine kernels — pull, push, count, grad, hess, pull_backward, push_backward, count_backward, grad_backward — now enumerate the support of the separable basis over D dimensions. Spatial derivatives are obtained by swapping the weight of the differentiated dimension for its derivative:

• coordinate gradients are accumulated over all D dimensions;
• grad_backward builds the full D x D symmetric matrix of second derivatives (d2/dx_d dx_e) instead of the six hard-coded 3D components, and contracts it with the incoming gradient;
• hess writes the compact symmetric layout used elsewhere in jitfields — the diagonal first, then the upper triangle in row-major order, i.e. [xx, yy, zz, xy, xz, yz] in 3D.

lib/pushpull.h now includes nd.h, so the file is compiled from here on.

Tests

jitfields/csrc/tests/test_pushpull_nd.cpp (new) checks every kernel against an independent reference: it enumerates the support explicitly and uses its own B-spline weight / first-derivative / second-derivative formulas, so the kernels are not validated against the very functions they call.

Coverage: D = 1..6, orders 0-3, boundary conditions dct1 / dct2 / dst2 / dft / replicate, mixed order+boundary across dimensions, several channels, non-cubic shapes, and the ABS=true variant. On top of that, the D == 3 results are cross-checked against the specialized 3D kernels.

jitfields/tests/test_pushpull_nd.py (new) compiles and runs it as part of the normal pytest suite; it skips itself if no C++ compiler is on the machine. It does not import torch or cppyy, so it runs anywhere.

$ pytest jitfields/tests/test_pushpull_nd.py
141/141 checks passed
PASSED

Verified with both g++ 13 and clang++. A mutation check (making node_grad always differentiate along dimension 0) turns 51 of the 141 checks red, so the suite does bite.

The existing cpp/pushpull.hpp translation unit was compiled before and after the change to confirm that including nd.h breaks nothing on the existing paths.

Two things deliberately left out of this PR

1. The N-D path is still not reachable from python.bindings/{cpp,cuda}/pushpull.py dispatch to jf::pushpull::pullnd, pushnd, countnd, gradnd, hessnd, pullnd_backward, ... when ndim > 3, but those loop wrappers do not exist in csrc/cpp/pushpull.hpp / csrc/cuda/pushpull.cu (and the template argument lists the bindings build for them do not match the ones the 1D/2D/3D wrappers use). Wiring that up — plus the CUDA side — is a bigger, separate piece of work; this PR fixes and tests the kernels themselves, which is what Fix pushpull backward when D > 3 #4 asks for. Happy to open a follow-up issue for the plumbing.

2. Two pre-existing bugs in 3d.h, spotted while cross-checking, left untouched so as not to bundle unrelated changes:

   • in the generic (any-order) hess, accxz is written to both the zz and the xz slot, so acczz is computed and then dropped — this is why hess is not part of the N-D vs 3D cross-check;
   • the lineargrad_backward returns a zero coordinate gradient. The diagonal second derivatives of a multilinear basis are indeed zero, but the mixed ones (d2/dx dy) are not, so the term is not identically zero. The N-D kernel computes them; the cross-check for the linear case therefore skips grad_backward's gout (flagged with a comment in the test).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XQTnUfyLAVMnY1xvRya8U1

---

Generated by Claude Code

[balbasty]

Closes#4

What was wrong

jitfields/csrc/lib/pushpull/nd.h was a stub. As reported in #4, the backward passes were copy-pasted verbatim from the 3D kernels: they still took the x, nx, sx, y, ny, sy, z, nz, sz argument convention, called utils_x/y/z::gindex, and wrote exactly three coordinate gradients (gout[0], gout[osg], gout[osg*2]) regardless of D.

The forward passes were in no better shape — they did not compile at all:

• sd = s + 8*d where s is never declared (the array is called f),
• db[d] used as the loop bound where the array is called l,
• coord referenced in push/count/grad, which don't have such a parameter,
• *out = static_cast<scalar_t>(acc) in grad, where acc is an array,
• an outright syntax error in grad: acc[dd] += val * (g[8*d + ] * weights[D-1];,
• #endif JF_PUSHPULL_ND (trailing tokens after #endif).

None of this was ever caught because nothing included the file: lib/pushpull.h only pulled in 1d.h, 2d.h and 3d.h.

What this PR does

nd.h is rewritten as a genuine N-dimensional implementation.

PushPullND<D, ABS> is now a class template of its own rather than a partial specialization of PushPull. Two reasons:

• PushPull only carries three (spline order, boundary condition) pairs, so for D > 3 they have to be runtime parameters anyway — they are read from inter[] / bnd[] arrays, exactly like resize::Multiscale;
• PushPull<D, Z, B0, Z, B0, Z, B0, ABS> would have been ambiguous with the 3D specialization PushPull<three, Z, BX, Z, BY, Z, BZ, ABS> as soon as the header was actually included.

PushPullND was already the name used to reach the fallback (it was an alias in pushpull/utils.h, unused anywhere), so call sites are unaffected.

All nine kernels — pull, push, count, grad, hess, pull_backward, push_backward, count_backward, grad_backward — now enumerate the support of the separable basis over D dimensions. Spatial derivatives are obtained by swapping the weight of the differentiated dimension for its derivative:

• coordinate gradients are accumulated over all D dimensions;
• grad_backward builds the full D x D symmetric matrix of second derivatives (d2/dx_d dx_e) instead of the six hard-coded 3D components, and contracts it with the incoming gradient;
• hess writes the compact symmetric layout used elsewhere in jitfields — the diagonal first, then the upper triangle in row-major order, i.e. [xx, yy, zz, xy, xz, yz] in 3D.

lib/pushpull.h now includes nd.h, so the file is compiled from here on.

Tests

jitfields/csrc/tests/test_pushpull_nd.cpp (new) checks every kernel against an independent reference: it enumerates the support explicitly and uses its own B-spline weight / first-derivative / second-derivative formulas, so the kernels are not validated against the very functions they call.

Coverage: D = 1..6, orders 0-3, boundary conditions dct1 / dct2 / dst2 / dft / replicate, mixed order+boundary across dimensions, several channels, non-cubic shapes, and the ABS=true variant. On top of that, the D == 3 results are cross-checked against the specialized 3D kernels.

jitfields/tests/test_pushpull_nd.py (new) compiles and runs it as part of the normal pytest suite; it skips itself if no C++ compiler is on the machine. It does not import torch or cppyy, so it runs anywhere.

$ pytest jitfields/tests/test_pushpull_nd.py
141/141 checks passed
PASSED

Verified with both g++ 13 and clang++. A mutation check (making node_grad always differentiate along dimension 0) turns 51 of the 141 checks red, so the suite does bite.

The existing cpp/pushpull.hpp translation unit was compiled before and after the change to confirm that including nd.h breaks nothing on the existing paths.

Two things deliberately left out of this PR

1. The N-D path is still not reachable from python.bindings/{cpp,cuda}/pushpull.py dispatch to jf::pushpull::pullnd, pushnd, countnd, gradnd, hessnd, pullnd_backward, ... when ndim > 3, but those loop wrappers do not exist in csrc/cpp/pushpull.hpp / csrc/cuda/pushpull.cu (and the template argument lists the bindings build for them do not match the ones the 1D/2D/3D wrappers use). Wiring that up — plus the CUDA side — is a bigger, separate piece of work; this PR fixes and tests the kernels themselves, which is what Fix pushpull backward when D > 3 #4 asks for. Happy to open a follow-up issue for the plumbing.

2. Two pre-existing bugs in 3d.h, spotted while cross-checking, left untouched so as not to bundle unrelated changes:

   • in the generic (any-order) hess, accxz is written to both the zz and the xz slot, so acczz is computed and then dropped — this is why hess is not part of the N-D vs 3D cross-check;
   • the lineargrad_backward returns a zero coordinate gradient. The diagonal second derivatives of a multilinear basis are indeed zero, but the mixed ones (d2/dx dy) are not, so the term is not identically zero. The N-D kernel computes them; the cross-check for the linear case therefore skips grad_backward's gout (flagged with a comment in the test).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XQTnUfyLAVMnY1xvRya8U1

---

Generated by Claude Code

[balbasty]

Closes#4

What was wrong

jitfields/csrc/lib/pushpull/nd.h was a stub. As reported in #4, the backward passes were copy-pasted verbatim from the 3D kernels: they still took the x, nx, sx, y, ny, sy, z, nz, sz argument convention, called utils_x/y/z::gindex, and wrote exactly three coordinate gradients (gout[0], gout[osg], gout[osg*2]) regardless of D.

The forward passes were in no better shape — they did not compile at all:

• sd = s + 8*d where s is never declared (the array is called f),
• db[d] used as the loop bound where the array is called l,
• coord referenced in push/count/grad, which don't have such a parameter,
• *out = static_cast<scalar_t>(acc) in grad, where acc is an array,
• an outright syntax error in grad: acc[dd] += val * (g[8*d + ] * weights[D-1];,
• #endif JF_PUSHPULL_ND (trailing tokens after #endif).

None of this was ever caught because nothing included the file: lib/pushpull.h only pulled in 1d.h, 2d.h and 3d.h.

What this PR does

nd.h is rewritten as a genuine N-dimensional implementation.

PushPullND<D, ABS> is now a class template of its own rather than a partial specialization of PushPull. Two reasons:

• PushPull only carries three (spline order, boundary condition) pairs, so for D > 3 they have to be runtime parameters anyway — they are read from inter[] / bnd[] arrays, exactly like resize::Multiscale;
• PushPull<D, Z, B0, Z, B0, Z, B0, ABS> would have been ambiguous with the 3D specialization PushPull<three, Z, BX, Z, BY, Z, BZ, ABS> as soon as the header was actually included.

PushPullND was already the name used to reach the fallback (it was an alias in pushpull/utils.h, unused anywhere), so call sites are unaffected.

All nine kernels — pull, push, count, grad, hess, pull_backward, push_backward, count_backward, grad_backward — now enumerate the support of the separable basis over D dimensions. Spatial derivatives are obtained by swapping the weight of the differentiated dimension for its derivative:

• coordinate gradients are accumulated over all D dimensions;
• grad_backward builds the full D x D symmetric matrix of second derivatives (d2/dx_d dx_e) instead of the six hard-coded 3D components, and contracts it with the incoming gradient;
• hess writes the compact symmetric layout used elsewhere in jitfields — the diagonal first, then the upper triangle in row-major order, i.e. [xx, yy, zz, xy, xz, yz] in 3D.

lib/pushpull.h now includes nd.h, so the file is compiled from here on.

Tests

jitfields/csrc/tests/test_pushpull_nd.cpp (new) checks every kernel against an independent reference: it enumerates the support explicitly and uses its own B-spline weight / first-derivative / second-derivative formulas, so the kernels are not validated against the very functions they call.

Coverage: D = 1..6, orders 0-3, boundary conditions dct1 / dct2 / dst2 / dft / replicate, mixed order+boundary across dimensions, several channels, non-cubic shapes, and the ABS=true variant. On top of that, the D == 3 results are cross-checked against the specialized 3D kernels.

jitfields/tests/test_pushpull_nd.py (new) compiles and runs it as part of the normal pytest suite; it skips itself if no C++ compiler is on the machine. It does not import torch or cppyy, so it runs anywhere.

$ pytest jitfields/tests/test_pushpull_nd.py
141/141 checks passed
PASSED

Verified with both g++ 13 and clang++. A mutation check (making node_grad always differentiate along dimension 0) turns 51 of the 141 checks red, so the suite does bite.

The existing cpp/pushpull.hpp translation unit was compiled before and after the change to confirm that including nd.h breaks nothing on the existing paths.

Two things deliberately left out of this PR

1. The N-D path is still not reachable from python.bindings/{cpp,cuda}/pushpull.py dispatch to jf::pushpull::pullnd, pushnd, countnd, gradnd, hessnd, pullnd_backward, ... when ndim > 3, but those loop wrappers do not exist in csrc/cpp/pushpull.hpp / csrc/cuda/pushpull.cu (and the template argument lists the bindings build for them do not match the ones the 1D/2D/3D wrappers use). Wiring that up — plus the CUDA side — is a bigger, separate piece of work; this PR fixes and tests the kernels themselves, which is what Fix pushpull backward when D > 3 #4 asks for. Happy to open a follow-up issue for the plumbing.

2. Two pre-existing bugs in 3d.h, spotted while cross-checking, left untouched so as not to bundle unrelated changes:

   • in the generic (any-order) hess, accxz is written to both the zz and the xz slot, so acczz is computed and then dropped — this is why hess is not part of the N-D vs 3D cross-check;
   • the lineargrad_backward returns a zero coordinate gradient. The diagonal second derivatives of a multilinear basis are indeed zero, but the mixed ones (d2/dx dy) are not, so the term is not identically zero. The N-D kernel computes them; the cross-check for the linear case therefore skips grad_backward's gout (flagged with a comment in the test).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XQTnUfyLAVMnY1xvRya8U1

---

Generated by Claude Code

[balbasty]

Closes#4

What was wrong

jitfields/csrc/lib/pushpull/nd.h was a stub. As reported in #4, the backward passes were copy-pasted verbatim from the 3D kernels: they still took the x, nx, sx, y, ny, sy, z, nz, sz argument convention, called utils_x/y/z::gindex, and wrote exactly three coordinate gradients (gout[0], gout[osg], gout[osg*2]) regardless of D.

The forward passes were in no better shape — they did not compile at all:

• sd = s + 8*d where s is never declared (the array is called f),
• db[d] used as the loop bound where the array is called l,
• coord referenced in push/count/grad, which don't have such a parameter,
• *out = static_cast<scalar_t>(acc) in grad, where acc is an array,
• an outright syntax error in grad: acc[dd] += val * (g[8*d + ] * weights[D-1];,
• #endif JF_PUSHPULL_ND (trailing tokens after #endif).

None of this was ever caught because nothing included the file: lib/pushpull.h only pulled in 1d.h, 2d.h and 3d.h.

What this PR does

nd.h is rewritten as a genuine N-dimensional implementation.

PushPullND<D, ABS> is now a class template of its own rather than a partial specialization of PushPull. Two reasons:

• PushPull only carries three (spline order, boundary condition) pairs, so for D > 3 they have to be runtime parameters anyway — they are read from inter[] / bnd[] arrays, exactly like resize::Multiscale;
• PushPull<D, Z, B0, Z, B0, Z, B0, ABS> would have been ambiguous with the 3D specialization PushPull<three, Z, BX, Z, BY, Z, BZ, ABS> as soon as the header was actually included.

PushPullND was already the name used to reach the fallback (it was an alias in pushpull/utils.h, unused anywhere), so call sites are unaffected.

All nine kernels — pull, push, count, grad, hess, pull_backward, push_backward, count_backward, grad_backward — now enumerate the support of the separable basis over D dimensions. Spatial derivatives are obtained by swapping the weight of the differentiated dimension for its derivative:

• coordinate gradients are accumulated over all D dimensions;
• grad_backward builds the full D x D symmetric matrix of second derivatives (d2/dx_d dx_e) instead of the six hard-coded 3D components, and contracts it with the incoming gradient;
• hess writes the compact symmetric layout used elsewhere in jitfields — the diagonal first, then the upper triangle in row-major order, i.e. [xx, yy, zz, xy, xz, yz] in 3D.

lib/pushpull.h now includes nd.h, so the file is compiled from here on.

Tests

jitfields/csrc/tests/test_pushpull_nd.cpp (new) checks every kernel against an independent reference: it enumerates the support explicitly and uses its own B-spline weight / first-derivative / second-derivative formulas, so the kernels are not validated against the very functions they call.

Coverage: D = 1..6, orders 0-3, boundary conditions dct1 / dct2 / dst2 / dft / replicate, mixed order+boundary across dimensions, several channels, non-cubic shapes, and the ABS=true variant. On top of that, the D == 3 results are cross-checked against the specialized 3D kernels.

jitfields/tests/test_pushpull_nd.py (new) compiles and runs it as part of the normal pytest suite; it skips itself if no C++ compiler is on the machine. It does not import torch or cppyy, so it runs anywhere.

$ pytest jitfields/tests/test_pushpull_nd.py
141/141 checks passed
PASSED

Verified with both g++ 13 and clang++. A mutation check (making node_grad always differentiate along dimension 0) turns 51 of the 141 checks red, so the suite does bite.

The existing cpp/pushpull.hpp translation unit was compiled before and after the change to confirm that including nd.h breaks nothing on the existing paths.

Two things deliberately left out of this PR

1. The N-D path is still not reachable from python.bindings/{cpp,cuda}/pushpull.py dispatch to jf::pushpull::pullnd, pushnd, countnd, gradnd, hessnd, pullnd_backward, ... when ndim > 3, but those loop wrappers do not exist in csrc/cpp/pushpull.hpp / csrc/cuda/pushpull.cu (and the template argument lists the bindings build for them do not match the ones the 1D/2D/3D wrappers use). Wiring that up — plus the CUDA side — is a bigger, separate piece of work; this PR fixes and tests the kernels themselves, which is what Fix pushpull backward when D > 3 #4 asks for. Happy to open a follow-up issue for the plumbing.

2. Two pre-existing bugs in 3d.h, spotted while cross-checking, left untouched so as not to bundle unrelated changes:

   • in the generic (any-order) hess, accxz is written to both the zz and the xz slot, so acczz is computed and then dropped — this is why hess is not part of the N-D vs 3D cross-check;
   • the lineargrad_backward returns a zero coordinate gradient. The diagonal second derivatives of a multilinear basis are indeed zero, but the mixed ones (d2/dx dy) are not, so the term is not identically zero. The N-D kernel computes them; the cross-check for the linear case therefore skips grad_backward's gout (flagged with a comment in the test).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XQTnUfyLAVMnY1xvRya8U1

---

Generated by Claude Code

[balbasty]

Closes#4

What was wrong

jitfields/csrc/lib/pushpull/nd.h was a stub. As reported in #4, the backward passes were copy-pasted verbatim from the 3D kernels: they still took the x, nx, sx, y, ny, sy, z, nz, sz argument convention, called utils_x/y/z::gindex, and wrote exactly three coordinate gradients (gout[0], gout[osg], gout[osg*2]) regardless of D.

The forward passes were in no better shape — they did not compile at all:

• sd = s + 8*d where s is never declared (the array is called f),
• db[d] used as the loop bound where the array is called l,
• coord referenced in push/count/grad, which don't have such a parameter,
• *out = static_cast<scalar_t>(acc) in grad, where acc is an array,
• an outright syntax error in grad: acc[dd] += val * (g[8*d + ] * weights[D-1];,
• #endif JF_PUSHPULL_ND (trailing tokens after #endif).

None of this was ever caught because nothing included the file: lib/pushpull.h only pulled in 1d.h, 2d.h and 3d.h.

What this PR does

nd.h is rewritten as a genuine N-dimensional implementation.

PushPullND<D, ABS> is now a class template of its own rather than a partial specialization of PushPull. Two reasons:

• PushPull only carries three (spline order, boundary condition) pairs, so for D > 3 they have to be runtime parameters anyway — they are read from inter[] / bnd[] arrays, exactly like resize::Multiscale;
• PushPull<D, Z, B0, Z, B0, Z, B0, ABS> would have been ambiguous with the 3D specialization PushPull<three, Z, BX, Z, BY, Z, BZ, ABS> as soon as the header was actually included.

PushPullND was already the name used to reach the fallback (it was an alias in pushpull/utils.h, unused anywhere), so call sites are unaffected.

All nine kernels — pull, push, count, grad, hess, pull_backward, push_backward, count_backward, grad_backward — now enumerate the support of the separable basis over D dimensions. Spatial derivatives are obtained by swapping the weight of the differentiated dimension for its derivative:

• coordinate gradients are accumulated over all D dimensions;
• grad_backward builds the full D x D symmetric matrix of second derivatives (d2/dx_d dx_e) instead of the six hard-coded 3D components, and contracts it with the incoming gradient;
• hess writes the compact symmetric layout used elsewhere in jitfields — the diagonal first, then the upper triangle in row-major order, i.e. [xx, yy, zz, xy, xz, yz] in 3D.

lib/pushpull.h now includes nd.h, so the file is compiled from here on.

Tests

jitfields/csrc/tests/test_pushpull_nd.cpp (new) checks every kernel against an independent reference: it enumerates the support explicitly and uses its own B-spline weight / first-derivative / second-derivative formulas, so the kernels are not validated against the very functions they call.

Coverage: D = 1..6, orders 0-3, boundary conditions dct1 / dct2 / dst2 / dft / replicate, mixed order+boundary across dimensions, several channels, non-cubic shapes, and the ABS=true variant. On top of that, the D == 3 results are cross-checked against the specialized 3D kernels.

jitfields/tests/test_pushpull_nd.py (new) compiles and runs it as part of the normal pytest suite; it skips itself if no C++ compiler is on the machine. It does not import torch or cppyy, so it runs anywhere.

$ pytest jitfields/tests/test_pushpull_nd.py
141/141 checks passed
PASSED

Verified with both g++ 13 and clang++. A mutation check (making node_grad always differentiate along dimension 0) turns 51 of the 141 checks red, so the suite does bite.

The existing cpp/pushpull.hpp translation unit was compiled before and after the change to confirm that including nd.h breaks nothing on the existing paths.

Two things deliberately left out of this PR

1. The N-D path is still not reachable from python.bindings/{cpp,cuda}/pushpull.py dispatch to jf::pushpull::pullnd, pushnd, countnd, gradnd, hessnd, pullnd_backward, ... when ndim > 3, but those loop wrappers do not exist in csrc/cpp/pushpull.hpp / csrc/cuda/pushpull.cu (and the template argument lists the bindings build for them do not match the ones the 1D/2D/3D wrappers use). Wiring that up — plus the CUDA side — is a bigger, separate piece of work; this PR fixes and tests the kernels themselves, which is what Fix pushpull backward when D > 3 #4 asks for. Happy to open a follow-up issue for the plumbing.

2. Two pre-existing bugs in 3d.h, spotted while cross-checking, left untouched so as not to bundle unrelated changes:

   • in the generic (any-order) hess, accxz is written to both the zz and the xz slot, so acczz is computed and then dropped — this is why hess is not part of the N-D vs 3D cross-check;
   • the lineargrad_backward returns a zero coordinate gradient. The diagonal second derivatives of a multilinear basis are indeed zero, but the mixed ones (d2/dx dy) are not, so the term is not identically zero. The N-D kernel computes them; the cross-check for the linear case therefore skips grad_backward's gout (flagged with a comment in the test).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XQTnUfyLAVMnY1xvRya8U1

---

Generated by Claude Code

[balbasty]

Closes#4

What was wrong

jitfields/csrc/lib/pushpull/nd.h was a stub. As reported in #4, the backward passes were copy-pasted verbatim from the 3D kernels: they still took the x, nx, sx, y, ny, sy, z, nz, sz argument convention, called utils_x/y/z::gindex, and wrote exactly three coordinate gradients (gout[0], gout[osg], gout[osg*2]) regardless of D.

The forward passes were in no better shape — they did not compile at all:

• sd = s + 8*d where s is never declared (the array is called f),
• db[d] used as the loop bound where the array is called l,
• coord referenced in push/count/grad, which don't have such a parameter,
• *out = static_cast<scalar_t>(acc) in grad, where acc is an array,
• an outright syntax error in grad: acc[dd] += val * (g[8*d + ] * weights[D-1];,
• #endif JF_PUSHPULL_ND (trailing tokens after #endif).

None of this was ever caught because nothing included the file: lib/pushpull.h only pulled in 1d.h, 2d.h and 3d.h.

What this PR does

nd.h is rewritten as a genuine N-dimensional implementation.

PushPullND<D, ABS> is now a class template of its own rather than a partial specialization of PushPull. Two reasons:

• PushPull only carries three (spline order, boundary condition) pairs, so for D > 3 they have to be runtime parameters anyway — they are read from inter[] / bnd[] arrays, exactly like resize::Multiscale;
• PushPull<D, Z, B0, Z, B0, Z, B0, ABS> would have been ambiguous with the 3D specialization PushPull<three, Z, BX, Z, BY, Z, BZ, ABS> as soon as the header was actually included.

PushPullND was already the name used to reach the fallback (it was an alias in pushpull/utils.h, unused anywhere), so call sites are unaffected.

All nine kernels — pull, push, count, grad, hess, pull_backward, push_backward, count_backward, grad_backward — now enumerate the support of the separable basis over D dimensions. Spatial derivatives are obtained by swapping the weight of the differentiated dimension for its derivative:

• coordinate gradients are accumulated over all D dimensions;
• grad_backward builds the full D x D symmetric matrix of second derivatives (d2/dx_d dx_e) instead of the six hard-coded 3D components, and contracts it with the incoming gradient;
• hess writes the compact symmetric layout used elsewhere in jitfields — the diagonal first, then the upper triangle in row-major order, i.e. [xx, yy, zz, xy, xz, yz] in 3D.

lib/pushpull.h now includes nd.h, so the file is compiled from here on.

Tests

jitfields/csrc/tests/test_pushpull_nd.cpp (new) checks every kernel against an independent reference: it enumerates the support explicitly and uses its own B-spline weight / first-derivative / second-derivative formulas, so the kernels are not validated against the very functions they call.

Coverage: D = 1..6, orders 0-3, boundary conditions dct1 / dct2 / dst2 / dft / replicate, mixed order+boundary across dimensions, several channels, non-cubic shapes, and the ABS=true variant. On top of that, the D == 3 results are cross-checked against the specialized 3D kernels.

jitfields/tests/test_pushpull_nd.py (new) compiles and runs it as part of the normal pytest suite; it skips itself if no C++ compiler is on the machine. It does not import torch or cppyy, so it runs anywhere.

$ pytest jitfields/tests/test_pushpull_nd.py
141/141 checks passed
PASSED

Verified with both g++ 13 and clang++. A mutation check (making node_grad always differentiate along dimension 0) turns 51 of the 141 checks red, so the suite does bite.

The existing cpp/pushpull.hpp translation unit was compiled before and after the change to confirm that including nd.h breaks nothing on the existing paths.

Two things deliberately left out of this PR

1. The N-D path is still not reachable from python.bindings/{cpp,cuda}/pushpull.py dispatch to jf::pushpull::pullnd, pushnd, countnd, gradnd, hessnd, pullnd_backward, ... when ndim > 3, but those loop wrappers do not exist in csrc/cpp/pushpull.hpp / csrc/cuda/pushpull.cu (and the template argument lists the bindings build for them do not match the ones the 1D/2D/3D wrappers use). Wiring that up — plus the CUDA side — is a bigger, separate piece of work; this PR fixes and tests the kernels themselves, which is what Fix pushpull backward when D > 3 #4 asks for. Happy to open a follow-up issue for the plumbing.

2. Two pre-existing bugs in 3d.h, spotted while cross-checking, left untouched so as not to bundle unrelated changes:

   • in the generic (any-order) hess, accxz is written to both the zz and the xz slot, so acczz is computed and then dropped — this is why hess is not part of the N-D vs 3D cross-check;
   • the lineargrad_backward returns a zero coordinate gradient. The diagonal second derivatives of a multilinear basis are indeed zero, but the mixed ones (d2/dx dy) are not, so the term is not identically zero. The N-D kernel computes them; the cross-check for the linear case therefore skips grad_backward's gout (flagged with a comment in the test).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XQTnUfyLAVMnY1xvRya8U1

---

Generated by Claude Code

[balbasty]

Closes#4

What was wrong

jitfields/csrc/lib/pushpull/nd.h was a stub. As reported in #4, the backward passes were copy-pasted verbatim from the 3D kernels: they still took the x, nx, sx, y, ny, sy, z, nz, sz argument convention, called utils_x/y/z::gindex, and wrote exactly three coordinate gradients (gout[0], gout[osg], gout[osg*2]) regardless of D.

The forward passes were in no better shape — they did not compile at all:

• sd = s + 8*d where s is never declared (the array is called f),
• db[d] used as the loop bound where the array is called l,
• coord referenced in push/count/grad, which don't have such a parameter,
• *out = static_cast<scalar_t>(acc) in grad, where acc is an array,
• an outright syntax error in grad: acc[dd] += val * (g[8*d + ] * weights[D-1];,
• #endif JF_PUSHPULL_ND (trailing tokens after #endif).

None of this was ever caught because nothing included the file: lib/pushpull.h only pulled in 1d.h, 2d.h and 3d.h.

What this PR does

nd.h is rewritten as a genuine N-dimensional implementation.

PushPullND<D, ABS> is now a class template of its own rather than a partial specialization of PushPull. Two reasons:

• PushPull only carries three (spline order, boundary condition) pairs, so for D > 3 they have to be runtime parameters anyway — they are read from inter[] / bnd[] arrays, exactly like resize::Multiscale;
• PushPull<D, Z, B0, Z, B0, Z, B0, ABS> would have been ambiguous with the 3D specialization PushPull<three, Z, BX, Z, BY, Z, BZ, ABS> as soon as the header was actually included.

PushPullND was already the name used to reach the fallback (it was an alias in pushpull/utils.h, unused anywhere), so call sites are unaffected.

All nine kernels — pull, push, count, grad, hess, pull_backward, push_backward, count_backward, grad_backward — now enumerate the support of the separable basis over D dimensions. Spatial derivatives are obtained by swapping the weight of the differentiated dimension for its derivative:

• coordinate gradients are accumulated over all D dimensions;
• grad_backward builds the full D x D symmetric matrix of second derivatives (d2/dx_d dx_e) instead of the six hard-coded 3D components, and contracts it with the incoming gradient;
• hess writes the compact symmetric layout used elsewhere in jitfields — the diagonal first, then the upper triangle in row-major order, i.e. [xx, yy, zz, xy, xz, yz] in 3D.

lib/pushpull.h now includes nd.h, so the file is compiled from here on.

Tests

jitfields/csrc/tests/test_pushpull_nd.cpp (new) checks every kernel against an independent reference: it enumerates the support explicitly and uses its own B-spline weight / first-derivative / second-derivative formulas, so the kernels are not validated against the very functions they call.

Coverage: D = 1..6, orders 0-3, boundary conditions dct1 / dct2 / dst2 / dft / replicate, mixed order+boundary across dimensions, several channels, non-cubic shapes, and the ABS=true variant. On top of that, the D == 3 results are cross-checked against the specialized 3D kernels.

jitfields/tests/test_pushpull_nd.py (new) compiles and runs it as part of the normal pytest suite; it skips itself if no C++ compiler is on the machine. It does not import torch or cppyy, so it runs anywhere.

$ pytest jitfields/tests/test_pushpull_nd.py
141/141 checks passed
PASSED

Verified with both g++ 13 and clang++. A mutation check (making node_grad always differentiate along dimension 0) turns 51 of the 141 checks red, so the suite does bite.

The existing cpp/pushpull.hpp translation unit was compiled before and after the change to confirm that including nd.h breaks nothing on the existing paths.

Two things deliberately left out of this PR

1. The N-D path is still not reachable from python.bindings/{cpp,cuda}/pushpull.py dispatch to jf::pushpull::pullnd, pushnd, countnd, gradnd, hessnd, pullnd_backward, ... when ndim > 3, but those loop wrappers do not exist in csrc/cpp/pushpull.hpp / csrc/cuda/pushpull.cu (and the template argument lists the bindings build for them do not match the ones the 1D/2D/3D wrappers use). Wiring that up — plus the CUDA side — is a bigger, separate piece of work; this PR fixes and tests the kernels themselves, which is what Fix pushpull backward when D > 3 #4 asks for. Happy to open a follow-up issue for the plumbing.

2. Two pre-existing bugs in 3d.h, spotted while cross-checking, left untouched so as not to bundle unrelated changes:

   • in the generic (any-order) hess, accxz is written to both the zz and the xz slot, so acczz is computed and then dropped — this is why hess is not part of the N-D vs 3D cross-check;
   • the lineargrad_backward returns a zero coordinate gradient. The diagonal second derivatives of a multilinear basis are indeed zero, but the mixed ones (d2/dx dy) are not, so the term is not identically zero. The N-D kernel computes them; the cross-check for the linear case therefore skips grad_backward's gout (flagged with a comment in the test).

🤖 Generated with Claude Code

https://claude.ai/code/session_01XQTnUfyLAVMnY1xvRya8U1

---

Generated by Claude Code