<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<p dir="auto">In essence, when an input is passed into a fan out, and then a fan in concat, regardless of n on the FanOut, element wise operations do not work.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.example_libraries.stax as stax
from jax import random
import jax.numpy as jnp


input_rows = 10
input_columns = 5
key = random.PRNGKey(42)
input = random.normal(key, (input_rows, input_columns))

out_stack = 5 # bugged with any value
input_comparison = jnp.column_stack([input]*out_stack)

s = stax.serial(
    stax.FanOut(out_stack),
    stax.FanInConcat(),
    stax.Selu, 
    # stax.Elu,
    # stax.Exp,
)
params, shape = s[0](key, input.shape)
assert jnp.isclose(s[1](params, input), input_comparison).all()"><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">example_libraries</span>.<span class="pl-s1">stax</span> <span class="pl-k">as</span> <span class="pl-s1">stax</span>
<span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">random</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">input_rows</span> <span class="pl-c1">=</span> <span class="pl-c1">10</span>
<span class="pl-s1">input_columns</span> <span class="pl-c1">=</span> <span class="pl-c1">5</span>
<span class="pl-s1">key</span> <span class="pl-c1">=</span> <span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">42</span>)
<span class="pl-s1">input</span> <span class="pl-c1">=</span> <span class="pl-s1">random</span>.<span class="pl-en">normal</span>(<span class="pl-s1">key</span>, (<span class="pl-s1">input_rows</span>, <span class="pl-s1">input_columns</span>))

<span class="pl-s1">out_stack</span> <span class="pl-c1">=</span> <span class="pl-c1">5</span> <span class="pl-c"># bugged with any value</span>
<span class="pl-s1">input_comparison</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">column_stack</span>([<span class="pl-s1">input</span>]<span class="pl-c1">*</span><span class="pl-s1">out_stack</span>)

<span class="pl-s1">s</span> <span class="pl-c1">=</span> <span class="pl-s1">stax</span>.<span class="pl-en">serial</span>(
    <span class="pl-s1">stax</span>.<span class="pl-v">FanOut</span>(<span class="pl-s1">out_stack</span>),
    <span class="pl-s1">stax</span>.<span class="pl-v">FanInConcat</span>(),
    <span class="pl-s1">stax</span>.<span class="pl-v">Selu</span>, 
    <span class="pl-c"># stax.Elu,</span>
    <span class="pl-c"># stax.Exp,</span>
)
<span class="pl-s1">params</span>, <span class="pl-s1">shape</span> <span class="pl-c1">=</span> <span class="pl-s1">s</span>[<span class="pl-c1">0</span>](<span class="pl-s1">key</span>, <span class="pl-s1">input</span>.<span class="pl-s1">shape</span>)
<span class="pl-k">assert</span> <span class="pl-s1">jnp</span>.<span class="pl-en">isclose</span>(<span class="pl-s1">s</span>[<span class="pl-c1">1</span>](<span class="pl-s1">params</span>, <span class="pl-s1">input</span>), <span class="pl-s1">input_comparison</span>).<span class="pl-en">all</span>()</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.<br>
N/A</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Python 3.9.7 (default, Sep 16 2021, 08:50:36)
[Clang 10.0.0 ] :: Anaconda, Inc. on darwin
Type &quot;help&quot;, &quot;copyright&quot;, &quot;credits&quot; or &quot;license&quot; for more information.
&gt;&gt;&gt; import jax
&gt;&gt;&gt; jax.__version__
'0.2.25'"><pre class="notranslate"><code class="notranslate">Python 3.9.7 (default, Sep 16 2021, 08:50:36)
[Clang 10.0.0 ] :: Anaconda, Inc. on darwin
Type "help", "copyright", "credits" or "license" for more information.
&gt;&gt;&gt; import jax
&gt;&gt;&gt; jax.__version__
'0.2.25'
</code></pre></div>
<p dir="auto">Please:</p>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Check for duplicate issues.</li>
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<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
def apply(mask, scores):
  attention_mask = -1. * 1e30 * (1.0 - mask)
  scores += attention_mask
  unnormalized = scores - scores.max()
  return unnormalized 
mask = np.array(1.0)
scores = jax.numpy.array([0.7, 0.5])
apply(mask, scores), jax.jit(apply)(mask, scores)"><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-k">def</span> <span class="pl-en">apply</span>(<span class="pl-s1">mask</span>, <span class="pl-s1">scores</span>):
  <span class="pl-s1">attention_mask</span> <span class="pl-c1">=</span> <span class="pl-c1">-</span><span class="pl-c1">1.</span> <span class="pl-c1">*</span> <span class="pl-c1">1e30</span> <span class="pl-c1">*</span> (<span class="pl-c1">1.0</span> <span class="pl-c1">-</span> <span class="pl-s1">mask</span>)
  <span class="pl-s1">scores</span> <span class="pl-c1">+=</span> <span class="pl-s1">attention_mask</span>
  <span class="pl-s1">unnormalized</span> <span class="pl-c1">=</span> <span class="pl-s1">scores</span> <span class="pl-c1">-</span> <span class="pl-s1">scores</span>.<span class="pl-en">max</span>()
  <span class="pl-k">return</span> <span class="pl-s1">unnormalized</span> 
<span class="pl-s1">mask</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-c1">1.0</span>)
<span class="pl-s1">scores</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-en">array</span>([<span class="pl-c1">0.7</span>, <span class="pl-c1">0.5</span>])
<span class="pl-en">apply</span>(<span class="pl-s1">mask</span>, <span class="pl-s1">scores</span>), <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">apply</span>)(<span class="pl-s1">mask</span>, <span class="pl-s1">scores</span>)</pre></div>
<ul class="contains-task-list">
<li class="task-list-item"><input type="checkbox" id="" disabled="" class="task-list-item-checkbox" checked=""> If applicable, include full error messages/tracebacks.</li>
</ul>
<p dir="auto">The code above produces:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="(DeviceArray([ 0.        , -0.19999999], dtype=float32),
 DeviceArray([0., 0.], dtype=float32))"><pre class="notranslate"><code class="notranslate">(DeviceArray([ 0.        , -0.19999999], dtype=float32),
 DeviceArray([0., 0.], dtype=float32))
</code></pre></div>
<p dir="auto">on a CPU machine but correctly produces:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="(DeviceArray([ 0.        , -0.19999999], dtype=float32),
 DeviceArray([ 0.        , -0.19999999], dtype=float32))"><pre class="notranslate"><code class="notranslate">(DeviceArray([ 0.        , -0.19999999], dtype=float32),
 DeviceArray([ 0.        , -0.19999999], dtype=float32))
</code></pre></div>
<p dir="auto">on GPU and TPU<br>
The issue seems to be around <code class="notranslate">-0.0</code> and adding it to the score.<br>
If the key part of the code snippet that adds the <code class="notranslate">-0.0</code> is rephrased as the equivalent:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content=" attention_mask = 1. * 1e30 * (1.0 - mask)
 scores -= attention_mask"><pre class="notranslate"><code class="notranslate"> attention_mask = 1. * 1e30 * (1.0 - mask)
 scores -= attention_mask
</code></pre></div>
<p dir="auto">Then the CPU produces the same results as GPU and TPU</p>