Issue 1:
import jax.numpy as np
import jax.ops as jo

f = lambda x, X: jo.index_update(X, jo.index[0], x)
x = np.zeros(2)
X = np.zeros(2)

# RuntimeError: Invalid argument: Updates tensor must be of rank 0; got 1.:
# This is a bug in JAX's shape-checking rules; please report it!
print(f(x, X))

The jaxpr for f in context of x, X is:

{ lambda c ; ; a b.
let d = broadcast[ sizes=() ] a
e = scatter[ updates_shape=(2,)
update_jaxpr={ lambda ; ; a b.
let
in [b] }
dimension_numbers=ScatterDimensionNumbers(update_window_dims=(), inserted_window_dims=(0,), scatter_dims_to_operand_dims=(0,))
update_consts=() ] b c d
in [e] }

Issue 2:
Similarly to #1379, shape checking should be done prior to execution with XLA:

import jax
import jax.numpy as np

jax.ops.index_update(np.zeros(3), 0, np.ones(1))

yields:

---------------------------------------------------------------------------
RuntimeError Traceback (most recent call last)
<ipython-input-2-d87d74931298> in <module>()
2 import jax.numpy as np
3
----> 4 jax.ops.index_update(np.zeros(3), 0, np.ones(1))

6 frames
google3/third_party/py/jax/ops/scatter.py in index_update(x, idx, y)
236 [1., 1., 1., 6., 6., 6.]], dtype=float32)
237 """
--> 238 return _scatter_update(x, idx, y, lax.scatter)
239
240 def segment_sum(data, segment_ids, num_segments=None):

google3/third_party/py/jax/ops/scatter.py in _scatter_update(x, idx, y, scatter_op)
70 scatter_dims_to_operand_dims=indexer.dnums.start_index_map
71 )
---> 72 return scatter_op(x, indexer.gather_indices, y, dnums)
73
74

google3/third_party/py/jax/lax/lax.py in scatter(operand, scatter_indices, updates, dimension_numbers)
813 operand, scatter_indices, updates, update_jaxpr=jaxpr,
814 update_consts=consts, dimension_numbers=dimension_numbers,
--> 815 updates_shape=updates.shape)
816
817 def index_take(src, idxs, axes):

google3/third_party/py/jax/core.py in bind(self, *args, **kwargs)
128 top_trace = find_top_trace(args)
129 if top_trace is None:
--> 130 return self.impl(*args, **kwargs)
131
132 tracers = map(top_trace.full_raise, args)

google3/third_party/py/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
121 """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
122 abstract_args = map(abstractify, args)
--> 123 compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
124 return compiled_fun(*args)
125

google3/third_party/py/jax/interpreters/xla.py in xla_primitive_callable(prim, *abstract_args, **params)
134 handle_result = aval_to_result_handler(aval_out)
135 xla_shapes = tuple(map(aval_to_xla_shape, abstract_args))
--> 136 built_c = primitive_computation(prim, *xla_shapes, **params)
137 compiled = built_c.Compile(xla_shapes, xb.get_compile_options(),
138 backend=xb.get_backend(backend))

google3/third_party/py/jax/interpreters/xla.py in primitive_computation(prim, *xla_shapes, **params)
166 "This is a bug in JAX's shape-checking rules; please report it!\n"
167 "https://github.com/google/jax/issues\n")
--> 168 raise RuntimeError(msg)
169
170 def _execute_compiled_primitive(prim, compiled, backend, result_handler, *args):

RuntimeError: Invalid argument: Updates tensor must be of rank 0; got 1.: @ 0x55983eb37dd5 xla::Scatter()
@ 0x7ff8c9077364 (unknown)
@ 0x7ff8c9082e6b (unknown)
@ 0x55984d75256d _PyCFunction_FastCallDict
@ 0x55984d033a82 _PyEval_EvalFrameDefault
@ 0x55984ca5d7ce _PyEval_EvalCodeWithName.llvm.293474707488941182
@ 0x55984d02f272 _PyEval_EvalFrameDefault
@ 0x55984ca5d7ce _PyEval_EvalCodeWithName.llvm.293474707488941182
@ 0x559850dded7c function_call
@ 0x55984d033564 _PyEval_EvalFrameDefault
@ 0x55984ca5d7ce _PyEval_EvalCodeWithName.llvm.293474707488941182
@ 0x559850dded7c function_call
@ 0x55984ca51daf PyObject_Call
@ 0x7ff8cacc0bda (unknown)
@ 0x55984d033564 _PyEval_EvalFrameDefault
@ 0x55984ca5d7ce _PyEval_EvalCodeWithName.llvm.293474707488941182
@ 0x559850dded7c function_call
@ 0x55984ca51daf PyObject_Call
@ 0x7ff8cacc0bda (unknown)
@ 0x55984d033564 _PyEval_EvalFrameDefault
@ 0x55984ca5d7ce _PyEval_EvalCodeWithName.llvm.293474707488941182
@ 0x559850dded7c function_call
@ 0x55984ca51daf PyObject_Call
@ 0x55984fa3f6ed partial_call
@ 0x55984d033564 _PyEval_EvalFrameDefault
@ 0x55984ca5d7ce _PyEval_EvalCodeWithName.llvm.293474707488941182
@ 0x55984d02f1c1 _PyEval_EvalFrameDefault
@ 0x55984d02c7d0 _PyEval_EvalFrameDefault
@ 0x55984d02c7d0 _PyEval_EvalFrameDefault
@ 0x55984d02c7d0 _PyEval_EvalFrameDefault
@ 0x55984ca5d7ce _PyEval_EvalCodeWithName.llvm.293474707488941182
@ 0x5598480ed679 builtin_exec

This is a bug in JAX's shape-checking rules; please report it!
https://github.com/google/jax/issues