<h3 dir="auto">Description</h3>
<p dir="auto">Checked for duplicates but I might have missed something.</p>
<p dir="auto">The following snippet returns inconsistent results between CPU and GPU runtimes. I checked this on Google Colab.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax import lax
import numpy as np

inp = np.array([0, 0, 0])
indices = np.array([[1], [2], [1]])
updates = np.array([3, 4, 2])

update_window_dims = ()
inserted_window_dims = (0,)
scatter_dims_to_operand_dims = (0,)

scatter_dim_nos = lax.ScatterDimensionNumbers(update_window_dims=update_window_dims,
                                              inserted_window_dims=inserted_window_dims,
                                              scatter_dims_to_operand_dims=scatter_dims_to_operand_dims)

lax.scatter(inp, indices, updates, scatter_dim_nos)"><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">lax</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">inp</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>([<span class="pl-c1">0</span>, <span class="pl-c1">0</span>, <span class="pl-c1">0</span>])
<span class="pl-s1">indices</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>([[<span class="pl-c1">1</span>], [<span class="pl-c1">2</span>], [<span class="pl-c1">1</span>]])
<span class="pl-s1">updates</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>([<span class="pl-c1">3</span>, <span class="pl-c1">4</span>, <span class="pl-c1">2</span>])

<span class="pl-s1">update_window_dims</span> <span class="pl-c1">=</span> ()
<span class="pl-s1">inserted_window_dims</span> <span class="pl-c1">=</span> (<span class="pl-c1">0</span>,)
<span class="pl-s1">scatter_dims_to_operand_dims</span> <span class="pl-c1">=</span> (<span class="pl-c1">0</span>,)

<span class="pl-s1">scatter_dim_nos</span> <span class="pl-c1">=</span> <span class="pl-s1">lax</span>.<span class="pl-v">ScatterDimensionNumbers</span>(<span class="pl-s1">update_window_dims</span><span class="pl-c1">=</span><span class="pl-s1">update_window_dims</span>,
                                              <span class="pl-s1">inserted_window_dims</span><span class="pl-c1">=</span><span class="pl-s1">inserted_window_dims</span>,
                                              <span class="pl-s1">scatter_dims_to_operand_dims</span><span class="pl-c1">=</span><span class="pl-s1">scatter_dims_to_operand_dims</span>)

<span class="pl-s1">lax</span>.<span class="pl-en">scatter</span>(<span class="pl-s1">inp</span>, <span class="pl-s1">indices</span>, <span class="pl-s1">updates</span>, <span class="pl-s1">scatter_dim_nos</span>)</pre></div>
<p dir="auto">On CPU it returns:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="DeviceArray([0, 2, 4], dtype=int32)"><pre class="notranslate"><code class="notranslate">DeviceArray([0, 2, 4], dtype=int32)
</code></pre></div>
<p dir="auto">On GPU it returns:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="DeviceArray([0, 3, 4], dtype=int32)"><pre class="notranslate"><code class="notranslate">DeviceArray([0, 3, 4], dtype=int32)
</code></pre></div>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">jax v0.3.17, jaxlib v0.3.15</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">CPU/GPU</p>
<h3 dir="auto">Additional System Info</h3>
<p dir="auto">Google Colab</p>
<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="from celery import Celery
from celery import shared_task

from functools import partial

import jax
import jax.numpy as jnp
import numba as nb
import numpy as np
from scipy.optimize import minimize

# app = Celery('test', broker='redis://localhost/3')

@jax.jit
def scale_down_2(x):
    return x[:-1]

@shared_task
def run_optimizer_2():
    nAC = 20
    guess_equal_wts = np.empty(nAC)
    for i in range(nAC):
        guess_equal_wts[i] = 1/nAC

    scale_down_2(guess_equal_wts)
    print('z')"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">celery</span> <span class="pl-k">import</span> <span class="pl-v">Celery</span>
<span class="pl-k">from</span> <span class="pl-s1">celery</span> <span class="pl-k">import</span> <span class="pl-s1">shared_task</span>

<span class="pl-k">from</span> <span class="pl-s1">functools</span> <span class="pl-k">import</span> <span class="pl-s1">partial</span>

<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">numba</span> <span class="pl-k">as</span> <span class="pl-s1">nb</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">from</span> <span class="pl-s1">scipy</span>.<span class="pl-s1">optimize</span> <span class="pl-k">import</span> <span class="pl-s1">minimize</span>

<span class="pl-c"># app = Celery('test', broker='redis://localhost/3')</span>

<span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">scale_down_2</span>(<span class="pl-s1">x</span>):
    <span class="pl-k">return</span> <span class="pl-s1">x</span>[:<span class="pl-c1">-</span><span class="pl-c1">1</span>]

<span class="pl-en">@<span class="pl-s1">shared_task</span></span>
<span class="pl-k">def</span> <span class="pl-en">run_optimizer_2</span>():
    <span class="pl-s1">nAC</span> <span class="pl-c1">=</span> <span class="pl-c1">20</span>
    <span class="pl-s1">guess_equal_wts</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">empty</span>(<span class="pl-s1">nAC</span>)
    <span class="pl-k">for</span> <span class="pl-s1">i</span> <span class="pl-c1">in</span> <span class="pl-en">range</span>(<span class="pl-s1">nAC</span>):
        <span class="pl-s1">guess_equal_wts</span>[<span class="pl-s1">i</span>] <span class="pl-c1">=</span> <span class="pl-c1">1</span><span class="pl-c1">/</span><span class="pl-s1">nAC</span>

    <span class="pl-en">scale_down_2</span>(<span class="pl-s1">guess_equal_wts</span>)
    <span class="pl-en">print</span>(<span class="pl-s">'z'</span>)</pre></div>
<p dir="auto">Calling <code class="notranslate">run_optimizer_2.delay()</code> or <code class="notranslate">run_optimizer_2()</code> from the debug console will cause the segfault.</p>
<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>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 1265, in mark_as_worker_lost
    raise WorkerLostError(
billiard.exceptions.WorkerLostError: Worker exited prematurely: signal 11 (SIGSEGV) Job: 0.
Fatal Python error: Segmentation fault

Thread 0x00007f1791ffb700 (most recent call first):
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 316 in wait
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 574 in wait
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/pydevd.py&quot;, line 245 in _on_run
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py&quot;, line 46 in run
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 973 in _bootstrap_inner
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 930 in _bootstrap
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py&quot;, line 1054 in __call__

Current thread 0x00007f17917fa700 (most recent call first):
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/interpreters/xla.py&quot;, line 874 in _execute_compiled
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/interpreters/xla.py&quot;, line 581 in _xla_call_impl
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/core.py&quot;, line 609 in process_call
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/core.py&quot;, line 1568 in process
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/core.py&quot;, line 1556 in call_bind
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/core.py&quot;, line 1565 in bind
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/_src/api.py&quot;, line 426 in cache_miss
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/_src/traceback_util.py&quot;, line 143 in reraise_with_filtered_traceback
  File &quot;/home/jorgen/django-2/caravel_cmpas/engine/optimizer.py&quot;, line 1130 in run_optimizer_2
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/app/task.py&quot;, line 389 in __call__
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/app/trace.py&quot;, line 732 in __protected_call__
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/local.py&quot;, line 188 in __call__
  File &quot;&lt;string&gt;&quot;, line 1 in &lt;module&gt;
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py&quot;, line 424 in evaluate_expression
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py&quot;, line 310 in _run_with_interrupt_thread
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py&quot;, line 339 in _run_with_unblock_threads
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py&quot;, line 371 in new_func
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py&quot;, line 1207 in internal_evaluate_expression_json
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_constants.py&quot;, line 511 in new_func
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py&quot;, line 542 in do_it
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py&quot;, line 574 in do_it
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/pydevd.py&quot;, line 1688 in process_internal_commands
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/pydevd.py&quot;, line 195 in _on_run
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py&quot;, line 46 in run
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 973 in _bootstrap_inner
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 930 in _bootstrap
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py&quot;, line 1054 in __call__

Thread 0x00007f1790ff9700 (most recent call first):
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 316 in wait
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 574 in wait
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_timeout.py&quot;, line 43 in _on_run
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py&quot;, line 46 in run
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 973 in _bootstrap_inner
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 930 in _bootstrap
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py&quot;, line 1054 in __call__

Thread 0x00007f177bfff700 (most recent call first):
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py&quot;, line 219 in _read_line
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py&quot;, line 237 in _on_run
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py&quot;, line 46 in run
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 973 in _bootstrap_inner
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 930 in _bootstrap
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py&quot;, line 1054 in __call__

Thread 0x00007f177b7fe700 (most recent call first):
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 316 in wait
  File &quot;/usr/lib/python3.9/queue.py&quot;, line 180 in get
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py&quot;, line 382 in _on_run
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py&quot;, line 46 in run
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 973 in _bootstrap_inner
  File &quot;/usr/lib/python3.9/threading.py&quot;, line 930 in _bootstrap
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py&quot;, line 1054 in __call__

Thread 0x00007f19e3595740 (most recent call first):
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/connection.py&quot;, line 422 in _recv
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/connection.py&quot;, line 456 in _recv_bytes
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/connection.py&quot;, line 243 in recv_bytes
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/queues.py&quot;, line 355 in get_payload
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 445 in _recv
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 473 in receive
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 351 in workloop
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 292 in __call__
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/process.py&quot;, line 114 in run
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/process.py&quot;, line 327 in _bootstrap
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/popen_fork.py&quot;, line 79 in _launch
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/popen_fork.py&quot;, line 24 in __init__
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/context.py&quot;, line 333 in _Popen
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/process.py&quot;, line 124 in start
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 1158 in _create_worker_process
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/concurrency/asynpool.py&quot;, line 477 in _create_worker_process
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 1328 in _repopulate_pool
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 1343 in _maintain_pool
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py&quot;, line 1351 in maintain_pool
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/concurrency/asynpool.py&quot;, line 482 in _event_process_exit
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/kombu/asynchronous/hub.py&quot;, line 361 in create_loop
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/worker/loops.py&quot;, line 81 in asynloop
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/worker/consumer/consumer.py&quot;, line 618 in start
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bootsteps.py&quot;, line 116 in start
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/worker/consumer/consumer.py&quot;, line 326 in start
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bootsteps.py&quot;, line 365 in start
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bootsteps.py&quot;, line 116 in start
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/worker/worker.py&quot;, line 203 in start
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bin/worker.py&quot;, line 346 in worker
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bin/base.py&quot;, line 133 in caller
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/decorators.py&quot;, line 21 in new_func
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py&quot;, line 610 in invoke
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py&quot;, line 1066 in invoke
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py&quot;, line 1259 in invoke
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py&quot;, line 782 in main
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py&quot;, line 829 in __call__
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bin/celery.py&quot;, line 213 in main
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/__main__.py&quot;, line 15 in main
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/__main__.py&quot;, line 19 in &lt;module&gt;
  File &quot;/usr/lib/python3.9/runpy.py&quot;, line 87 in _run_code
  File &quot;/usr/lib/python3.9/runpy.py&quot;, line 197 in _run_module_as_main
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/server/cli.py&quot;, line 331 in run_module
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/server/cli.py&quot;, line 444 in main
  File &quot;/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/__main__.py&quot;, line 45 in &lt;module&gt;
  File &quot;/usr/lib/python3.9/runpy.py&quot;, line 87 in _run_code
  File &quot;/usr/lib/python3.9/runpy.py&quot;, line 197 in _run_module_as_main
Process 'ForkPoolWorker-25' pid:59626 exited with 'signal 11 (SIGSEGV)'"><pre lang="ERROR:celery.worker.request:Task" class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 1265, in mark_as_worker_lost
    raise WorkerLostError(
billiard.exceptions.WorkerLostError: Worker exited prematurely: signal 11 (SIGSEGV) Job: 0.
Fatal Python error: Segmentation fault

Thread 0x00007f1791ffb700 (most recent call first):
  File "/usr/lib/python3.9/threading.py", line 316 in wait
  File "/usr/lib/python3.9/threading.py", line 574 in wait
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/pydevd.py", line 245 in _on_run
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py", line 46 in run
  File "/usr/lib/python3.9/threading.py", line 973 in _bootstrap_inner
  File "/usr/lib/python3.9/threading.py", line 930 in _bootstrap
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py", line 1054 in __call__

Current thread 0x00007f17917fa700 (most recent call first):
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/interpreters/xla.py", line 874 in _execute_compiled
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/interpreters/xla.py", line 581 in _xla_call_impl
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/core.py", line 609 in process_call
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/core.py", line 1568 in process
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/core.py", line 1556 in call_bind
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/core.py", line 1565 in bind
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/_src/api.py", line 426 in cache_miss
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 143 in reraise_with_filtered_traceback
  File "/home/jorgen/django-2/caravel_cmpas/engine/optimizer.py", line 1130 in run_optimizer_2
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/app/task.py", line 389 in __call__
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/app/trace.py", line 732 in __protected_call__
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/local.py", line 188 in __call__
  File "&lt;string&gt;", line 1 in &lt;module&gt;
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py", line 424 in evaluate_expression
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py", line 310 in _run_with_interrupt_thread
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py", line 339 in _run_with_unblock_threads
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_vars.py", line 371 in new_func
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py", line 1207 in internal_evaluate_expression_json
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_constants.py", line 511 in new_func
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py", line 542 in do_it
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py", line 574 in do_it
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/pydevd.py", line 1688 in process_internal_commands
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/pydevd.py", line 195 in _on_run
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py", line 46 in run
  File "/usr/lib/python3.9/threading.py", line 973 in _bootstrap_inner
  File "/usr/lib/python3.9/threading.py", line 930 in _bootstrap
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py", line 1054 in __call__

Thread 0x00007f1790ff9700 (most recent call first):
  File "/usr/lib/python3.9/threading.py", line 316 in wait
  File "/usr/lib/python3.9/threading.py", line 574 in wait
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_timeout.py", line 43 in _on_run
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py", line 46 in run
  File "/usr/lib/python3.9/threading.py", line 973 in _bootstrap_inner
  File "/usr/lib/python3.9/threading.py", line 930 in _bootstrap
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py", line 1054 in __call__

Thread 0x00007f177bfff700 (most recent call first):
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py", line 219 in _read_line
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py", line 237 in _on_run
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py", line 46 in run
  File "/usr/lib/python3.9/threading.py", line 973 in _bootstrap_inner
  File "/usr/lib/python3.9/threading.py", line 930 in _bootstrap
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py", line 1054 in __call__

Thread 0x00007f177b7fe700 (most recent call first):
  File "/usr/lib/python3.9/threading.py", line 316 in wait
  File "/usr/lib/python3.9/queue.py", line 180 in get
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_comm.py", line 382 in _on_run
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_daemon_thread.py", line 46 in run
  File "/usr/lib/python3.9/threading.py", line 973 in _bootstrap_inner
  File "/usr/lib/python3.9/threading.py", line 930 in _bootstrap
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/_vendored/pydevd/_pydev_bundle/pydev_monkey.py", line 1054 in __call__

Thread 0x00007f19e3595740 (most recent call first):
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/connection.py", line 422 in _recv
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/connection.py", line 456 in _recv_bytes
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/connection.py", line 243 in recv_bytes
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/queues.py", line 355 in get_payload
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 445 in _recv
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 473 in receive
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 351 in workloop
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 292 in __call__
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/process.py", line 114 in run
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/process.py", line 327 in _bootstrap
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/popen_fork.py", line 79 in _launch
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/popen_fork.py", line 24 in __init__
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/context.py", line 333 in _Popen
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/process.py", line 124 in start
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 1158 in _create_worker_process
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/concurrency/asynpool.py", line 477 in _create_worker_process
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 1328 in _repopulate_pool
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 1343 in _maintain_pool
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/billiard/pool.py", line 1351 in maintain_pool
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/concurrency/asynpool.py", line 482 in _event_process_exit
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/kombu/asynchronous/hub.py", line 361 in create_loop
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/worker/loops.py", line 81 in asynloop
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/worker/consumer/consumer.py", line 618 in start
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bootsteps.py", line 116 in start
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/worker/consumer/consumer.py", line 326 in start
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bootsteps.py", line 365 in start
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bootsteps.py", line 116 in start
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/worker/worker.py", line 203 in start
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bin/worker.py", line 346 in worker
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bin/base.py", line 133 in caller
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/decorators.py", line 21 in new_func
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py", line 610 in invoke
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py", line 1066 in invoke
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py", line 1259 in invoke
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py", line 782 in main
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/click/core.py", line 829 in __call__
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/bin/celery.py", line 213 in main
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/__main__.py", line 15 in main
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/celery/__main__.py", line 19 in &lt;module&gt;
  File "/usr/lib/python3.9/runpy.py", line 87 in _run_code
  File "/usr/lib/python3.9/runpy.py", line 197 in _run_module_as_main
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/server/cli.py", line 331 in run_module
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/server/cli.py", line 444 in main
  File "/home/jorgen/django-2/venv/lib/python3.9/site-packages/debugpy/__main__.py", line 45 in &lt;module&gt;
  File "/usr/lib/python3.9/runpy.py", line 87 in _run_code
  File "/usr/lib/python3.9/runpy.py", line 197 in _run_module_as_main
Process 'ForkPoolWorker-25' pid:59626 exited with 'signal 11 (SIGSEGV)'
</code></pre></div>
<p dir="auto">Not sure what's up here. Potentially not an issue with JAX, but it works fine with @jax.jit disba</p>