<p dir="auto">When using <code class="notranslate">jax.numpy.histogram2d</code> together with the <code class="notranslate">range</code> keyword argument, an error occurs.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp
import numpy as np

samples = jax.random.uniform(jax.random.PRNGKey(123), shape=(2, 30))
lims = [[-.5, .5], [-.5, .5]]

# numpy version works fine
hist = np.histogram2d(samples[0, :], samples[1, :], range=lims)
# jnp version fail with error below
hist = jnp.histogram2d(samples[0, :], samples[1, :], range=lims)"><pre class="notranslate"><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-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>

<span class="pl-s1">samples</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">123</span>), <span class="pl-s1">shape</span><span class="pl-c1">=</span>(<span class="pl-c1">2</span>, <span class="pl-c1">30</span>))
<span class="pl-s1">lims</span> <span class="pl-c1">=</span> [[<span class="pl-c1">-</span><span class="pl-c1">.5</span>, <span class="pl-c1">.5</span>], [<span class="pl-c1">-</span><span class="pl-c1">.5</span>, <span class="pl-c1">.5</span>]]

<span class="pl-c"># numpy version works fine</span>
<span class="pl-s1">hist</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">histogram2d</span>(<span class="pl-s1">samples</span>[<span class="pl-c1">0</span>, :], <span class="pl-s1">samples</span>[<span class="pl-c1">1</span>, :], <span class="pl-s1">range</span><span class="pl-c1">=</span><span class="pl-s1">lims</span>)
<span class="pl-c"># jnp version fail with error below</span>
<span class="pl-s1">hist</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">histogram2d</span>(<span class="pl-s1">samples</span>[<span class="pl-c1">0</span>, :], <span class="pl-s1">samples</span>[<span class="pl-c1">1</span>, :], <span class="pl-s1">range</span><span class="pl-c1">=</span><span class="pl-s1">lims</span>)</pre></div>
<p dir="auto">Error:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Traceback (most recent call last):
  File &quot;/home/jnphist2d.py&quot;, line 11, in &lt;module&gt;
    hist = jnp.histogram2d(samples[0, :], samples[1, :], range=lims)
  File &quot;/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py&quot;, line 1005, in histogram2d
    hist, edges = histogramdd(sample, bins, range, weights, density)
  File &quot;/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py&quot;, line 1031, in histogramdd
    bin_idx = searchsorted(bin_edges, sample[:, i], side='right')
  File &quot;/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py&quot;, line 5717, in searchsorted
    raise ValueError(&quot;a should be 1-dimensional&quot;)
ValueError: a should be 1-dimensional"><pre class="notranslate"><code class="notranslate">WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Traceback (most recent call last):
  File "/home/jnphist2d.py", line 11, in &lt;module&gt;
    hist = jnp.histogram2d(samples[0, :], samples[1, :], range=lims)
  File "/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py", line 1005, in histogram2d
    hist, edges = histogramdd(sample, bins, range, weights, density)
  File "/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py", line 1031, in histogramdd
    bin_idx = searchsorted(bin_edges, sample[:, i], side='right')
  File "/home/usr/anaconda3/lib/python3.7/site-packages/jax/_src/numpy/lax_numpy.py", line 5717, in searchsorted
    raise ValueError("a should be 1-dimensional")
ValueError: a should be 1-dimensional
</code></pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax import numpy as jnp
a = [-2,2]
b=jnp.array([a,a])
print(b)
jnp.histogram2d(jnp.ones(10), jnp.ones(10), bins=10, range=b)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">jnp</span>
<span class="pl-s1">a</span> <span class="pl-c1">=</span> [<span class="pl-c1">-</span><span class="pl-c1">2</span>,<span class="pl-c1">2</span>]
<span class="pl-s1">b</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-s1">a</span>,<span class="pl-s1">a</span>])
<span class="pl-en">print</span>(<span class="pl-s1">b</span>)
<span class="pl-s1">jnp</span>.<span class="pl-en">histogram2d</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-c1">10</span>), <span class="pl-s1">jnp</span>.<span class="pl-en">ones</span>(<span class="pl-c1">10</span>), <span class="pl-s1">bins</span><span class="pl-c1">=</span><span class="pl-c1">10</span>, <span class="pl-s1">range</span><span class="pl-c1">=</span><span class="pl-s1">b</span>)</pre></div>
<p dir="auto">Suggested fix from the user: <a href="https://github.com/google/jax/blob/7a40aa0114e46a7a8975e04b175dfbc7b0b388ef/jax/_src/numpy/lax_numpy.py#L1054">this line</a> should probably have <code class="notranslate">range[i]</code>. (That seems to fix it, but I didn't try writing tests!)</p>
<p dir="auto"><a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/users/jakevdp/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/jakevdp">@jakevdp</a> could you take a look?</p>