<h3 dir="auto">Description</h3>
<p dir="auto">Hello! I recently was struggling to find a bug in my code, when I realized the problem came from some weird behavior from Jax. Below, note the inconsistency in evaluation when mutliplying <code class="notranslate">jnp.array(False)</code> with arrays of length &gt;1:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp

f = jnp.array(False)
t = jnp.array(True)

inf=jnp.inf
a = jnp.array([inf])
b = jnp.array([inf, inf]) 

# Scalar multiplication (what I used as a baseline for 'normal')
print( t * inf )  # -&gt; Array(inf, dtype=float32)
print( f * inf )  # -&gt; Array(0., dtype=float32)

# Array multiplication
print( t * a )  # normal behavior -&gt; Array([inf], dtype=float32)
print( f * a )  # normal behavior -&gt; Array([0.], dtype=float32)
print( t * b ) # normal behavior -&gt; Array([inf, inf], dtype=float32)
print( f * b ) # unexpected behavior! -&gt; Array([nan, nan], dtype=float32)"><pre class="notranslate"><code class="notranslate">import jax.numpy as jnp

f = jnp.array(False)
t = jnp.array(True)

inf=jnp.inf
a = jnp.array([inf])
b = jnp.array([inf, inf]) 

# Scalar multiplication (what I used as a baseline for 'normal')
print( t * inf )  # -&gt; Array(inf, dtype=float32)
print( f * inf )  # -&gt; Array(0., dtype=float32)

# Array multiplication
print( t * a )  # normal behavior -&gt; Array([inf], dtype=float32)
print( f * a )  # normal behavior -&gt; Array([0.], dtype=float32)
print( t * b ) # normal behavior -&gt; Array([inf, inf], dtype=float32)
print( f * b ) # unexpected behavior! -&gt; Array([nan, nan], dtype=float32)
</code></pre></div>
<p dir="auto">It seems like there is a low probability that this behavior is intentional, so I decided to open this issue. Thanks in advance for any help!</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax 0.4.8, jaxlib 0.4.7</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">GPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">Python 3.9.12, using Ubuntu on WSL2</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.65       Driver Version: 527.37       CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0 Off |                  N/A |
| N/A   53C    P8     2W /  N/A |   3850MiB /  4096MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     11014      C   /python3.9                      N/A      |
+-----------------------------------------------------------------------------+"><pre class="notranslate"><code class="notranslate">+-----------------------------------------------------------------------------+
| NVIDIA-SMI 525.65       Driver Version: 527.37       CUDA Version: 12.0     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  NVIDIA GeForce ...  On   | 00000000:01:00.0 Off |                  N/A |
| N/A   53C    P8     2W /  N/A |   3850MiB /  4096MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     11014      C   /python3.9                      N/A      |
+-----------------------------------------------------------------------------+
</code></pre></div>
<h3 dir="auto">Description</h3>
<p dir="auto">Multiplying a NaN by False in jax gives a 0. See example below:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="nan = jnp.array(float('nan'), dtype=jnp.float32)
print(nan)
print(nan * jnp.array(0., dtype=jnp.float32))
print(nan * jnp.array(False, dtype=jnp.bool_))"><pre class="notranslate"><span class="pl-s1">nan</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-en">float</span>(<span class="pl-s">'nan'</span>), <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">float32</span>)
<span class="pl-en">print</span>(<span class="pl-s1">nan</span>)
<span class="pl-en">print</span>(<span class="pl-s1">nan</span> <span class="pl-c1">*</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-c1">0.</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">float32</span>))
<span class="pl-en">print</span>(<span class="pl-s1">nan</span> <span class="pl-c1">*</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-c1">False</span>, <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">jnp</span>.<span class="pl-s1">bool_</span>))</pre></div>
<p dir="auto">Output:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="nan
nan
0.0"><pre class="notranslate"><code class="notranslate">nan
nan
0.0
</code></pre></div>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto"><em>No response</em></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">Linux</p>