We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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]
```

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `jax/scipy/ndimage.py`
Content:
```
1 # Copyright 2019 Google LLC
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     https://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 
16 import functools
17 import itertools
18 import operator
19 import textwrap
20 
21 import scipy.ndimage
22 
23 from .. import api
24 from ..numpy import lax_numpy as jnp
25 from ..numpy._util import _wraps
26 from ..util import safe_zip as zip
27 
28 
29 _nonempty_prod = functools.partial(functools.reduce, operator.mul)
30 _nonempty_sum = functools.partial(functools.reduce, operator.add)
31 
32 _INDEX_FIXERS = {
33     'constant': lambda index, size: index,
34     'nearest': lambda index, size: jnp.clip(index, 0, size - 1),
35     'wrap': lambda index, size: index % size,
36 }
37 
38 
39 def _nearest_indices_and_weights(coordinate):
40   index = jnp.around(coordinate).astype(jnp.int32)
41   weight = coordinate.dtype.type(1)
42   return [(index, weight)]
43 
44 
45 def _linear_indices_and_weights(coordinate):
46   lower = jnp.floor(coordinate)
47   upper = jnp.ceil(coordinate)
48   l_index = lower.astype(jnp.int32)
49   u_index = upper.astype(jnp.int32)
50   one = coordinate.dtype.type(1)
51   l_weight = one - (coordinate - lower)
52   u_weight = one - l_weight  # handles the edge case lower==upper
53   return [(l_index, l_weight), (u_index, u_weight)]
54 
55 
56 @functools.partial(api.jit, static_argnums=(2, 3, 4))
57 def _map_coordinates(input, coordinates, order, mode, cval):
58   input = jnp.asarray(input)
59   coordinates = [jnp.asarray(c, input.dtype) for c in coordinates]
60   cval = jnp.asarray(cval, input.dtype)
61 
62   if len(coordinates) != input.ndim:
63     raise ValueError('coordinates must be a sequence of length input.ndim, but '
64                      '{} != {}'.format(len(coordinates), input.ndim))
65 
66   index_fixer = _INDEX_FIXERS.get(mode)
67   if index_fixer is None:
68     raise NotImplementedError(
69         'jax.scipy.ndimage.map_coordinates does not yet support mode {}. '
70         'Currently supported modes are {}.'.format(mode, set(_INDEX_FIXERS)))
71 
72   if mode == 'constant':
73     is_valid = lambda index, size: (0 <= index) & (index < size)
74   else:
75     is_valid = lambda index, size: True
76 
77   if order == 0:
78     interp_fun = _nearest_indices_and_weights
79   elif order == 1:
80     interp_fun = _linear_indices_and_weights
81   else:
82     raise NotImplementedError(
83         'jax.scipy.ndimage.map_coordinates currently requires order<=1')
84 
85   valid_1d_interpolations = []
86   for coordinate, size in zip(coordinates, input.shape):
87     interp_nodes = interp_fun(coordinate)
88     valid_interp = []
89     for index, weight in interp_nodes:
90       fixed_index = index_fixer(index, size)
91       valid = is_valid(index, size)
92       valid_interp.append((fixed_index, valid, weight))
93     valid_1d_interpolations.append(valid_interp)
94 
95   outputs = []
96   for items in itertools.product(*valid_1d_interpolations):
97     indices, validities, weights = zip(*items)
98     if any(valid is not True for valid in validities):
99       all_valid = functools.reduce(operator.and_, validities)
100       contribution = jnp.where(all_valid, input[indices], cval)
101     else:
102       contribution = input[indices]
103     outputs.append(_nonempty_prod(weights) * contribution)
104   result = _nonempty_sum(outputs)
105   return result
106 
107 
108 @_wraps(scipy.ndimage.map_coordinates, lax_description=textwrap.dedent("""\
109     Only linear interpolation (``order=1``) and modes ``'constant'``,
110     ``'nearest'`` and ``'wrap'`` are currently supported. Note that
111     interpolation near boundaries differs from the scipy function, because we
112     fixed an outstanding bug (https://github.com/scipy/scipy/issues/2640);
113     this function interprets the ``mode`` argument as documented by SciPy, but
114     not as implemented by SciPy.
115     """))
116 def map_coordinates(
117     input, coordinates, order, mode='constant', cval=0.0,
118 ):
119   return _map_coordinates(input, coordinates, order, mode, cval)
120
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 _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 @@
   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

{"golden_diff": "diff --git a/jax/scipy/ndimage.py b/jax/scipy/ndimage.py\n--- a/jax/scipy/ndimage.py\n+++ b/jax/scipy/ndimage.py\n@@ -44,13 +44,10 @@\n \n def _linear_indices_and_weights(coordinate):\n   lower = jnp.floor(coordinate)\n-  upper = jnp.ceil(coordinate)\n-  l_index = lower.astype(jnp.int32)\n-  u_index = upper.astype(jnp.int32)\n-  one = coordinate.dtype.type(1)\n-  l_weight = one - (coordinate - lower)\n-  u_weight = one - l_weight  # handles the edge case lower==upper\n-  return [(l_index, l_weight), (u_index, u_weight)]\n+  upper_weight = coordinate - lower\n+  lower_weight = 1 - upper_weight\n+  index = lower.astype(jnp.int32)\n+  return [(index, lower_weight), (index + 1, upper_weight)]\n \n \n @functools.partial(api.jit, static_argnums=(2, 3, 4))\n@@ -95,11 +92,12 @@\n   outputs = []\n   for items in itertools.product(*valid_1d_interpolations):\n     indices, validities, weights = zip(*items)\n-    if any(valid is not True for valid in validities):\n+    if all(valid is True for valid in validities):\n+      # fast path\n+      contribution = input[indices]\n+    else:\n       all_valid = functools.reduce(operator.and_, validities)\n       contribution = jnp.where(all_valid, input[indices], cval)\n-    else:\n-      contribution = input[indices]\n     outputs.append(_nonempty_prod(weights) * contribution)\n   result = _nonempty_sum(outputs)\n   return result\n", "issue": "Zero gradient when resampling an image at grid location using map_coordinates\nSorry if this not a very minimal test case but let me explain my use case and issue as I believe the context will help.\r\n\r\nI 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`.\r\n\r\nThe 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)\u2207x2` and computing `\u2207x2` with something similar to `np.gradient`.\r\n\r\nTrying 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.\r\n\r\nBelow 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**.\r\n\r\nI 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.\r\n\r\n```python\r\nimport jax\r\nimport jax.numpy as jnp\r\nimport numpy as onp\r\nfrom jax.scipy import ndimage as jndimage\r\nfrom scipy import ndimage as ondimage\r\nimport scipy as oscipy\r\n\r\n# This needs to run at startup\r\n# https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#Double-(64bit)-precision\r\njax.config.update('jax_enable_x64', True)\r\n\r\n# Exclude a border in hte computation of the loss to try and avoid numerical issues\r\nexcl_border = 2\r\n\r\ndef run(np,ndimage):\r\n  print(f\"\\nRunning on {np}\")\r\n  onp.random.seed(0)\r\n  x1 = onp.random.randn(20,10)\r\n  x2 = onp.random.randn(20,10)\r\n\r\n  grid_x, grid_y = np.meshgrid(np.arange(x1.shape[1]), np.arange(x1.shape[0]))\r\n\r\n  def lossfunc(du):\r\n    # Get translated grid\r\n    def_grid_x = grid_x + du[0]\r\n    def_grid_y = grid_y + du[1]\r\n\r\n    # Resample image\r\n    tmpx2_warped = jndimage.map_coordinates(x2, [def_grid_y, def_grid_x], order=1)\r\n\r\n    # Compute the MSE between the warped image an the fixed image\r\n    diff_im = tmpx2_warped[excl_border:-excl_border,excl_border:-excl_border]-x1[excl_border:-excl_border,excl_border:-excl_border]\r\n\r\n    imloss = np.mean((diff_im)**2)\r\n\r\n    return imloss\r\n\r\n  def mg_lossfunc(du):\r\n    #Get translated grid\r\n    def_grid_x = grid_x + du[0]\r\n    def_grid_y = grid_y + du[1]\r\n\r\n    # Resample image\r\n    tmpx2_warped = ndimage.map_coordinates(x2, [def_grid_y, def_grid_x], order=1)\r\n\r\n    # Compute the MSE between the warped image an the fixed image\r\n    diff_im = tmpx2_warped[excl_border:-excl_border,excl_border:-excl_border]-x1[excl_border:-excl_border,excl_border:-excl_border]\r\n    jm = np.gradient(tmpx2_warped)\r\n    jmrx = jm[0][excl_border:-excl_border,excl_border:-excl_border]\r\n    jmry = jm[1][excl_border:-excl_border,excl_border:-excl_border]\r\n    jdmx = -2.*diff_im*jmrx\r\n    jdmy = -2.*diff_im*jmry\r\n\r\n    return np.array([np.mean(jdmx), np.mean(jdmy)])\r\n\r\n  print(f\"loss at 0,0: {lossfunc([0., 0.])}\")\r\n  print(f\"loss at 0.1,0.1: {lossfunc([0.1, 0.1])}\")\r\n\r\n  print(f\"Manual approx gradient at 0,0: {mg_lossfunc([0., 0.])}\")\r\n  print(f\"Manual approx gradient at 0.1,0.1: {mg_lossfunc([0.1, 0.1])}\")\r\n  # Finite difference gradient with a large step as the image is sampled on a grid and interpolated\r\n  epsgrad = 0.1\r\n  print(f\"Numerical approx gradient at 0,0: {oscipy.optimize.approx_fprime([0., 0.],lossfunc,epsgrad)}\")\r\n  print(f\"Numerical approx gradient at 0.1,0.1: {oscipy.optimize.approx_fprime([0.1, 0.1],lossfunc,epsgrad)}\")\r\n  if np==jnp:\r\n    jg_lossfunc = lambda du:np.asarray(jax.jit(jax.grad(lossfunc))(du))\r\n    print(f\"Jax gradient at 0,0: {jg_lossfunc([0., 0.])}\")\r\n    print(f\"Jax gradient at 0.1,0.1: {jg_lossfunc([0.1, 0.1])}\")\r\n\r\nrun(onp,ondimage)\r\nrun(jnp,jndimage)\r\n```\r\nOutputs:\r\n```\r\nRunning on <module 'numpy' from '/usr/local/lib/python3.6/dist-packages/numpy/__init__.py'>\r\nloss at 0,0: 1.9514397057693333\r\nloss at 0.1,0.1: 1.6584230856711766\r\nManual approx gradient at 0,0: [ 0.01189463 -0.21033731]\r\nManual approx gradient at 0.1,0.1: [-0.00626396 -0.12798883]\r\nNumerical approx gradient at 0,0: [-1.72714353 -1.47912221]\r\nNumerical approx gradient at 0.1,0.1: [-1.15234194 -0.9063109 ]\r\n\r\nRunning on <module 'jax.numpy' from '/usr/local/lib/python3.6/dist-packages/jax/numpy/__init__.py'>\r\nloss at 0,0: 1.9514397057693333\r\nloss at 0.1,0.1: 1.6584230856711766\r\nManual approx gradient at 0,0: [ 0.01189463 -0.21033731]\r\nManual approx gradient at 0.1,0.1: [-0.00626396 -0.12798883]\r\nNumerical approx gradient at 0,0: [-1.72714353 -1.47912221]\r\nNumerical approx gradient at 0.1,0.1: [-1.15234194 -0.9063109 ]\r\nJax gradient at 0,0: [0. 0.]\r\nJax gradient at 0.1,0.1: [-1.30169297 -1.05466678]\r\n```\n", "before_files": [{"content": "# Copyright 2019 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport functools\nimport itertools\nimport operator\nimport textwrap\n\nimport scipy.ndimage\n\nfrom .. import api\nfrom ..numpy import lax_numpy as jnp\nfrom ..numpy._util import _wraps\nfrom ..util import safe_zip as zip\n\n\n_nonempty_prod = functools.partial(functools.reduce, operator.mul)\n_nonempty_sum = functools.partial(functools.reduce, operator.add)\n\n_INDEX_FIXERS = {\n    'constant': lambda index, size: index,\n    'nearest': lambda index, size: jnp.clip(index, 0, size - 1),\n    'wrap': lambda index, size: index % size,\n}\n\n\ndef _nearest_indices_and_weights(coordinate):\n  index = jnp.around(coordinate).astype(jnp.int32)\n  weight = coordinate.dtype.type(1)\n  return [(index, weight)]\n\n\ndef _linear_indices_and_weights(coordinate):\n  lower = jnp.floor(coordinate)\n  upper = jnp.ceil(coordinate)\n  l_index = lower.astype(jnp.int32)\n  u_index = upper.astype(jnp.int32)\n  one = coordinate.dtype.type(1)\n  l_weight = one - (coordinate - lower)\n  u_weight = one - l_weight  # handles the edge case lower==upper\n  return [(l_index, l_weight), (u_index, u_weight)]\n\n\n@functools.partial(api.jit, static_argnums=(2, 3, 4))\ndef _map_coordinates(input, coordinates, order, mode, cval):\n  input = jnp.asarray(input)\n  coordinates = [jnp.asarray(c, input.dtype) for c in coordinates]\n  cval = jnp.asarray(cval, input.dtype)\n\n  if len(coordinates) != input.ndim:\n    raise ValueError('coordinates must be a sequence of length input.ndim, but '\n                     '{} != {}'.format(len(coordinates), input.ndim))\n\n  index_fixer = _INDEX_FIXERS.get(mode)\n  if index_fixer is None:\n    raise NotImplementedError(\n        'jax.scipy.ndimage.map_coordinates does not yet support mode {}. '\n        'Currently supported modes are {}.'.format(mode, set(_INDEX_FIXERS)))\n\n  if mode == 'constant':\n    is_valid = lambda index, size: (0 <= index) & (index < size)\n  else:\n    is_valid = lambda index, size: True\n\n  if order == 0:\n    interp_fun = _nearest_indices_and_weights\n  elif order == 1:\n    interp_fun = _linear_indices_and_weights\n  else:\n    raise NotImplementedError(\n        'jax.scipy.ndimage.map_coordinates currently requires order<=1')\n\n  valid_1d_interpolations = []\n  for coordinate, size in zip(coordinates, input.shape):\n    interp_nodes = interp_fun(coordinate)\n    valid_interp = []\n    for index, weight in interp_nodes:\n      fixed_index = index_fixer(index, size)\n      valid = is_valid(index, size)\n      valid_interp.append((fixed_index, valid, weight))\n    valid_1d_interpolations.append(valid_interp)\n\n  outputs = []\n  for items in itertools.product(*valid_1d_interpolations):\n    indices, validities, weights = zip(*items)\n    if any(valid is not True for valid in validities):\n      all_valid = functools.reduce(operator.and_, validities)\n      contribution = jnp.where(all_valid, input[indices], cval)\n    else:\n      contribution = input[indices]\n    outputs.append(_nonempty_prod(weights) * contribution)\n  result = _nonempty_sum(outputs)\n  return result\n\n\n@_wraps(scipy.ndimage.map_coordinates, lax_description=textwrap.dedent(\"\"\"\\\n    Only linear interpolation (``order=1``) and modes ``'constant'``,\n    ``'nearest'`` and ``'wrap'`` are currently supported. Note that\n    interpolation near boundaries differs from the scipy function, because we\n    fixed an outstanding bug (https://github.com/scipy/scipy/issues/2640);\n    this function interprets the ``mode`` argument as documented by SciPy, but\n    not as implemented by SciPy.\n    \"\"\"))\ndef map_coordinates(\n    input, coordinates, order, mode='constant', cval=0.0,\n):\n  return _map_coordinates(input, coordinates, order, mode, cval)\n", "path": "jax/scipy/ndimage.py"}], "after_files": [{"content": "# Copyright 2019 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport functools\nimport itertools\nimport operator\nimport textwrap\n\nimport scipy.ndimage\n\nfrom .. import api\nfrom ..numpy import lax_numpy as jnp\nfrom ..numpy._util import _wraps\nfrom ..util import safe_zip as zip\n\n\n_nonempty_prod = functools.partial(functools.reduce, operator.mul)\n_nonempty_sum = functools.partial(functools.reduce, operator.add)\n\n_INDEX_FIXERS = {\n    'constant': lambda index, size: index,\n    'nearest': lambda index, size: jnp.clip(index, 0, size - 1),\n    'wrap': lambda index, size: index % size,\n}\n\n\ndef _nearest_indices_and_weights(coordinate):\n  index = jnp.around(coordinate).astype(jnp.int32)\n  weight = coordinate.dtype.type(1)\n  return [(index, weight)]\n\n\ndef _linear_indices_and_weights(coordinate):\n  lower = jnp.floor(coordinate)\n  upper_weight = coordinate - lower\n  lower_weight = 1 - upper_weight\n  index = lower.astype(jnp.int32)\n  return [(index, lower_weight), (index + 1, upper_weight)]\n\n\n@functools.partial(api.jit, static_argnums=(2, 3, 4))\ndef _map_coordinates(input, coordinates, order, mode, cval):\n  input = jnp.asarray(input)\n  coordinates = [jnp.asarray(c, input.dtype) for c in coordinates]\n  cval = jnp.asarray(cval, input.dtype)\n\n  if len(coordinates) != input.ndim:\n    raise ValueError('coordinates must be a sequence of length input.ndim, but '\n                     '{} != {}'.format(len(coordinates), input.ndim))\n\n  index_fixer = _INDEX_FIXERS.get(mode)\n  if index_fixer is None:\n    raise NotImplementedError(\n        'jax.scipy.ndimage.map_coordinates does not yet support mode {}. '\n        'Currently supported modes are {}.'.format(mode, set(_INDEX_FIXERS)))\n\n  if mode == 'constant':\n    is_valid = lambda index, size: (0 <= index) & (index < size)\n  else:\n    is_valid = lambda index, size: True\n\n  if order == 0:\n    interp_fun = _nearest_indices_and_weights\n  elif order == 1:\n    interp_fun = _linear_indices_and_weights\n  else:\n    raise NotImplementedError(\n        'jax.scipy.ndimage.map_coordinates currently requires order<=1')\n\n  valid_1d_interpolations = []\n  for coordinate, size in zip(coordinates, input.shape):\n    interp_nodes = interp_fun(coordinate)\n    valid_interp = []\n    for index, weight in interp_nodes:\n      fixed_index = index_fixer(index, size)\n      valid = is_valid(index, size)\n      valid_interp.append((fixed_index, valid, weight))\n    valid_1d_interpolations.append(valid_interp)\n\n  outputs = []\n  for items in itertools.product(*valid_1d_interpolations):\n    indices, validities, weights = zip(*items)\n    if all(valid is True for valid in validities):\n      # fast path\n      contribution = input[indices]\n    else:\n      all_valid = functools.reduce(operator.and_, validities)\n      contribution = jnp.where(all_valid, input[indices], cval)\n    outputs.append(_nonempty_prod(weights) * contribution)\n  result = _nonempty_sum(outputs)\n  return result\n\n\n@_wraps(scipy.ndimage.map_coordinates, lax_description=textwrap.dedent(\"\"\"\\\n    Only linear interpolation (``order=1``) and modes ``'constant'``,\n    ``'nearest'`` and ``'wrap'`` are currently supported. Note that\n    interpolation near boundaries differs from the scipy function, because we\n    fixed an outstanding bug (https://github.com/scipy/scipy/issues/2640);\n    this function interprets the ``mode`` argument as documented by SciPy, but\n    not as implemented by SciPy.\n    \"\"\"))\ndef map_coordinates(\n    input, coordinates, order, mode='constant', cval=0.0,\n):\n  return _map_coordinates(input, coordinates, order, mode, cval)\n", "path": "jax/scipy/ndimage.py"}]}