<h3 dir="auto">Description</h3>
<p dir="auto">When integer inputs are passed in to <code class="notranslate">jax.numpy.divmod</code>, a division/modulo by zero results in inconsistent results when compared with <code class="notranslate">numpy.divmod</code> or the Python operators <code class="notranslate">// and %</code> as shown below:</p>
<p dir="auto"><strong><code class="notranslate">jax.numpy.divmod</code></strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
x1 = jnp.array(3, dtype=&quot;int32&quot;)
x2 = jnp.array(0, dtype=&quot;int32&quot;)
print(x1)
print(x2)
print(jnp.divmod(x1, x2))"><pre class="notranslate"><code class="notranslate">import jax.numpy as jnp
x1 = jnp.array(3, dtype="int32")
x2 = jnp.array(0, dtype="int32")
print(x1)
print(x2)
print(jnp.divmod(x1, x2))
</code></pre></div>
<p dir="auto"><em>Captured Output:</em></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="3
0
(DeviceArray(-2, dtype=int32), DeviceArray(3, dtype=int32))"><pre class="notranslate"><code class="notranslate">3
0
(DeviceArray(-2, dtype=int32), DeviceArray(3, dtype=int32))
</code></pre></div>
<p dir="auto">whereas with <code class="notranslate">numpy.divmod</code>, the output comes out to be:</p>
<p dir="auto"><strong><code class="notranslate">numpy.divmod</code></strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import numpy as np
x1 = np.array(3, dtype=&quot;int32&quot;)
x2 = np.array(0, dtype=&quot;int32&quot;)
print(x1)
print(x2)
print(np.divmod(x1, x2))"><pre class="notranslate"><code class="notranslate">import numpy as np
x1 = np.array(3, dtype="int32")
x2 = np.array(0, dtype="int32")
print(x1)
print(x2)
print(np.divmod(x1, x2))
</code></pre></div>
<p dir="auto"><em>Captured Output:</em></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="3
0
(0, 0)"><pre class="notranslate"><code class="notranslate">3
0
(0, 0)
</code></pre></div>
<p dir="auto">and with Python's divmod operators:</p>
<p dir="auto"><strong><code class="notranslate">//</code> and <code class="notranslate">%</code></strong></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="x1 = np.array(3, dtype=&quot;int32&quot;)
x2 = np.array(0, dtype=&quot;int32&quot;)
print(x1)
print(x2)
print((x1//x2, x1%x2))"><pre class="notranslate"><code class="notranslate">x1 = np.array(3, dtype="int32")
x2 = np.array(0, dtype="int32")
print(x1)
print(x2)
print((x1//x2, x1%x2))
</code></pre></div>
<p dir="auto"><em>Captured Output:</em></p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="3
0
(0, 0)"><pre class="notranslate"><code class="notranslate">3
0
(0, 0)
</code></pre></div>
<p dir="auto">This behaviour does not follow numpy's or python's functional equivalents, even though the same has been claimed in the documentation. This behaviour is also reproducible when the above code is ran with <code class="notranslate">jax.numpy.floor_divide</code>, <code class="notranslate">jax.numpy.divide</code> or <code class="notranslate">jax.numpy.mod</code> methods as those methods probably delegate to the same API in the backend implementation.</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax v0.3.14</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Python (v3.8.10), OS (Linux WSL2 Ubuntu 20.04 LTS)</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto"><em>No response</em></p>
<h3 dir="auto">Description</h3>
<p dir="auto">Jax and numpy differ in their behavior with regards to computing a modulus by zero. Jax treats it as a no-op, and numpy treats it as the zero function.</p>
<h2 dir="auto">Numpy</h2>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="Python 3.10.5 (main, Jul 15 2022, 03:56:49) [GCC 11.3.0] on linux
&gt;&gt;&gt; import numpy as np
&gt;&gt;&gt; a = np.array([1,2,3], dtype=np.uint32)
&gt;&gt;&gt; np.mod(a, 0)
&lt;stdin&gt;:1: RuntimeWarning: divide by zero encountered in remainder
array([0, 0, 0], dtype=uint32)"><pre class="notranslate"><span class="pl-v">Python</span> <span class="pl-c1">3.10</span><span class="pl-c1">.5</span> (<span class="pl-s1">main</span>, <span class="pl-v">Jul</span> <span class="pl-c1">15</span> <span class="pl-c1">2022</span>, <span class="pl-c1">03</span>:<span class="pl-c1">56</span>:<span class="pl-c1">49</span>) [<span class="pl-v">GCC</span> <span class="pl-c1">11.3</span><span class="pl-c1">.0</span>] <span class="pl-s1">on</span> <span class="pl-s1">linux</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">a</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>,<span class="pl-c1">2</span>,<span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">np</span>.<span class="pl-s1">uint32</span>)
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">np</span>.<span class="pl-en">mod</span>(<span class="pl-s1">a</span>, <span class="pl-c1">0</span>)
<span class="pl-c1">&lt;</span><span class="pl-s1">stdin</span><span class="pl-c1">&gt;</span>:<span class="pl-c1">1</span>: <span class="pl-v">RuntimeWarning</span>: <span class="pl-s1">divide</span> <span class="pl-s1">by</span> <span class="pl-s1">zero</span> <span class="pl-s1">encountered</span> <span class="pl-c1">in</span> <span class="pl-s1">remainder</span>
<span class="pl-en">array</span>([<span class="pl-c1">0</span>, <span class="pl-c1">0</span>, <span class="pl-c1">0</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">uint32</span>)</pre></div>
<h2 dir="auto">JAX</h2>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import jax.numpy as jnp
&gt;&gt;&gt; a = jnp.array([1,2,3], dtype=jnp.uint32)
&gt;&gt;&gt; jnp.mod(a, 0)
DeviceArray([1, 2, 3], dtype=uint32)"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</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-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">a</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1</span>,<span class="pl-c1">2</span>,<span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">uint32</span>)
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">jnp</span>.<span class="pl-en">mod</span>(<span class="pl-s1">a</span>, <span class="pl-c1">0</span>)
<span class="pl-v">DeviceArray</span>([<span class="pl-c1">1</span>, <span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">uint32</span>)</pre></div>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">HEAD</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU</p>
<h3 dir="auto">Additional System Info</h3>
<p dir="auto">Python 3.9, 3.10.5</p>