Preserve coordinate gradients at integer interpolation points

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.

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()