Issue 1:
Description

Checked for duplicates but I might have missed something.

The following snippet returns inconsistent results between CPU and GPU runtimes. I checked this on Google Colab.

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)

On CPU it returns:

DeviceArray([0, 2, 4], dtype=int32)

On GPU it returns:

DeviceArray([0, 3, 4], dtype=int32)

What jax/jaxlib version are you using?

jax v0.3.17, jaxlib v0.3.15

Which accelerator(s) are you using?

CPU/GPU

Additional System Info

Google Colab

Issue 2:
Please:

Check for duplicate issues.

Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

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')

Calling run_optimizer_2.delay() or run_optimizer_2() from the debug console will cause the segfault.

If applicable, include full error messages/tracebacks.

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 "<string>", line 1 in <module>
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 <module>
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 <module>
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)'

Not sure what's up here. Potentially not an issue with JAX, but it works fine with @jax.jit disba