<p dir="auto">There's inconsistent behavior when using jax.numpy arrays vs numpy arrays when plotting histograms in matplotlib. I'm guessing this is a result of some differing property of the two objects thats not immediately obvious to me. They have the same shape. Its easy to fix by just converting to a regular numpy array (as below), but I thought I would report in case this is indicative of something that might cause problems in other situations</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import matplotlib.pyplot as plt
import jax.numpy as np
import numpy as onp

a = onp.random.rand(100)

plt.figure()
_ = plt.hist(np.array(a))

plt.figure()
_ = plt.hist(onp.array(a))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">matplotlib</span>.<span class="pl-s1">pyplot</span> <span class="pl-k">as</span> <span class="pl-s1">plt</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">np</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">onp</span>

<span class="pl-s1">a</span> <span class="pl-c1">=</span> <span class="pl-s1">onp</span>.<span class="pl-s1">random</span>.<span class="pl-en">rand</span>(<span class="pl-c1">100</span>)

<span class="pl-s1">plt</span>.<span class="pl-en">figure</span>()
<span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">plt</span>.<span class="pl-en">hist</span>(<span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">a</span>))

<span class="pl-s1">plt</span>.<span class="pl-en">figure</span>()
<span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">plt</span>.<span class="pl-en">hist</span>(<span class="pl-s1">onp</span>.<span class="pl-en">array</span>(<span class="pl-s1">a</span>))</pre></div>
<p dir="auto">gives:<br>
<a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/7969470/146586234-29ea6f14-11d0-4085-9582-695391900812.png"><img src="https://user-images.githubusercontent.com/7969470/146586234-29ea6f14-11d0-4085-9582-695391900812.png" alt="Screen Shot 2021-12-17 at 9 43 40 AM" style="max-width: 100%;"></a></p>
<p dir="auto">For some usages of JAX, the elements of an array seem to be transposed relative to NumPy leading to unexpected result. This is for example the case for matplotlib histograms.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

fig, axs = plt.subplots(1, 3)
axs.flat[0].hist(np.linspace(0., 1.))
axs.flat[0].set_title(&quot;Expected Result (NumPy)&quot;)
axs.flat[1].hist(jnp.linspace(0., 1.))
axs.flat[1].set_title(&quot;JAX Result&quot;)
axs.flat[2].hist(np.linspace(0., 1.).reshape(1, -1))
axs.flat[2].set_title(&quot;JAX Result w/ NumPy&quot;)
fig.tight_layout()
plt.show()"><pre class="notranslate"><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">matplotlib</span>.<span class="pl-s1">pyplot</span> <span class="pl-k">as</span> <span class="pl-s1">plt</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">fig</span>, <span class="pl-s1">axs</span> <span class="pl-c1">=</span> <span class="pl-s1">plt</span>.<span class="pl-en">subplots</span>(<span class="pl-c1">1</span>, <span class="pl-c1">3</span>)
<span class="pl-s1">axs</span>.<span class="pl-s1">flat</span>[<span class="pl-c1">0</span>].<span class="pl-en">hist</span>(<span class="pl-s1">np</span>.<span class="pl-en">linspace</span>(<span class="pl-c1">0.</span>, <span class="pl-c1">1.</span>))
<span class="pl-s1">axs</span>.<span class="pl-s1">flat</span>[<span class="pl-c1">0</span>].<span class="pl-en">set_title</span>(<span class="pl-s">"Expected Result (NumPy)"</span>)
<span class="pl-s1">axs</span>.<span class="pl-s1">flat</span>[<span class="pl-c1">1</span>].<span class="pl-en">hist</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">linspace</span>(<span class="pl-c1">0.</span>, <span class="pl-c1">1.</span>))
<span class="pl-s1">axs</span>.<span class="pl-s1">flat</span>[<span class="pl-c1">1</span>].<span class="pl-en">set_title</span>(<span class="pl-s">"JAX Result"</span>)
<span class="pl-s1">axs</span>.<span class="pl-s1">flat</span>[<span class="pl-c1">2</span>].<span class="pl-en">hist</span>(<span class="pl-s1">np</span>.<span class="pl-en">linspace</span>(<span class="pl-c1">0.</span>, <span class="pl-c1">1.</span>).<span class="pl-en">reshape</span>(<span class="pl-c1">1</span>, <span class="pl-c1">-</span><span class="pl-c1">1</span>))
<span class="pl-s1">axs</span>.<span class="pl-s1">flat</span>[<span class="pl-c1">2</span>].<span class="pl-en">set_title</span>(<span class="pl-s">"JAX Result w/ NumPy"</span>)
<span class="pl-s1">fig</span>.<span class="pl-en">tight_layout</span>()
<span class="pl-s1">plt</span>.<span class="pl-en">show</span>()</pre></div>
<p dir="auto"><a target="_blank" rel="noopener noreferrer nofollow" href="https://user-images.githubusercontent.com/7541458/145623461-026faef5-8e48-49da-a00f-cbb27c10284f.png"><img src="https://user-images.githubusercontent.com/7541458/145623461-026faef5-8e48-49da-a00f-cbb27c10284f.png" alt="JAX_transpose_bug" style="max-width: 100%;"></a></p>