Repository: google/jax

Issue:
Zero gradient when resampling an image at grid location using map_coordinates
Sorry if this not a very minimal test case but let me explain my use case and issue as I believe the context will help.

I am trying to use jax for a toy image registration problem. Given two images `x1` and `x2` I want to find the translation `u` that minimises the difference between `x1(.)` and `x2(.+u)` as measured in terms of mean square error (MSE). The resampled version of `x2` after the (non-integer) translation is computed with `map_coordinates`.

The computation of the gradient of the cost funcrtion in this context is usually done by assuming the images are continuous, computing the gradient of the MSE as `-2(x1-x2)∇x2` and computing `∇x2` with something similar to `np.gradient`.

Trying to mimick this setup with jax to avoid computing the gradient manually (this would be useful for example as soon as one wants to change the MSE loss for something else) fails to converge to a suitable translation (at least if initialised with an integer translation) as the gradient from jax is **exactly** zero for integer translations.

Below is a test case to illustrate the zero gradient issue. I understand there is a discontinuity of the gradient at such integer points but I was expecting to nonetheless get a proper **sub-gradient**.

I am not sure if I am doing something wrong or simply misunderstanding something but so far I haven't managed to get the image registration to converge with jax derivatives even though it does with a numerical approximation of the gradient, or the classical continous approximation I was refering to.

```python
import jax
import jax.numpy as jnp
import numpy as onp
from jax.scipy import ndimage as jndimage
from scipy import ndimage as ondimage
import scipy as oscipy

# This needs to run at startup
# https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#Double-(64bit)-precision
jax.config.update('jax_enable_x64', True)

# Exclude a border in hte computation of the loss to try and avoid numerical issues
excl_border = 2

def run(np,ndimage):
  print(f"\nRunning on {np}")
  onp.random.seed(0)
  x1 = onp.random.randn(20,10)
  x2 = onp.random.randn(20,10)

  grid_x, grid_y = np.meshgrid(np.arange(x1.shape[1]), np.arange(x1.shape[0]))

  def lossfunc(du):
    # Get translated grid
    def_grid_x = grid_x + du[0]
    def_grid_y = grid_y + du[1]

    # Resample image
    tmpx2_warped = jndimage.map_coordinates(x2, [def_grid_y, def_grid_x], order=1)

    # Compute the MSE between the warped image an the fixed image
    diff_im = tmpx2_warped[excl_border:-excl_border,excl_border:-excl_border]-x1[excl_border:-excl_border,excl_border:-excl_border]

    imloss = np.mean((diff_im)**2)

    return imloss

  def mg_lossfunc(du):
    #Get translated grid
    def_grid_x = grid_x + du[0]
    def_grid_y = grid_y + du[1]

    # Resample image
    tmpx2_warped = ndimage.map_coordinates(x2, [def_grid_y, def_grid_x], order=1)

    # Compute the MSE between the warped image an the fixed image
    diff_im = tmpx2_warped[excl_border:-excl_border,excl_border:-excl_border]-x1[excl_border:-excl_border,excl_border:-excl_border]
    jm = np.gradient(tmpx2_warped)
    jmrx = jm[0][excl_border:-excl_border,excl_border:-excl_border]
    jmry = jm[1][excl_border:-excl_border,excl_border:-excl_border]
    jdmx = -2.*diff_im*jmrx
    jdmy = -2.*diff_im*jmry

    return np.array([np.mean(jdmx), np.mean(jdmy)])

  print(f"loss at 0,0: {lossfunc([0., 0.])}")
  print(f"loss at 0.1,0.1: {lossfunc([0.1, 0.1])}")

  print(f"Manual approx gradient at 0,0: {mg_lossfunc([0., 0.])}")
  print(f"Manual approx gradient at 0.1,0.1: {mg_lossfunc([0.1, 0.1])}")
  # Finite difference gradient with a large step as the image is sampled on a grid and interpolated
  epsgrad = 0.1
  print(f"Numerical approx gradient at 0,0: {oscipy.optimize.approx_fprime([0., 0.],lossfunc,epsgrad)}")
  print(f"Numerical approx gradient at 0.1,0.1: {oscipy.optimize.approx_fprime([0.1, 0.1],lossfunc,epsgrad)}")
  if np==jnp:
    jg_lossfunc = lambda du:np.asarray(jax.jit(jax.grad(lossfunc))(du))
    print(f"Jax gradient at 0,0: {jg_lossfunc([0., 0.])}")
    print(f"Jax gradient at 0.1,0.1: {jg_lossfunc([0.1, 0.1])}")

run(onp,ondimage)
run(jnp,jndimage)
```
Outputs:
```
Running on <module 'numpy' from '/usr/local/lib/python3.6/dist-packages/numpy/__init__.py'>
loss at 0,0: 1.9514397057693333
loss at 0.1,0.1: 1.6584230856711766
Manual approx gradient at 0,0: [ 0.01189463 -0.21033731]
Manual approx gradient at 0.1,0.1: [-0.00626396 -0.12798883]
Numerical approx gradient at 0,0: [-1.72714353 -1.47912221]
Numerical approx gradient at 0.1,0.1: [-1.15234194 -0.9063109 ]

Running on <module 'jax.numpy' from '/usr/local/lib/python3.6/dist-packages/jax/numpy/__init__.py'>
loss at 0,0: 1.9514397057693333
loss at 0.1,0.1: 1.6584230856711766
Manual approx gradient at 0,0: [ 0.01189463 -0.21033731]
Manual approx gradient at 0.1,0.1: [-0.00626396 -0.12798883]
Numerical approx gradient at 0,0: [-1.72714353 -1.47912221]
Numerical approx gradient at 0.1,0.1: [-1.15234194 -0.9063109 ]
Jax gradient at 0,0: [0. 0.]
Jax gradient at 0.1,0.1: [-1.30169297 -1.05466678]
```


Discussion:
I think it may be helpful to look at the function you're optimizing:
```python
import matplotlib.pyplot as plt

plt.figure()
x = jnp.linspace(-1, 2, num=91)
losses = jax.jit(jax.vmap(lossfunc))(jnp.stack([x, x], axis=1))
plt.plot(x, jax.device_get(losses), '-s')
plt.ylabel('loss')

plt.figure()
grads = jax.jit(jax.vmap(jax.grad(lossfunc)))(jnp.stack([x, x], axis=1))
plt.plot(x, jax.device_get(grads), '-s')
plt.ylabel('grad(loss)')
```
![image](https://user-images.githubusercontent.com/1217238/81532180-d2999480-9318-11ea-949e-d99d56e9c475.png)
![image](https://user-images.githubusercontent.com/1217238/81532186-d4fbee80-9318-11ea-86a6-9e15aa72ab45.png)

So yes, it's a little weird that the gradient is _exactly_ zero at these points instead of picking the value from one of the sides (which _might_ make more sense?) but you're going to have trouble optimizing this function with gradient based methods no matter how you calculate them, because the gradients on either side of this point go in opposite directions!

Maybe there's something different about your real use-case here?
Thanks @shoyer the only obvious difference with the real use case is that images are not random. A more complete example that includes the optimisation over the translation can be found here:
https://colab.research.google.com/drive/1lkV7zBPL4YLiKwTxz1uL9K188uiF6BLI?usp=sharing

I guess the classical approximation of the gradient (`-2(x1-x2)∇x2`) has a smoothing effect without which local minima at grid points are an issues. Swiching the interpolation order to 3 instead of 1 might help (at the cost of computational time) but for now this is not implemented in jax:
```
NotImplementedError: jax.scipy.ndimage.map_coordinates currently requires order<=1
```
I don't know the details from the signal processing literature, but I suspect adding some sort of a low-pass filter, either before or after resampling, is important to avoid [anti-aliasing](https://en.wikipedia.org/wiki/Anti-aliasing_filter) issues when doing this sort of alignment.
Many thanks @shoyer for spending time on this. Anti-aliasing is not typically needed in such a simple translation problem. I have expanded a bit the example (a fixed a small bug in the manual approximate gradient) to rely on a very smooth set of images:
![image](https://user-images.githubusercontent.com/1614505/81612933-82343c80-93d5-11ea-9049-38c612ae4007.png)

The following graph shows the gradients along a diagonal translation as computed with the manual approximate gradient (`(mx,my)`), the finite difference one (`(ax,ay)`) and the jax one (`(jx,jy)` with the spikes):
![image](https://user-images.githubusercontent.com/1614505/81613101-c6274180-93d5-11ea-8c0b-792dbef32b60.png)

Code snippet:
<details>
 <summary>Click to expand code</summary>

```python
import jax
import jax.numpy as jnp
import numpy as onp
from jax.scipy import ndimage as jndimage
from scipy import ndimage as ondimage
import scipy as oscipy
import matplotlib.pyplot as plt

# This needs to run at startup
# https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#Double-(64bit)-precision
jax.config.update('jax_enable_x64', True)

# Exclude a border in the computation of the loss to try and avoid numerical issues
excl_border = 2

def run(np,ndimage):
  print(f"\n===\nRunning on {np}")

  h=30
  w=40
  grid_x, grid_y = np.meshgrid(np.arange(w), np.arange(h))

  cx1 = np.round(w/2.)
  cy1 = np.round(h/2.)
  sx = w/4.
  sy = h/4.
  cx2 = cx1+2.
  cy2 = cy1+2.

  x1 = 10*np.exp( -( ((grid_x-cx1)/sx)**2 + ((grid_y-cy1)/sy)**2 ) )
  x2 = 10*np.exp( -( ((grid_x-cx2)/sx)**2 + ((grid_y-cy2)/sy)**2 ) )

  plt.figure()
  fig, axs = plt.subplots(1,2)
  axs[0].set_title('x1')
  axs[0].imshow(x1)

  axs[1].set_title('x2')
  axs[1].imshow(x2)
  plt.show()

  def lossfunc(du):
    # Get translated grid
    def_grid_x = grid_x + du[0]
    def_grid_y = grid_y + du[1]

    # Resample image
    tmpx2_warped = jndimage.map_coordinates(x2, [def_grid_y, def_grid_x], order=1)

    # Compute the MSE between the warped image an the fixed image
    diff_im = tmpx2_warped[excl_border:-excl_border,excl_border:-excl_border]-x1[excl_border:-excl_border,excl_border:-excl_border]

    imloss = np.mean((diff_im)**2)

    return imloss

  def mg_lossfunc(du):
    #Get translated grid
    def_grid_x = grid_x + du[0]
    def_grid_y = grid_y + du[1]

    # Resample image
    tmpx2_warped = ndimage.map_coordinates(x2, [def_grid_y, def_grid_x], order=1)

    # Compute the MSE between the warped image an the fixed image
    diff_im = tmpx2_warped[excl_border:-excl_border,excl_border:-excl_border]-x1[excl_border:-excl_border,excl_border:-excl_border]
    jm = np.gradient(tmpx2_warped)
    jmrx = jm[1][excl_border:-excl_border,excl_border:-excl_border]
    jmry = jm[0][excl_border:-excl_border,excl_border:-excl_border]
    jdmx = 2.*diff_im*jmrx
    jdmy = 2.*diff_im*jmry

    return np.array([np.mean(jdmx), np.mean(jdmy)])

  print(f"loss at 0,0: {lossfunc([0., 0.])}")
  print(f"loss at 0.1,0.1: {lossfunc([0.1, 0.1])}")
  print(f"loss at -0.1,-0.1: {lossfunc([-0.1, -0.1])}")

  plt.figure(figsize=(30, 2))
  uu = np.arange(-1.1, 3.5, 0.05)
  losses = list(map(lossfunc, np.stack([uu, uu], axis=1)))
  plt.plot(uu, losses, '-s')
  plt.ylabel('loss')
  plt.show()

  print(f"Manual approx gradient at 0,0: {mg_lossfunc([0., 0.])}")
  print(f"Manual approx gradient at 0.1,0.1: {mg_lossfunc([0.1, 0.1])}")
  print(f"Manual approx gradient at -0.1,-0.1: {mg_lossfunc([-0.1, -0.1])}")
  print(f"Manual approx gradient at 2,-1: {mg_lossfunc([2., -1.])}")
  print(f"Manual approx gradient at 2.1,-1.1: {mg_lossfunc([2.1, -1.1])}")
  # Finite difference gradient with a large step as the image is sampled on a grid and interpolated
  epsgrad = 0.1
  ag_lossfunc = lambda du: oscipy.optimize.approx_fprime(du,lossfunc,epsgrad)
  print(f"Numerical approx gradient at 0,0: {ag_lossfunc([0., 0.])}")
  print(f"Numerical approx gradient at 0.1,0.1: {ag_lossfunc([0.1, 0.1])}")
  print(f"Numerical approx gradient at -0.1,-0.1: {ag_lossfunc([-0.1, -0.1])}")
  print(f"Numerical approx gradient at 2.,-1.: {ag_lossfunc([2., -1.])}")
  print(f"Numerical approx gradient at 2.1,-1.1: {ag_lossfunc([2.1, -1.1])}")
  if np==jnp:
    jg_lossfunc = lambda du:np.asarray(jax.jit(jax.grad(lossfunc))(du))
    print(f"Jax gradient at 0,0: {jg_lossfunc([0., 0.])}")
    print(f"Jax gradient at 0.1,0.1: {jg_lossfunc([0.1, 0.1])}")
    print(f"Jax gradient at -0.1,-0.1: {jg_lossfunc([-0.1, -0.1])}")
    print(f"Jax gradient at 2.1,-1.1: {jg_lossfunc([2., -1.])}")
    print(f"Jax gradient at 2.1,-1.1: {jg_lossfunc([2.1, -1.1])}")

  plt.figure(figsize=(30, 2))
  mg_grads = list(map(mg_lossfunc, np.stack([uu, uu], axis=1)))
  plt.plot(uu, mg_grads, '-o')
  ag_grads = list(map(ag_lossfunc, np.stack([uu, uu], axis=1)))
  plt.plot(uu, ag_grads, '-+')
  plt.legend(['mx', 'my', 'ax', 'ay'])
  if np==jnp:
    jg_lossfunc = lambda du:np.asarray(jax.jit(jax.grad(lossfunc))(du))
    jg_grads = list(map(jg_lossfunc, np.stack([uu, uu], axis=1)))
    plt.plot(uu, jg_grads, '-d')
    plt.legend(['mx', 'my', 'ax', 'ay', 'jx', 'jy'])
  plt.ylabel('grad(loss)')
  plt.show()

  print("\nBFGS optimisation\n")
  u = np.zeros(2)
  opt_opt={'disp': True, 'maxiter': 200, 'eps': epsgrad, 'gtol':1e-9}
  res = oscipy.optimize.minimize(lossfunc, u, method="BFGS", options=opt_opt)
  print(f"\nNumerical approx gradient - loss at optim end {res.x}: {lossfunc(res.x)}")

  res = oscipy.optimize.minimize(lossfunc, u, jac=mg_lossfunc, method="BFGS", options=opt_opt)
  print(f"\nManual approx gradient - loss at optim end {res.x}: {lossfunc(res.x)}")

  if np==jnp:
    jg_lossfunc = lambda du:np.asarray(jax.jit(jax.grad(lossfunc))(du))
    res = oscipy.optimize.minimize(lossfunc, u, jac=jg_lossfunc, method="BFGS", options=opt_opt)
    print(f"\nJax grad - loss at optim end {res.x}: {lossfunc(res.x)}")

run(onp,ondimage)
run(jnp,jndimage)
```
</details>


<details>
<summary>Click to expand output</summary>

```
===
Running on <module 'numpy' from '/usr/local/lib/python3.6/dist-packages/numpy/__init__.py'>

<Figure size 432x288 with 0 Axes>

loss at 0,0: 1.3439400676144586
loss at 0.1,0.1: 1.214026608789585
loss at -0.1,-0.1: 1.4720960079911962

Manual approx gradient at 0,0: [-0.4687936  -0.82254061]
Manual approx gradient at 0.1,0.1: [-0.44707712 -0.78433099]
Manual approx gradient at -0.1,-0.1: [-0.48807353 -0.85495853]
Manual approx gradient at 2,-1: [-7.75443519e-05 -1.18900303e+00]
Manual approx gradient at 2.1,-1.1: [ 0.02652188 -1.21587652]
Numerical approx gradient at 0,0: [-0.46826943 -0.83002476]
Numerical approx gradient at 0.1,0.1: [-0.44431941 -0.78719149]
Numerical approx gradient at -0.1,-0.1: [-0.46686583 -0.81231689]
Numerical approx gradient at 2.,-1.: [ 0.00328248 -1.22299187]
Numerical approx gradient at 2.1,-1.1: [ 0.02816145 -1.15391493]


BFGS optimisation

Warning: Desired error not necessarily achieved due to precision loss.
         Current function value: 0.000708
         Iterations: 4
         Function evaluations: 324
         Gradient evaluations: 78

Numerical approx gradient - loss at optim end [1.95190614 1.95647443]: 0.0007083630658263352
Warning: Desired error not necessarily achieved due to precision loss.
         Current function value: 0.000000
         Iterations: 8
         Function evaluations: 50
         Gradient evaluations: 39

Manual approx gradient - loss at optim end [1.99890299 1.99994993]: 1.507664203680706e-07

===
Running on <module 'jax.numpy' from '/usr/local/lib/python3.6/dist-packages/jax/numpy/__init__.py'>

<Figure size 432x288 with 0 Axes>

loss at 0,0: 1.3439400676144586
loss at 0.1,0.1: 1.214026608789585
loss at -0.1,-0.1: 1.4720960079911962

Manual approx gradient at 0,0: [-0.4687936  -0.82254061]
Manual approx gradient at 0.1,0.1: [-0.44707712 -0.78433099]
Manual approx gradient at -0.1,-0.1: [-0.48807353 -0.85495853]
Manual approx gradient at 2,-1: [-7.75443519e-05 -1.18900303e+00]
Manual approx gradient at 2.1,-1.1: [ 0.02296511 -1.21586866]
Numerical approx gradient at 0,0: [-0.46826943 -0.83002476]
Numerical approx gradient at 0.1,0.1: [-0.44431941 -0.78719149]
Numerical approx gradient at -0.1,-0.1: [-0.46686583 -0.81231689]
Numerical approx gradient at 2.,-1.: [ 0.00328248 -1.22299187]
Numerical approx gradient at 2.1,-1.1: [ 0.02816145 -1.15391493]
Jax gradient at 0,0: [0. 0.]
Jax gradient at 0.1,0.1: [-0.45671462 -0.80902832]
Jax gradient at -0.1,-0.1: [-0.47918073 -0.83378842]
Jax gradient at 2.1,-1.1: [0. 0.]
Jax gradient at 2.1,-1.1: [ 0.01547956 -1.17478728]


BFGS optimisation

Warning: Desired error not necessarily achieved due to precision loss.
         Current function value: 0.000708
         Iterations: 4
         Function evaluations: 324
         Gradient evaluations: 78

Numerical approx gradient - loss at optim end [1.95190614 1.95647443]: 0.0007083630658263352
Optimization terminated successfully.
         Current function value: 0.000000
         Iterations: 9
         Function evaluations: 10
         Gradient evaluations: 10

Manual approx gradient - loss at optim end [2. 2.]: 5.814722077624332e-20
Optimization terminated successfully.
         Current function value: 1.343940
         Iterations: 0
         Function evaluations: 1
         Gradient evaluations: 1

Jax grad - loss at optim end [0. 0.]: 1.3439400676144586
```
</details>
OK, I agree this looks pretty bad! We shouldn't have the gradient deviate at a single point.
I bet the problem is when `lower == upper` on these lines:
https://github.com/google/jax/blob/db71f3c5fc5226c3e9c87dd9f056d1b63cfa0286/jax/scipy/ndimage.py#L45-L53

Rather than computing `upper = jnp.ceil(coordinate)`, we should probably just set `upper = lower + 1`

Resolving patch:
diff --git a/jax/scipy/ndimage.py b/jax/scipy/ndimage.py
--- a/jax/scipy/ndimage.py
+++ b/jax/scipy/ndimage.py
@@ -44,13 +44,10 @@ def _nearest_indices_and_weights(coordinate):
 
 def _linear_indices_and_weights(coordinate):
   lower = jnp.floor(coordinate)
-  upper = jnp.ceil(coordinate)
-  l_index = lower.astype(jnp.int32)
-  u_index = upper.astype(jnp.int32)
-  one = coordinate.dtype.type(1)
-  l_weight = one - (coordinate - lower)
-  u_weight = one - l_weight  # handles the edge case lower==upper
-  return [(l_index, l_weight), (u_index, u_weight)]
+  upper_weight = coordinate - lower
+  lower_weight = 1 - upper_weight
+  index = lower.astype(jnp.int32)
+  return [(index, lower_weight), (index + 1, upper_weight)]
 
 
 @functools.partial(api.jit, static_argnums=(2, 3, 4))
@@ -95,11 +92,12 @@ def _map_coordinates(input, coordinates, order, mode, cval):
   outputs = []
   for items in itertools.product(*valid_1d_interpolations):
     indices, validities, weights = zip(*items)
-    if any(valid is not True for valid in validities):
+    if all(valid is True for valid in validities):
+      # fast path
+      contribution = input[indices]
+    else:
       all_valid = functools.reduce(operator.and_, validities)
       contribution = jnp.where(all_valid, input[indices], cval)
-    else:
-      contribution = input[indices]
     outputs.append(_nonempty_prod(weights) * contribution)
   result = _nonempty_sum(outputs)
   return result


Test patch:
diff --git a/tests/scipy_ndimage_test.py b/tests/scipy_ndimage_test.py
--- a/tests/scipy_ndimage_test.py
+++ b/tests/scipy_ndimage_test.py
@@ -21,6 +21,7 @@
 from absl.testing import parameterized
 import scipy.ndimage as osp_ndimage
 
+from jax import grad
 from jax import test_util as jtu
 from jax import dtypes
 from jax.scipy import ndimage as lsp_ndimage
@@ -119,6 +120,21 @@ def testMapCoordinateDocstring(self):
     self.assertIn("Only linear interpolation",
                   lsp_ndimage.map_coordinates.__doc__)
 
+  def testContinuousGradients(self):
+    # regression test for https://github.com/google/jax/issues/3024
+
+    def loss(delta):
+      x = onp.arange(100.0)
+      border = 10
+      indices = onp.arange(x.size) + delta
+      # linear interpolation of the linear function y=x should be exact
+      shifted = lsp_ndimage.map_coordinates(x, [indices], order=1)
+      return ((x - shifted) ** 2)[border:-border].mean()
+
+    # analytical gradient of (x - (x - delta)) ** 2 is 2 * delta
+    self.assertAllClose(grad(loss)(0.5), 1.0, check_dtypes=False)
+    self.assertAllClose(grad(loss)(1.0), 2.0, check_dtypes=False)
+
 
 if __name__ == "__main__":
   absltest.main()