Repository: google/jax

Issue:
jax.scipy.ndimage.map_coordinates round half is different than scipy
`jax.scipy.ndimage.map_coordinates` implements 'round half to even' when using nearest-neighbor interpolation and the coordinat values are exactly half (i.e. both 1.5 and 2.5 get rounded to 2). This is the default behavior for jax and ordinary numpy rounding (round, around, rint). However scipy.ndimage.map_coordinates appears to implement 'round half up'.

The difference between jax and scipy is illustrated with the following commands.
`>>> x = np.arange(10).astype(np.float32)`
`>>> coords = [[3.5, 4.5, 5.5, 6.5, 7.5]]`
`>>> jax.scipy.ndimage.map_coordinates(x, coords, order=0)`
`DeviceArray([4., 4., 6., 6., 8.], dtype=float32)`
`>>> scipy.ndimage.map_coordinates(x, coords, order=0)`
`array([4., 5., 6., 7., 8.], dtype=float32)`

Not sure which behavior is desired, but if the default numpy behavior is desired, the jax.scipy documentation should mention this. If the scipy behavior is desired, the rounding function at 
[https://github.com/google/jax/blob/7d157c71815e5386f2b449bac69ef9e3ab7e701f/jax/scipy/ndimage.py#L40](https://github.com/google/jax/blob/7d157c71815e5386f2b449bac69ef9e3ab7e701f/jax/scipy/ndimage.py#L40)
should be changed to something like:
`index = jnp.floor(coordinate + .5).astype(jnp.int32)`


Discussion:

Resolving patch:
diff --git a/jax/scipy/ndimage.py b/jax/scipy/ndimage.py
--- a/jax/scipy/ndimage.py
+++ b/jax/scipy/ndimage.py
@@ -21,6 +21,7 @@
 import scipy.ndimage
 
 from .. import api
+from .. import lax
 from ..numpy import lax_numpy as jnp
 from ..numpy._util import _wraps
 from ..util import safe_zip as zip
@@ -36,8 +37,12 @@
 }
 
 
+def _round_half_away_from_zero(a):
+  return a if jnp.issubdtype(a.dtype, jnp.integer) else lax.round(a)
+
+
 def _nearest_indices_and_weights(coordinate):
-  index = jnp.around(coordinate).astype(jnp.int32)
+  index = _round_half_away_from_zero(coordinate).astype(jnp.int32)
   weight = coordinate.dtype.type(1)
   return [(index, weight)]
 
@@ -53,7 +58,7 @@ def _linear_indices_and_weights(coordinate):
 @functools.partial(api.jit, static_argnums=(2, 3, 4))
 def _map_coordinates(input, coordinates, order, mode, cval):
   input = jnp.asarray(input)
-  coordinates = [jnp.asarray(c, input.dtype) for c in coordinates]
+  coordinates = [jnp.asarray(c) for c in coordinates]
   cval = jnp.asarray(cval, input.dtype)
 
   if len(coordinates) != input.ndim:
@@ -100,7 +105,9 @@ def _map_coordinates(input, coordinates, order, mode, cval):
       contribution = jnp.where(all_valid, input[indices], cval)
     outputs.append(_nonempty_prod(weights) * contribution)
   result = _nonempty_sum(outputs)
-  return result
+  if jnp.issubdtype(input.dtype, jnp.integer):
+    result = _round_half_away_from_zero(result)
+  return result.astype(input.dtype)
 
 
 @_wraps(scipy.ndimage.map_coordinates, lax_description=textwrap.dedent("""\


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
@@ -74,7 +74,7 @@ class NdimageTest(jtu.JaxTestCase):
        "cval": cval, "impl": impl, "round_": round_}
       for shape in [(5,), (3, 4), (3, 4, 5)]
       for coords_shape in [(7,), (2, 3, 4)]
-      for dtype in float_dtypes
+      for dtype in float_dtypes + int_dtypes
       for coords_dtype in float_dtypes
       for order in [0, 1]
       for mode in ['wrap', 'constant', 'nearest']
@@ -100,10 +100,14 @@ def args_maker():
     impl_fun = (osp_ndimage.map_coordinates if impl == "original"
                 else _fixed_ref_map_coordinates)
     osp_op = lambda x, c: impl_fun(x, c, order=order, mode=mode, cval=cval)
-    epsilon = max([dtypes.finfo(dtypes.canonicalize_dtype(d)).eps
-                   for d in [dtype, coords_dtype]])
-    self._CheckAgainstNumpy(lsp_op, osp_op, args_maker, tol=100*epsilon,
-                            check_dtypes=True)
+    if dtype in float_dtypes:
+      epsilon = max([dtypes.finfo(dtypes.canonicalize_dtype(d)).eps
+                     for d in [dtype, coords_dtype]])
+      self._CheckAgainstNumpy(lsp_op, osp_op, args_maker, tol=100*epsilon,
+                              check_dtypes=True)
+    else:
+      self._CheckAgainstNumpy(lsp_op, osp_op, args_maker, tol=0,
+                              check_dtypes=True)
 
   def testMapCoordinatesErrors(self):
     x = onp.arange(5.0)
@@ -120,6 +124,21 @@ def testMapCoordinateDocstring(self):
     self.assertIn("Only linear interpolation",
                   lsp_ndimage.map_coordinates.__doc__)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_{}_order={}".format(onp.dtype(dtype), order),
+       "dtype": dtype, "order": order}
+      for dtype in float_dtypes + int_dtypes
+      for order in [0, 1]))
+  def testMapCoordinatesRoundHalf(self, dtype, order):
+    x = onp.arange(-3, 3, dtype=dtype)
+    c = onp.array([[.5, 1.5, 2.5, 3.5]])
+    def args_maker():
+      return x, c
+
+    lsp_op = lambda x, c: lsp_ndimage.map_coordinates(x, c, order=order)
+    osp_op = lambda x, c: osp_ndimage.map_coordinates(x, c, order=order)
+    self._CheckAgainstNumpy(lsp_op, osp_op, args_maker, check_dtypes=True)
+
   def testContinuousGradients(self):
     # regression test for https://github.com/google/jax/issues/3024