From dc14fccae0d67121f47db3e590e30747e1181d75 Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Tue, 22 Sep 2026 16:44:54 +0100 Subject: [PATCH 1/2] Preserve coordinate gradients at integer interpolation points --- dm_pix/_src/interpolation.py | 2 +- dm_pix/_src/interpolation_gradient_test.py | 89 ++++++++++++++++++++++ 2 files changed, 90 insertions(+), 1 deletion(-) create mode 100644 dm_pix/_src/interpolation_gradient_test.py diff --git a/dm_pix/_src/interpolation.py b/dm_pix/_src/interpolation.py index 9a023d1..669a217 100644 --- a/dm_pix/_src/interpolation.py +++ b/dm_pix/_src/interpolation.py @@ -44,7 +44,7 @@ def _make_linear_interpolation_indices_nd( The lower and upper indices of `coordinates` and their weights. """ lower = jnp.floor(coordinates).astype(jnp.int32) - upper = jnp.ceil(coordinates).astype(jnp.int32) + upper = (jnp.floor(coordinates) + 1).astype(jnp.int32) weights = coordinates - lower # Expand dimensions for `shape` to allow broadcasting it to every coordinate. diff --git a/dm_pix/_src/interpolation_gradient_test.py b/dm_pix/_src/interpolation_gradient_test.py new file mode 100644 index 0000000..a340e13 --- /dev/null +++ b/dm_pix/_src/interpolation_gradient_test.py @@ -0,0 +1,89 @@ +# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Coordinate gradients of linear interpolation at exact grid points.""" + +from absl.testing import absltest +from absl.testing import parameterized +import dm_pix +import jax +import jax.numpy as jnp +import numpy as np + + +class InterpolationGradientTest(parameterized.TestCase): + + @parameterized.product( + ndim=(1, 2, 3), flattened=(False, True), constant=(False, True) + ) + def test_affine_ramp_gradient(self, ndim, flattened, constant): + shape = (5,) * ndim + slopes = jnp.arange(1, ndim + 1, dtype=jnp.float32) / 32 + grid = jnp.indices(shape, dtype=jnp.float32) + volume = jnp.einsum("i,i...->...", slopes, grid) + coordinates = jnp.full((ndim,), 2.0, dtype=jnp.float32) + fn = ( + dm_pix.flat_nd_linear_interpolate_constant + if constant + else dm_pix.flat_nd_linear_interpolate + ) + kwargs = {"unflattened_vol_shape": shape} if flattened else {} + values = volume.reshape(-1) if flattened else volume + + def interpolate(x): + return fn(values, x, **kwargs) + + expected = jnp.dot(slopes, coordinates) + np.testing.assert_allclose(interpolate(coordinates), expected, atol=1e-7) + for grad in (jax.grad(interpolate), jax.jit(jax.grad(interpolate))): + np.testing.assert_allclose(grad(coordinates), slopes, atol=1e-7) + tangent = jnp.ones(ndim) + _, derivative = jax.jvp(interpolate, (coordinates,), (tangent,)) + np.testing.assert_allclose(derivative, jnp.sum(slopes), atol=1e-7) + + @parameterized.parameters("nearest", "constant") + def test_identity_affine_translation_has_gradient(self, mode): + rows, columns = jnp.indices((5, 5), dtype=jnp.float32) + image = (rows / 16 + columns / 8)[..., None] + + def sample(offset): + return dm_pix.affine_transform( + image, jnp.eye(3), offset=offset, mode=mode + )[2, 2, 0] + + offset = jnp.zeros(3) + for grad in (jax.grad(sample), jax.jit(jax.grad(sample))): + np.testing.assert_allclose(grad(offset), [1 / 16, 1 / 8, 0], atol=1e-7) + + def test_clipped_coordinates_and_singleton_axes(self): + volume = jnp.arange(5, dtype=jnp.float32)[None, :] / 4 + + def fn(x): + return dm_pix.flat_nd_linear_interpolate(volume, x) + + for coordinate, expected in (([-1.0, -1.0], 0.0), ([1.0, 6.0], 1.0)): + coordinate = jnp.asarray(coordinate) + np.testing.assert_allclose(fn(coordinate), expected) + np.testing.assert_array_equal(jax.grad(fn)(coordinate), jnp.zeros(2)) + + def test_integer_samples_preserve_value_gradients(self): + volume = jnp.arange(5, dtype=jnp.float32) / 4 + + def fn(values): + return dm_pix.flat_nd_linear_interpolate(values, jnp.array([2.0])) + + np.testing.assert_array_equal(jax.grad(fn)(volume), [0, 0, 1, 0, 0]) + + +if __name__ == "__main__": + absltest.main() From a2fc9aa6119f0e78417a91c02d8bc33c8eeabc80 Mon Sep 17 00:00:00 2001 From: Sylvester Kaczmarek <16242628+sylvesterkaczmarek@users.noreply.github.com> Date: Sat, 24 Oct 2026 14:09:14 +0100 Subject: [PATCH 2/2] Place interpolation gradient regressions in existing test module --- dm_pix/_src/interpolation_gradient_test.py | 89 ---------------------- dm_pix/_src/interpolation_test.py | 65 ++++++++++++++++ 2 files changed, 65 insertions(+), 89 deletions(-) delete mode 100644 dm_pix/_src/interpolation_gradient_test.py diff --git a/dm_pix/_src/interpolation_gradient_test.py b/dm_pix/_src/interpolation_gradient_test.py deleted file mode 100644 index a340e13..0000000 --- a/dm_pix/_src/interpolation_gradient_test.py +++ /dev/null @@ -1,89 +0,0 @@ -# Copyright 2021 DeepMind Technologies Limited. All Rights Reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Coordinate gradients of linear interpolation at exact grid points.""" - -from absl.testing import absltest -from absl.testing import parameterized -import dm_pix -import jax -import jax.numpy as jnp -import numpy as np - - -class InterpolationGradientTest(parameterized.TestCase): - - @parameterized.product( - ndim=(1, 2, 3), flattened=(False, True), constant=(False, True) - ) - def test_affine_ramp_gradient(self, ndim, flattened, constant): - shape = (5,) * ndim - slopes = jnp.arange(1, ndim + 1, dtype=jnp.float32) / 32 - grid = jnp.indices(shape, dtype=jnp.float32) - volume = jnp.einsum("i,i...->...", slopes, grid) - coordinates = jnp.full((ndim,), 2.0, dtype=jnp.float32) - fn = ( - dm_pix.flat_nd_linear_interpolate_constant - if constant - else dm_pix.flat_nd_linear_interpolate - ) - kwargs = {"unflattened_vol_shape": shape} if flattened else {} - values = volume.reshape(-1) if flattened else volume - - def interpolate(x): - return fn(values, x, **kwargs) - - expected = jnp.dot(slopes, coordinates) - np.testing.assert_allclose(interpolate(coordinates), expected, atol=1e-7) - for grad in (jax.grad(interpolate), jax.jit(jax.grad(interpolate))): - np.testing.assert_allclose(grad(coordinates), slopes, atol=1e-7) - tangent = jnp.ones(ndim) - _, derivative = jax.jvp(interpolate, (coordinates,), (tangent,)) - np.testing.assert_allclose(derivative, jnp.sum(slopes), atol=1e-7) - - @parameterized.parameters("nearest", "constant") - def test_identity_affine_translation_has_gradient(self, mode): - rows, columns = jnp.indices((5, 5), dtype=jnp.float32) - image = (rows / 16 + columns / 8)[..., None] - - def sample(offset): - return dm_pix.affine_transform( - image, jnp.eye(3), offset=offset, mode=mode - )[2, 2, 0] - - offset = jnp.zeros(3) - for grad in (jax.grad(sample), jax.jit(jax.grad(sample))): - np.testing.assert_allclose(grad(offset), [1 / 16, 1 / 8, 0], atol=1e-7) - - def test_clipped_coordinates_and_singleton_axes(self): - volume = jnp.arange(5, dtype=jnp.float32)[None, :] / 4 - - def fn(x): - return dm_pix.flat_nd_linear_interpolate(volume, x) - - for coordinate, expected in (([-1.0, -1.0], 0.0), ([1.0, 6.0], 1.0)): - coordinate = jnp.asarray(coordinate) - np.testing.assert_allclose(fn(coordinate), expected) - np.testing.assert_array_equal(jax.grad(fn)(coordinate), jnp.zeros(2)) - - def test_integer_samples_preserve_value_gradients(self): - volume = jnp.arange(5, dtype=jnp.float32) / 4 - - def fn(values): - return dm_pix.flat_nd_linear_interpolate(values, jnp.array([2.0])) - - np.testing.assert_array_equal(jax.grad(fn)(volume), [0, 0, 1, 0, 0]) - - -if __name__ == "__main__": - absltest.main() diff --git a/dm_pix/_src/interpolation_test.py b/dm_pix/_src/interpolation_test.py index 4343cda..8150369 100644 --- a/dm_pix/_src/interpolation_test.py +++ b/dm_pix/_src/interpolation_test.py @@ -18,7 +18,9 @@ from absl.testing import absltest from absl.testing import parameterized import chex +import dm_pix from dm_pix._src import interpolation +import jax import jax.numpy as jnp import numpy as np @@ -188,6 +190,69 @@ def test_flat_nd_linear_interpolate_constant(self, shape_coordinates, cval): unflattened_vol_shape=volume.shape), expected) +class InterpolationGradientTest(parameterized.TestCase): + @parameterized.product( + ndim=(1, 2, 3), flattened=(False, True), constant=(False, True) + ) + def test_affine_ramp_gradient(self, ndim, flattened, constant): + shape = (5,) * ndim + slopes = jnp.arange(1, ndim + 1, dtype=jnp.float32) / 32 + grid = jnp.indices(shape, dtype=jnp.float32) + volume = jnp.einsum("i,i...->...", slopes, grid) + coordinates = jnp.full((ndim,), 2.0, dtype=jnp.float32) + fn = ( + dm_pix.flat_nd_linear_interpolate_constant + if constant + else dm_pix.flat_nd_linear_interpolate + ) + kwargs = {"unflattened_vol_shape": shape} if flattened else {} + values = volume.reshape(-1) if flattened else volume + + def interpolate(x): + return fn(values, x, **kwargs) + + expected = jnp.dot(slopes, coordinates) + np.testing.assert_allclose(interpolate(coordinates), expected, atol=1e-7) + for grad in (jax.grad(interpolate), jax.jit(jax.grad(interpolate))): + np.testing.assert_allclose(grad(coordinates), slopes, atol=1e-7) + tangent = jnp.ones(ndim) + _, derivative = jax.jvp(interpolate, (coordinates,), (tangent,)) + np.testing.assert_allclose(derivative, jnp.sum(slopes), atol=1e-7) + + @parameterized.parameters("nearest", "constant") + def test_identity_affine_translation_has_gradient(self, mode): + rows, columns = jnp.indices((5, 5), dtype=jnp.float32) + image = (rows / 16 + columns / 8)[..., None] + + def sample(offset): + return dm_pix.affine_transform( + image, jnp.eye(3), offset=offset, mode=mode + )[2, 2, 0] + + offset = jnp.zeros(3) + for grad in (jax.grad(sample), jax.jit(jax.grad(sample))): + np.testing.assert_allclose(grad(offset), [1 / 16, 1 / 8, 0], atol=1e-7) + + def test_clipped_coordinates_and_singleton_axes(self): + volume = jnp.arange(5, dtype=jnp.float32)[None, :] / 4 + + def fn(x): + return dm_pix.flat_nd_linear_interpolate(volume, x) + + for coordinate, expected in (([-1.0, -1.0], 0.0), ([1.0, 6.0], 1.0)): + coordinate = jnp.asarray(coordinate) + np.testing.assert_allclose(fn(coordinate), expected) + np.testing.assert_array_equal(jax.grad(fn)(coordinate), jnp.zeros(2)) + + def test_integer_samples_preserve_value_gradients(self): + volume = jnp.arange(5, dtype=jnp.float32) / 4 + + def fn(values): + return dm_pix.flat_nd_linear_interpolate(values, jnp.array([2.0])) + + np.testing.assert_array_equal(jax.grad(fn)(volume), [0, 0, 1, 0, 0]) + + if __name__ == "__main__": os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" absltest.main()