<p dir="auto">LayerError                                Traceback (most recent call last)<br>
 in <br>
1 # this takes around 1 minute<br>
2 # print(next(data_generator(Q1_test,Q2_test,512,pad=vocab[''],shuffle=False)))<br>
----&gt; 3 accuracy = classify(Q1_test,Q2_test, y_test, 0.7, model, vocab, batch_size = 512)<br>
4 print("Accuracy", accuracy)</p>
<p dir="auto"> in classify(test_Q1, test_Q2, y, threshold, model, vocab, data_generator, batch_size)<br>
27 #        print(f'shape of y_test = {y_test.shape}\n')<br>
28         # Call the model<br>
---&gt; 29         v1, v2 = model((q1,q2))<br>
30<br>
31         for j in range(batch_size):</p>
<p dir="auto">~/anaconda3/lib/python3.8/site-packages/trax/layers/base.py in <strong>call</strong>(self, x, weights, state, rng)<br>
171       self.state = state  # Needed if the model wasn't fully initialized.<br>
172     state = self.state<br>
--&gt; 173     outputs, new_state = self.pure_fn(x, weights, state, rng)<br>
174     self.state = new_state<br>
175     self.weights = weights</p>
<p dir="auto">~/anaconda3/lib/python3.8/site-packages/trax/layers/base.py in pure_fn(self, x, weights, state, rng, use_cache)<br>
520       # Skipping 3 lines as it's always the uninteresting internal call.<br>
521       name, trace = self._name, _short_traceback(skip=3)<br>
--&gt; 522       raise LayerError(name, 'pure_fn',<br>
523                        self._caller, signature(x), trace) from None<br>
524</p>
<p dir="auto">LayerError: Exception passing through layer Parallel (in pure_fn):<br>
layer created in file [...]/, line 27<br>
layer input shapes: (ShapeDtype{shape:(512, 64), dtype:int64}, ShapeDtype{shape:(512, 64), dtype:int64})</p>
<p dir="auto">File [...]/trax/layers/combinators.py, line 209, in forward<br>
sub_outputs, sub_state = layer.pure_fn(x, w, s, r, use_cache=True)</p>
<p dir="auto">LayerError: Exception passing through layer Serial (in pure_fn):<br>
layer created in file [...]/, line 17<br>
layer input shapes: ShapeDtype{shape:(512, 64), dtype:int64}</p>
<p dir="auto">File [...]/trax/layers/combinators.py, line 88, in forward<br>
outputs, s = layer.pure_fn(inputs, w, s, rng, use_cache=True)</p>
<p dir="auto">LayerError: Exception passing through layer Embedding_41699_128 (in pure_fn):<br>
layer created in file [...]/, line 18<br>
layer input shapes: ShapeDtype{shape:(512, 64), dtype:int64}</p>
<p dir="auto">File [...]/trax/layers/core.py, line 150, in forward<br>
return jnp.take(self.weights, x, axis=0)</p>
<p dir="auto">File [...]/jax/numpy/lax_numpy.py, line 3422, in take<br>
return lax.gather(a, indices[..., None], dimension_numbers=dnums,</p>
<p dir="auto">File [...]/jax/lax/lax.py, line 807, in gather<br>
return gather_p.bind(</p>
<p dir="auto">File [...]/site-packages/jax/core.py, line 276, in bind<br>
return self.impl(*args, **kwargs)</p>
<p dir="auto">File [...]/jax/interpreters/xla.py, line 224, in apply_primitive<br>
compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)</p>
<p dir="auto">File [...]/jax/interpreters/xla.py, line 257, in xla_primitive_callable<br>
built_c = primitive_computation(prim, AxisEnv(nreps, (), (), None), backend,</p>
<p dir="auto">File [...]/jax/interpreters/xla.py, line 316, in primitive_computation<br>
raise RuntimeError(msg) from e</p>
<p dir="auto">RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.:<br>
This is a bug in JAX's shape-checking rules; please report it!</p>
<p dir="auto">code:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="

def gumbel_sample(log_probs, temperature=1.0):
    &quot;&quot;&quot;Gumbel sampling from a categorical distribution.&quot;&quot;&quot;
    u = numpy.random.uniform(low=1e-6, high=1.0 - 1e-6, size=log_probs.shape)
    g = -np.log(-np.log(u))
    return np.argmax(log_probs + g * temperature, axis=-1)

def predict(num_chars, prefix):
    inp = [ord(c) for c in prefix]
    result = [c for c in prefix]
    max_len = len(prefix) + num_chars
    for _ in range(num_chars):
        cur_inp = np.array(inp + [0] * (max_len - len(inp)))
        outp = model(cur_inp[None, :])  # Add batch dim.
        next_char = gumbel_sample(outp[0, len(inp)])
        inp += [int(next_char)]
       
        if inp[-1] == 1:
            break  # EOS
        result.append(chr(int(next_char)))
    
    return &quot;&quot;.join(result)

print(predict(32, &quot;&quot;))"><pre class="notranslate"><span class="pl-k">def</span> <span class="pl-en">gumbel_sample</span>(<span class="pl-s1">log_probs</span>, <span class="pl-s1">temperature</span><span class="pl-c1">=</span><span class="pl-c1">1.0</span>):
    <span class="pl-s">"""Gumbel sampling from a categorical distribution."""</span>
    <span class="pl-s1">u</span> <span class="pl-c1">=</span> <span class="pl-s1">numpy</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">low</span><span class="pl-c1">=</span><span class="pl-c1">1e-6</span>, <span class="pl-s1">high</span><span class="pl-c1">=</span><span class="pl-c1">1.0</span> <span class="pl-c1">-</span> <span class="pl-c1">1e-6</span>, <span class="pl-s1">size</span><span class="pl-c1">=</span><span class="pl-s1">log_probs</span>.<span class="pl-s1">shape</span>)
    <span class="pl-s1">g</span> <span class="pl-c1">=</span> <span class="pl-c1">-</span><span class="pl-s1">np</span>.<span class="pl-en">log</span>(<span class="pl-c1">-</span><span class="pl-s1">np</span>.<span class="pl-en">log</span>(<span class="pl-s1">u</span>))
    <span class="pl-k">return</span> <span class="pl-s1">np</span>.<span class="pl-en">argmax</span>(<span class="pl-s1">log_probs</span> <span class="pl-c1">+</span> <span class="pl-s1">g</span> <span class="pl-c1">*</span> <span class="pl-s1">temperature</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">-</span><span class="pl-c1">1</span>)

<span class="pl-k">def</span> <span class="pl-en">predict</span>(<span class="pl-s1">num_chars</span>, <span class="pl-s1">prefix</span>):
    <span class="pl-s1">inp</span> <span class="pl-c1">=</span> [<span class="pl-en">ord</span>(<span class="pl-s1">c</span>) <span class="pl-k">for</span> <span class="pl-s1">c</span> <span class="pl-c1">in</span> <span class="pl-s1">prefix</span>]
    <span class="pl-s1">result</span> <span class="pl-c1">=</span> [<span class="pl-s1">c</span> <span class="pl-k">for</span> <span class="pl-s1">c</span> <span class="pl-c1">in</span> <span class="pl-s1">prefix</span>]
    <span class="pl-s1">max_len</span> <span class="pl-c1">=</span> <span class="pl-en">len</span>(<span class="pl-s1">prefix</span>) <span class="pl-c1">+</span> <span class="pl-s1">num_chars</span>
    <span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_chars</span>):
        <span class="pl-s1">cur_inp</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">inp</span> <span class="pl-c1">+</span> [<span class="pl-c1">0</span>] <span class="pl-c1">*</span> (<span class="pl-s1">max_len</span> <span class="pl-c1">-</span> <span class="pl-en">len</span>(<span class="pl-s1">inp</span>)))
        <span class="pl-s1">outp</span> <span class="pl-c1">=</span> <span class="pl-en">model</span>(<span class="pl-s1">cur_inp</span>[<span class="pl-c1">None</span>, :])  <span class="pl-c"># Add batch dim.</span>
        <span class="pl-s1">next_char</span> <span class="pl-c1">=</span> <span class="pl-en">gumbel_sample</span>(<span class="pl-s1">outp</span>[<span class="pl-c1">0</span>, <span class="pl-en">len</span>(<span class="pl-s1">inp</span>)])
        <span class="pl-s1">inp</span> <span class="pl-c1">+=</span> [<span class="pl-en">int</span>(<span class="pl-s1">next_char</span>)]
       
        <span class="pl-k">if</span> <span class="pl-s1">inp</span>[<span class="pl-c1">-</span><span class="pl-c1">1</span>] <span class="pl-c1">==</span> <span class="pl-c1">1</span>:
            <span class="pl-k">break</span>  <span class="pl-c"># EOS</span>
        <span class="pl-s1">result</span>.<span class="pl-en">append</span>(<span class="pl-en">chr</span>(<span class="pl-en">int</span>(<span class="pl-s1">next_char</span>)))
    
    <span class="pl-k">return</span> <span class="pl-s">""</span>.<span class="pl-en">join</span>(<span class="pl-s1">result</span>)

<span class="pl-en">print</span>(<span class="pl-en">predict</span>(<span class="pl-c1">32</span>, <span class="pl-s">""</span>))</pre></div>
<p dir="auto">Error:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="LayerError                                Traceback (most recent call last)
&lt;ipython-input-27-6f9f9afc30e6&gt; in &lt;module&gt;
     22     return &quot;&quot;.join(result)
     23 
---&gt; 24 print(predict(32, &quot;&quot;))

&lt;ipython-input-27-6f9f9afc30e6&gt; in predict(num_chars, prefix)
     12     for _ in range(num_chars):
     13         cur_inp = np.array(inp + [0] * (max_len - len(inp)))
---&gt; 14         outp = model(cur_inp[None, :])  # Add batch dim.
     15         next_char = gumbel_sample(outp[0, len(inp)])
     16         inp += [int(next_char)]

/opt/conda/lib/python3.7/site-packages/trax/layers/base.py in __call__(self, x, weights, state, rng)
    165       self.state = state  # Needed if the model wasn't fully initialized.
    166     state = self.state
--&gt; 167     outputs, new_state = self.pure_fn(x, weights, state, rng)
    168     self.state = new_state
    169     self.weights = weights

/opt/conda/lib/python3.7/site-packages/trax/layers/base.py in pure_fn(self, x, weights, state, rng, use_cache)
    448       name, trace = self._name, _short_traceback(skip=3)
    449       raise LayerError(name, 'pure_fn',
--&gt; 450                        self._caller, signature(x), trace) from None
    451 
    452   def output_signature(self, input_signature):

LayerError: Exception passing through layer Serial (in pure_fn):
  layer created in file [...]/&lt;ipython-input-13-6f7ffe0c061e&gt;, line 21
  layer input shapes: ShapeDtype{shape:(1, 32), dtype:int32}

  File [...]/trax/layers/combinators.py, line 88, in forward
    outputs, s = layer.pure_fn(inputs, w, s, rng, use_cache=True)

LayerError: Exception passing through layer Embedding_256_512 (in pure_fn):
  layer created in file [...]/&lt;ipython-input-13-6f7ffe0c061e&gt;, line 18
  layer input shapes: ShapeDtype{shape:(1, 32), dtype:int32}

  File [...]/trax/layers/core.py, line 150, in forward
    return jnp.take(self.weights, x, axis=0)

  File [...]/jax/numpy/lax_numpy.py, line 3298, in take
    slice_sizes=tuple(slice_sizes))

  File [...]/jax/lax/lax.py, line 835, in gather
    slice_sizes=canonicalize_shape(slice_sizes))

  File [...]/site-packages/jax/core.py, line 273, in bind
    return self.impl(*args, **kwargs)

  File [...]/jax/interpreters/xla.py, line 228, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args), **params)

  File [...]/jax/interpreters/xla.py, line 262, in xla_primitive_callable
    *avals, **params)

  File [...]/jax/interpreters/xla.py, line 320, in primitive_computation
    raise RuntimeError(msg) from e

RuntimeError: Invalid argument: Slice size at index 0 in gather op is out of range, must be within [0, 1), got 1.: 
This is a bug in JAX's shape-checking rules; please report it!"><pre class="notranslate"><span class="pl-v">LayerError</span>                                <span class="pl-v">Traceback</span> (<span class="pl-s1">most</span> <span class="pl-s1">recent</span> <span class="pl-s1">call</span> <span class="pl-s1">last</span>)
<span class="pl-c1">&lt;</span><span class="pl-s1">ipython</span><span class="pl-c1">-</span><span class="pl-s1">input</span><span class="pl-c1">-</span><span class="pl-c1">27</span><span class="pl-c1">-</span><span class="pl-c1">6</span><span class="pl-s1">f9f9afc30e6</span><span class="pl-c1">&gt;</span> <span class="pl-s1">in</span> <span class="pl-c1">&lt;</span><span class="pl-s1">module</span><span class="pl-c1">&gt;</span>
     <span class="pl-c1">22</span>     <span class="pl-k">return</span> <span class="pl-s">""</span>.<span class="pl-en">join</span>(<span class="pl-s1">result</span>)
     <span class="pl-c1">23</span> 
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">24</span> <span class="pl-en">print</span>(<span class="pl-en">predict</span>(<span class="pl-c1">32</span>, <span class="pl-s">""</span>))

<span class="pl-c1">&lt;</span><span class="pl-s1">ipython</span><span class="pl-c1">-</span><span class="pl-s1">input</span><span class="pl-c1">-</span><span class="pl-c1">27</span><span class="pl-c1">-</span><span class="pl-c1">6</span><span class="pl-s1">f9f9afc30e6</span><span class="pl-c1">&gt;</span> <span class="pl-s1">in</span> <span class="pl-en">predict</span>(<span class="pl-s1">num_chars</span>, <span class="pl-s1">prefix</span>)
     <span class="pl-c1">12</span>     <span class="pl-k">for</span> <span class="pl-s1">_</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">num_chars</span>):
     <span class="pl-c1">13</span>         <span class="pl-s1">cur_inp</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>(<span class="pl-s1">inp</span> <span class="pl-c1">+</span> [<span class="pl-c1">0</span>] <span class="pl-c1">*</span> (<span class="pl-s1">max_len</span> <span class="pl-c1">-</span> <span class="pl-en">len</span>(<span class="pl-s1">inp</span>)))
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">14</span>         <span class="pl-s1">outp</span> <span class="pl-c1">=</span> <span class="pl-en">model</span>(<span class="pl-s1">cur_inp</span>[<span class="pl-c1">None</span>, :])  <span class="pl-c"># Add batch dim.</span>
     <span class="pl-c1">15</span>         <span class="pl-s1">next_char</span> <span class="pl-c1">=</span> <span class="pl-en">gumbel_sample</span>(<span class="pl-s1">outp</span>[<span class="pl-c1">0</span>, <span class="pl-en">len</span>(<span class="pl-s1">inp</span>)])
     <span class="pl-c1">16</span>         <span class="pl-s1">inp</span> <span class="pl-c1">+=</span> [<span class="pl-en">int</span>(<span class="pl-s1">next_char</span>)]

<span class="pl-c1">/</span><span class="pl-s1">opt</span><span class="pl-c1">/</span><span class="pl-s1">conda</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">7</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">trax</span><span class="pl-c1">/</span><span class="pl-s1">layers</span><span class="pl-c1">/</span><span class="pl-s1">base</span>.<span class="pl-s1">py</span> <span class="pl-c1">in</span> <span class="pl-en">__call__</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">weights</span>, <span class="pl-s1">state</span>, <span class="pl-s1">rng</span>)
    <span class="pl-c1">165</span>       <span class="pl-s1">self</span>.<span class="pl-s1">state</span> <span class="pl-c1">=</span> <span class="pl-s1">state</span>  <span class="pl-c"># Needed if the model wasn't fully initialized.</span>
    <span class="pl-c1">166</span>     <span class="pl-s1">state</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">state</span>
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">167</span>     <span class="pl-s1">outputs</span>, <span class="pl-s1">new_state</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-en">pure_fn</span>(<span class="pl-s1">x</span>, <span class="pl-s1">weights</span>, <span class="pl-s1">state</span>, <span class="pl-s1">rng</span>)
    <span class="pl-c1">168</span>     <span class="pl-s1">self</span>.<span class="pl-s1">state</span> <span class="pl-c1">=</span> <span class="pl-s1">new_state</span>
    <span class="pl-c1">169</span>     <span class="pl-s1">self</span>.<span class="pl-s1">weights</span> <span class="pl-c1">=</span> <span class="pl-s1">weights</span>

<span class="pl-c1">/</span><span class="pl-s1">opt</span><span class="pl-c1">/</span><span class="pl-s1">conda</span><span class="pl-c1">/</span><span class="pl-s1">lib</span><span class="pl-c1">/</span><span class="pl-s1">python3</span>.<span class="pl-c1">7</span><span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">trax</span><span class="pl-c1">/</span><span class="pl-s1">layers</span><span class="pl-c1">/</span><span class="pl-s1">base</span>.<span class="pl-s1">py</span> <span class="pl-c1">in</span> <span class="pl-en">pure_fn</span>(<span class="pl-s1">self</span>, <span class="pl-s1">x</span>, <span class="pl-s1">weights</span>, <span class="pl-s1">state</span>, <span class="pl-s1">rng</span>, <span class="pl-s1">use_cache</span>)
    <span class="pl-c1">448</span>       <span class="pl-s1">name</span>, <span class="pl-s1">trace</span> <span class="pl-c1">=</span> <span class="pl-s1">self</span>.<span class="pl-s1">_name</span>, <span class="pl-en">_short_traceback</span>(<span class="pl-s1">skip</span><span class="pl-c1">=</span><span class="pl-c1">3</span>)
    <span class="pl-c1">449</span>       <span class="pl-k">raise</span> <span class="pl-v">LayerError</span>(<span class="pl-s1">name</span>, <span class="pl-s">'pure_fn'</span>,
<span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-c1">&gt;</span> <span class="pl-c1">450</span>                        <span class="pl-s1">self</span>.<span class="pl-s1">_caller</span>, <span class="pl-en">signature</span>(<span class="pl-s1">x</span>), <span class="pl-s1">trace</span>) <span class="pl-k">from</span> <span class="pl-c1">None</span>
    <span class="pl-c1">451</span> 
    <span class="pl-c1">452</span>   <span class="pl-k">def</span> <span class="pl-en">output_signature</span>(<span class="pl-s1">self</span>, <span class="pl-s1">input_signature</span>):

<span class="pl-v">LayerError</span>: <span class="pl-v">Exception</span> <span class="pl-s1">passing</span> <span class="pl-s1">through</span> <span class="pl-s1">layer</span> <span class="pl-v">Serial</span> (<span class="pl-s1">in</span> <span class="pl-s1">pure_fn</span>):
  <span class="pl-s1">layer</span> <span class="pl-s1">created</span> <span class="pl-c1">in</span> <span class="pl-s1">file</span> [...]<span class="pl-c1">/</span><span class="pl-c1">&lt;</span><span class="pl-s1">ipython</span><span class="pl-c1">-</span><span class="pl-s1">input</span><span class="pl-c1">-</span><span class="pl-c1">13</span><span class="pl-c1">-</span><span class="pl-c1">6</span><span class="pl-s1">f7ffe0c061e</span><span class="pl-c1">&gt;</span>, <span class="pl-s1">line</span> <span class="pl-c1">21</span>
  <span class="pl-s1">layer</span> <span class="pl-s1">input</span> <span class="pl-s1">shapes</span>: <span class="pl-v">ShapeDtype</span>{<span class="pl-s1">shape</span>:(<span class="pl-c1">1</span>, <span class="pl-c1">32</span>), <span class="pl-s1">dtype</span>:<span class="pl-s1">int32</span>}

  <span class="pl-v">File</span> [...]<span class="pl-c1">/</span><span class="pl-s1">trax</span><span class="pl-c1">/</span><span class="pl-s1">layers</span><span class="pl-c1">/</span><span class="pl-s1">combinators</span>.<span class="pl-s1">py</span>, <span class="pl-s1">line</span> <span class="pl-c1">88</span>, <span class="pl-s1">in</span> <span class="pl-s1">forward</span>
    <span class="pl-s1">outputs</span>, <span class="pl-s1">s</span> <span class="pl-c1">=</span> <span class="pl-s1">layer</span>.<span class="pl-en">pure_fn</span>(<span class="pl-s1">inputs</span>, <span class="pl-s1">w</span>, <span class="pl-s1">s</span>, <span class="pl-s1">rng</span>, <span class="pl-s1">use_cache</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)

<span class="pl-v">LayerError</span>: <span class="pl-v">Exception</span> <span class="pl-s1">passing</span> <span class="pl-s1">through</span> <span class="pl-s1">layer</span> <span class="pl-v">Embedding_256_512</span> (<span class="pl-s1">in</span> <span class="pl-s1">pure_fn</span>):
  <span class="pl-s1">layer</span> <span class="pl-s1">created</span> <span class="pl-c1">in</span> <span class="pl-s1">file</span> [...]<span class="pl-c1">/</span><span class="pl-c1">&lt;</span><span class="pl-s1">ipython</span><span class="pl-c1">-</span><span class="pl-s1">input</span><span class="pl-c1">-</span><span class="pl-c1">13</span><span class="pl-c1">-</span><span class="pl-c1">6</span><span class="pl-s1">f7ffe0c061e</span><span class="pl-c1">&gt;</span>, <span class="pl-s1">line</span> <span class="pl-c1">18</span>
  <span class="pl-s1">layer</span> <span class="pl-s1">input</span> <span class="pl-s1">shapes</span>: <span class="pl-v">ShapeDtype</span>{<span class="pl-s1">shape</span>:(<span class="pl-c1">1</span>, <span class="pl-c1">32</span>), <span class="pl-s1">dtype</span>:<span class="pl-s1">int32</span>}

  <span class="pl-v">File</span> [...]<span class="pl-c1">/</span><span class="pl-s1">trax</span><span class="pl-c1">/</span><span class="pl-s1">layers</span><span class="pl-c1">/</span><span class="pl-s1">core</span>.<span class="pl-s1">py</span>, <span class="pl-s1">line</span> <span class="pl-c1">150</span>, <span class="pl-s1">in</span> <span class="pl-s1">forward</span>
    <span class="pl-k">return</span> <span class="pl-s1">jnp</span>.<span class="pl-en">take</span>(<span class="pl-s1">self</span>.<span class="pl-s1">weights</span>, <span class="pl-s1">x</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">0</span>)

  <span class="pl-v">File</span> [...]<span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">/</span><span class="pl-s1">numpy</span><span class="pl-c1">/</span><span class="pl-s1">lax_numpy</span>.<span class="pl-s1">py</span>, <span class="pl-s1">line</span> <span class="pl-c1">3298</span>, <span class="pl-s1">in</span> <span class="pl-s1">take</span>
    <span class="pl-s1">slice_sizes</span><span class="pl-c1">=</span><span class="pl-en">tuple</span>(<span class="pl-s1">slice_sizes</span>))

  <span class="pl-v">File</span> [...]<span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">/</span><span class="pl-s1">lax</span><span class="pl-c1">/</span><span class="pl-s1">lax</span>.<span class="pl-s1">py</span>, <span class="pl-s1">line</span> <span class="pl-c1">835</span>, <span class="pl-s1">in</span> <span class="pl-s1">gather</span>
    <span class="pl-s1">slice_sizes</span><span class="pl-c1">=</span><span class="pl-en">canonicalize_shape</span>(<span class="pl-s1">slice_sizes</span>))

  <span class="pl-v">File</span> [...]<span class="pl-c1">/</span><span class="pl-s1">site</span><span class="pl-c1">-</span><span class="pl-s1">packages</span><span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">/</span><span class="pl-s1">core</span>.<span class="pl-s1">py</span>, <span class="pl-s1">line</span> <span class="pl-c1">273</span>, <span class="pl-s1">in</span> <span class="pl-s1">bind</span>
    <span class="pl-k">return</span> <span class="pl-s1">self</span>.<span class="pl-en">impl</span>(<span class="pl-c1">*</span><span class="pl-s1">args</span>, <span class="pl-c1">**</span><span class="pl-s1">kwargs</span>)

  <span class="pl-v">File</span> [...]<span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">/</span><span class="pl-s1">interpreters</span><span class="pl-c1">/</span><span class="pl-s1">xla</span>.<span class="pl-s1">py</span>, <span class="pl-s1">line</span> <span class="pl-c1">228</span>, <span class="pl-s1">in</span> <span class="pl-s1">apply_primitive</span>
    <span class="pl-s1">compiled_fun</span> <span class="pl-c1">=</span> <span class="pl-en">xla_primitive_callable</span>(<span class="pl-s1">prim</span>, <span class="pl-c1">*</span><span class="pl-en">unsafe_map</span>(<span class="pl-s1">arg_spec</span>, <span class="pl-s1">args</span>), <span class="pl-c1">**</span><span class="pl-s1">params</span>)

  <span class="pl-v">File</span> [...]<span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">/</span><span class="pl-s1">interpreters</span><span class="pl-c1">/</span><span class="pl-s1">xla</span>.<span class="pl-s1">py</span>, <span class="pl-s1">line</span> <span class="pl-c1">262</span>, <span class="pl-s1">in</span> <span class="pl-s1">xla_primitive_callable</span>
    <span class="pl-c1">*</span><span class="pl-s1">avals</span>, <span class="pl-c1">**</span><span class="pl-s1">params</span>)

  <span class="pl-v">File</span> [...]<span class="pl-c1">/</span><span class="pl-s1">jax</span><span class="pl-c1">/</span><span class="pl-s1">interpreters</span><span class="pl-c1">/</span><span class="pl-s1">xla</span>.<span class="pl-s1">py</span>, <span class="pl-s1">line</span> <span class="pl-c1">320</span>, <span class="pl-s1">in</span> <span class="pl-s1">primitive_computation</span>
    <span class="pl-k">raise</span> <span class="pl-v">RuntimeError</span>(<span class="pl-s1">msg</span>) <span class="pl-k">from</span> <span class="pl-s1">e</span>

<span class="pl-v">RuntimeError</span>: <span class="pl-v">Invalid</span> <span class="pl-s1">argument</span>: <span class="pl-v">Slice</span> <span class="pl-s1">size</span> <span class="pl-s1">at</span> <span class="pl-s1">index</span> <span class="pl-c1">0</span> <span class="pl-c1">in</span> <span class="pl-s1">gather</span> <span class="pl-s1">op</span> <span class="pl-c1">is</span> <span class="pl-s1">out</span> <span class="pl-s1">of</span> <span class="pl-s1">range</span>, <span class="pl-s1">must</span> <span class="pl-s1">be</span> <span class="pl-s1">within</span> [<span class="pl-c1">0</span>, <span class="pl-c1">1</span>), <span class="pl-s1">got</span> <span class="pl-c1">1.</span>: 
<span class="pl-v">This</span> <span class="pl-c1">is</span> <span class="pl-s1">a</span> <span class="pl-s1">bug</span> <span class="pl-c1">in</span> <span class="pl-v">JAX</span>'<span class="pl-s1">s</span> <span class="pl-s1">shape</span><span class="pl-c1">-</span><span class="pl-s1">checking</span> <span class="pl-s1">rules</span>; <span class="pl-s1">please</span> <span class="pl-s1">report</span> <span class="pl-s1">it</span>!</pre></div>