From 4e5bfea9b647ba511a73edab1c5e341d3c868af2 Mon Sep 17 00:00:00 2001 From: Claude Date: Sat, 1 Aug 2026 06:09:19 +0000 Subject: [PATCH] Fix field_kernel_* dispatching to flow_kernel, and kernel_sub not subtracting Two independent bugs in the `*_kernel_add` / `*_kernel_sub` wrappers, found while porting these kernels to fastfields (https://github.com/fastfields). 1. `field_kernel_add` and `field_kernel_add_` called `impl.flow_kernel` instead of `impl.field_kernel`. This is a hard failure, not a silent one: the backend signature is flow_kernel(out, bound, voxel_size, absolute, membrane, bending, shears, div, op='') -- 8 required positional parameters -- and the call site passes 7, so all four public entry points raised TypeError: flow_kernel() missing 1 required positional argument: 'div' The stale "See `flow_kernel`" docstrings on the two `_add` variants (every sibling says "See `field_kernel`") look like the origin of the slip, and are corrected too. Note that `field_kernel` and `flow_kernel` produce bit-identical channel-0 kernels under the default isotropic voxel size and only diverge once `voxel_size` is anisotropic -- so even a call with the right arity would have gone unnoticed under a default-argument test. 2. `field_kernel_sub`, `field_kernel_sub_`, `flow_kernel_sub` and `flow_kernel_sub_` forwarded to their `_add` sibling without the trailing `_sub=True` flag, so `op` was computed as 'add' and they added the kernel instead of subtracting it. Every other `_sub` wrapper in these two modules (field/flow x matvec/diag x in-place/out-of-place) does pass `True`; these four were the only omissions. Adds `test_kernel_add_sub` to tests/test_reg_field.py and tests/test_reg_flow.py, asserting `kernel_add(inp) == inp + kernel` and `kernel_sub(inp) == inp - kernel` for the out-of-place and in-place variants. The voxel size is deliberately anisotropic ([1.5, 1., 2.5][:dim]) so that a field-vs-flow cross-wiring cannot hide behind the isotropic coincidence described above. Before: 12 failed (field: TypeError; flow: the `sub` assertion). After: 36 passed across both reg test modules. Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_016AjQcY78NgbagPSbPJRr6Z --- jitfields/_regularisers_fields.py | 20 +++++++------- jitfields/_regularisers_flows.py | 4 +-- jitfields/tests/test_reg_field.py | 45 ++++++++++++++++++++++++++++++- jitfields/tests/test_reg_flow.py | 45 ++++++++++++++++++++++++++++++- 4 files changed, 100 insertions(+), 14 deletions(-) diff --git a/jitfields/_regularisers_fields.py b/jitfields/_regularisers_fields.py index 41d0473..9bb6f88 100755 --- a/jitfields/_regularisers_fields.py +++ b/jitfields/_regularisers_fields.py @@ -336,7 +336,7 @@ def field_kernel_add( out: Optional[Tensor] = None, _sub: bool = False, ) -> Tensor: - """See `flow_kernel`""" + """See `field_kernel`""" # allocate output if out is None: out = inp.clone() @@ -352,9 +352,9 @@ def field_kernel_add( # forward impl = cuda_impl if out.is_cuda else cpu_impl - impl.flow_kernel(out, bound, voxel_size, - absolute, membrane, bending, - 'sub' if _sub else 'add') + impl.field_kernel(out, bound, voxel_size, + absolute, membrane, bending, + 'sub' if _sub else 'add') return out @@ -369,7 +369,7 @@ def field_kernel_add_( voxel_size: OneOrSeveral[float] = 1, _sub: bool = False, ) -> Tensor: - """See `flow_kernel`""" + """See `field_kernel`""" nc = inp.shape[-1] bound = ensure_list(bound, ndim) voxel_size = make_vector(voxel_size, ndim).tolist() @@ -379,9 +379,9 @@ def field_kernel_add_( # forward impl = cuda_impl if inp.is_cuda else cpu_impl - impl.flow_kernel(inp, bound, voxel_size, - absolute, membrane, bending, - 'sub' if _sub else 'add') + impl.field_kernel(inp, bound, voxel_size, + absolute, membrane, bending, + 'sub' if _sub else 'add') return inp @@ -397,7 +397,7 @@ def field_kernel_sub( ) -> Tensor: """See `field_kernel`""" return field_kernel_add(ndim, inp, absolute, membrane, bending, - bound, voxel_size, out) + bound, voxel_size, out, True) def field_kernel_sub_( @@ -411,7 +411,7 @@ def field_kernel_sub_( ) -> Tensor: """See `field_kernel`""" return field_kernel_add_(ndim, inp, absolute, membrane, bending, - bound, voxel_size) + bound, voxel_size, True) def field_diag( diff --git a/jitfields/_regularisers_flows.py b/jitfields/_regularisers_flows.py index 766b2f4..014b2d2 100755 --- a/jitfields/_regularisers_flows.py +++ b/jitfields/_regularisers_flows.py @@ -425,7 +425,7 @@ def flow_kernel_sub( from `inp`. """ return flow_kernel_add(inp, absolute, membrane, bending, shears, div, - bound, voxel_size, out) + bound, voxel_size, out, True) def flow_kernel_sub_( @@ -443,7 +443,7 @@ def flow_kernel_sub_( from `inp` (inplace). """ return flow_kernel_add_(inp, absolute, membrane, bending, shears, div, - bound, voxel_size) + bound, voxel_size, True) def flow_diag( diff --git a/jitfields/tests/test_reg_field.py b/jitfields/tests/test_reg_field.py index f61466a..f33d9a7 100755 --- a/jitfields/tests/test_reg_field.py +++ b/jitfields/tests/test_reg_field.py @@ -1,4 +1,8 @@ -from jitfields.regularisers import field_kernel, field_matvec +from jitfields.regularisers import ( + field_kernel, field_matvec, + field_kernel_add, field_kernel_add_, + field_kernel_sub, field_kernel_sub_, +) from .ref_kernels import kernels1, kernels2, kernels3 from .utils import get_test_devices, init_device import torch @@ -44,6 +48,45 @@ def test_kernel(device, dim): f"{pred_bending}\n{kernel_bending}" +@pytest.mark.parametrize("device", devices) +@pytest.mark.parametrize("dim", dims) +def test_kernel_add_sub(device, dim): + """`field_kernel_{add,sub}[_]` must agree with `field_kernel`. + + The voxel size is deliberately anisotropic: under the default + (isotropic) voxel size the field and flow kernels happen to be + numerically identical, so a variant that computes the *flow* kernel + instead of the *field* kernel would go unnoticed. + """ + device = init_device(device) + backend = dict(device=device, dtype=torch.float32) + + voxel_size = [1.5, 1., 2.5][:dim] + absolute, membrane, bending = 0.3, 1., 0.2 + + kernel = field_kernel([5]*dim, absolute=absolute, membrane=membrane, + bending=bending, voxel_size=voxel_size, **backend) + inp = torch.randn(kernel.shape, **backend) + + add, sub = inp + kernel, inp - kernel + + pred = field_kernel_add(dim, inp, absolute, membrane, bending, + voxel_size=voxel_size) + assert torch.allclose(pred, add), f"{pred}\n{add}" + + pred = field_kernel_sub(dim, inp, absolute, membrane, bending, + voxel_size=voxel_size) + assert torch.allclose(pred, sub), f"{pred}\n{sub}" + + pred = field_kernel_add_(dim, inp.clone(), absolute, membrane, bending, + voxel_size=voxel_size) + assert torch.allclose(pred, add), f"{pred}\n{add}" + + pred = field_kernel_sub_(dim, inp.clone(), absolute, membrane, bending, + voxel_size=voxel_size) + assert torch.allclose(pred, sub), f"{pred}\n{sub}" + + @pytest.mark.parametrize("device", devices) @pytest.mark.parametrize("dim", dims) def test_matvec(device, dim): diff --git a/jitfields/tests/test_reg_flow.py b/jitfields/tests/test_reg_flow.py index b2416b4..2e6a6d2 100755 --- a/jitfields/tests/test_reg_flow.py +++ b/jitfields/tests/test_reg_flow.py @@ -1,4 +1,8 @@ -from jitfields.regularisers import flow_kernel, flow_matvec +from jitfields.regularisers import ( + flow_kernel, flow_matvec, + flow_kernel_add, flow_kernel_add_, + flow_kernel_sub, flow_kernel_sub_, +) from .ref_kernels import kernels1, kernels2, kernels3 from .utils import get_test_devices, init_device import torch @@ -55,6 +59,45 @@ def test_kernel(device, dim): assert torch.allclose(pred_div, kernel_div), f"{pred_div}\n{kernel_div}" +@pytest.mark.parametrize("device", devices) +@pytest.mark.parametrize("dim", dims) +def test_kernel_add_sub(device, dim): + """`flow_kernel_{add,sub}[_]` must agree with `flow_kernel`. + + The voxel size is deliberately anisotropic: under the default + (isotropic) voxel size the field and flow kernels happen to be + numerically identical, so a variant that computes the wrong kernel + would go unnoticed. + """ + device = init_device(device) + backend = dict(device=device, dtype=torch.float32) + + voxel_size = [1.5, 1., 2.5][:dim] + absolute, membrane, bending = 0.3, 1., 0.2 + + kernel = flow_kernel([5]*dim, absolute=absolute, membrane=membrane, + bending=bending, voxel_size=voxel_size, **backend) + inp = torch.randn(kernel.shape, **backend) + + add, sub = inp + kernel, inp - kernel + + pred = flow_kernel_add(inp, absolute, membrane, bending, + voxel_size=voxel_size) + assert torch.allclose(pred, add), f"{pred}\n{add}" + + pred = flow_kernel_sub(inp, absolute, membrane, bending, + voxel_size=voxel_size) + assert torch.allclose(pred, sub), f"{pred}\n{sub}" + + pred = flow_kernel_add_(inp.clone(), absolute, membrane, bending, + voxel_size=voxel_size) + assert torch.allclose(pred, add), f"{pred}\n{add}" + + pred = flow_kernel_sub_(inp.clone(), absolute, membrane, bending, + voxel_size=voxel_size) + assert torch.allclose(pred, sub), f"{pred}\n{sub}" + + @pytest.mark.parametrize("device", devices) @pytest.mark.parametrize("dim", dims) def test_matvec(device, dim):