<h3 dir="auto">Description</h3>
<p dir="auto">Code to reproduce and compare to scipy below.<br>
I think there might be a bug map_coordinates when using the default arguments of mode='constant' and cval=0.0.<br>
The function seems to be doing some sort of extrapolation in beyond one of the corners of the image even though it should just return the constant 0.0. Whats more, this seems to happen only on one corner.</p>
<p dir="auto">Related to <a href="https://github.com/google/jax/issues/5687" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/5687/hovercard">#5687</a>, but with different behavior in different corners.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '' # Doesnt impact the bug

import jax
import jax.numpy as jnp
import scipy as sp

kernel = jnp.ones((8, 8), dtype=jnp.float32)
in_pos = jnp.array([[0., 0.], [4., 4.], [7., 7.], [7.01, 7.01], [7.5, 7.5], [-7.5, 7.5], [7.5, -7.5], [-7.5, -7.5]])
out = jax.scipy.ndimage.map_coordinates(kernel, in_pos.T, order=1, mode='constant', cval=0.0)
print(out)  # [1.         1.         1.         0.98009956 0.25       0.        0.         0.        ]

out2 = sp.ndimage.map_coordinates(kernel, in_pos.T, order=1)
print(out2)  # [1. 1. 1. 0. 0. 0. 0. 0.]
"><pre class="notranslate"><code class="notranslate">import os
# os.environ['CUDA_VISIBLE_DEVICES'] = '' # Doesnt impact the bug

import jax
import jax.numpy as jnp
import scipy as sp

kernel = jnp.ones((8, 8), dtype=jnp.float32)
in_pos = jnp.array([[0., 0.], [4., 4.], [7., 7.], [7.01, 7.01], [7.5, 7.5], [-7.5, 7.5], [7.5, -7.5], [-7.5, -7.5]])
out = jax.scipy.ndimage.map_coordinates(kernel, in_pos.T, order=1, mode='constant', cval=0.0)
print(out)  # [1.         1.         1.         0.98009956 0.25       0.        0.         0.        ]

out2 = sp.ndimage.map_coordinates(kernel, in_pos.T, order=1)
print(out2)  # [1. 1. 1. 0. 0. 0. 0. 0.]

</code></pre></div>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">0.4.1</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU and GPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Ubuntu 22.05</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto"><em>No response</em></p>
<p dir="auto">jax.scipy.ndimage.map_coordinates behaves differently than scipy.ndimage.map_coordinates when using mode="constant". In this example we can see that 2.9 which should be as "off the array" as 3, produces a different output which is in fact wrong (1) but should in fact produce what the original implementation produces (0, which is cval).</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from scipy import ndimage
import jax

import jax.numpy as jnp

src = jnp.array([5, 1, 6])  # Fails!
# src = jnp.array([5, 1, 5]) # Doesn't fail

# Since 2.9 is first jnp.floored, this gives &quot;2&quot;. Which is a valid index.
coords = [[2, 3, 2.9]]


def scipy_map_coordinates():
  return ndimage.map_coordinates(src, coords, order=1, mode=&quot;constant&quot;)


def jax_map_coordinates():
  return jax.scipy.ndimage.map_coordinates(
      src, coords, order=1, mode=&quot;constant&quot;)


print(scipy_map_coordinates())
print(jax_map_coordinates())

assert jnp.array_equal(scipy_map_coordinates(), jax_map_coordinates())"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">scipy</span> <span class="pl-k">import</span> <span class="pl-s1">ndimage</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>

<span class="pl-s1">src</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">5</span>, <span class="pl-c1">1</span>, <span class="pl-c1">6</span>])  <span class="pl-c"># Fails!</span>
<span class="pl-c"># src = jnp.array([5, 1, 5]) # Doesn't fail</span>

<span class="pl-c"># Since 2.9 is first jnp.floored, this gives "2". Which is a valid index.</span>
<span class="pl-s1">coords</span> <span class="pl-c1">=</span> [[<span class="pl-c1">2</span>, <span class="pl-c1">3</span>, <span class="pl-c1">2.9</span>]]


<span class="pl-k">def</span> <span class="pl-en">scipy_map_coordinates</span>():
  <span class="pl-k">return</span> <span class="pl-s1">ndimage</span>.<span class="pl-en">map_coordinates</span>(<span class="pl-s1">src</span>, <span class="pl-s1">coords</span>, <span class="pl-s1">order</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">mode</span><span class="pl-c1">=</span><span class="pl-s">"constant"</span>)


<span class="pl-k">def</span> <span class="pl-en">jax_map_coordinates</span>():
  <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">scipy</span>.<span class="pl-s1">ndimage</span>.<span class="pl-en">map_coordinates</span>(
      <span class="pl-s1">src</span>, <span class="pl-s1">coords</span>, <span class="pl-s1">order</span><span class="pl-c1">=</span><span class="pl-c1">1</span>, <span class="pl-s1">mode</span><span class="pl-c1">=</span><span class="pl-s">"constant"</span>)


<span class="pl-en">print</span>(<span class="pl-en">scipy_map_coordinates</span>())
<span class="pl-en">print</span>(<span class="pl-en">jax_map_coordinates</span>())

<span class="pl-k">assert</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array_equal</span>(<span class="pl-en">scipy_map_coordinates</span>(), <span class="pl-en">jax_map_coordinates</span>())</pre></div>
<p dir="auto">Yielding:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[6 0 0]
[6 0 1]
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
&lt;ipython-input-6-b13f0479754d&gt; in &lt;module&gt;()
     20 print(jax_map_coordinates())
     21 
---&gt; 22 assert jnp.array_equal(scipy_map_coordinates(), jax_map_coordinates())

AssertionError: "><pre class="notranslate"><code class="notranslate">[6 0 0]
[6 0 1]
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
&lt;ipython-input-6-b13f0479754d&gt; in &lt;module&gt;()
     20 print(jax_map_coordinates())
     21 
---&gt; 22 assert jnp.array_equal(scipy_map_coordinates(), jax_map_coordinates())

AssertionError: 
</code></pre></div>
<p dir="auto">CC: <a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/claudiofantacci/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/claudiofantacci">@claudiofantacci</a></p>