<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="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))"><pre class="notranslate"><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">np</span>
<span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">ops</span> <span class="pl-k">as</span> <span class="pl-s1">jo</span>

<span class="pl-s1">f</span> <span class="pl-c1">=</span> <span class="pl-k">lambda</span> <span class="pl-s1">x</span>, <span class="pl-v">X</span>: <span class="pl-s1">jo</span>.<span class="pl-en">index_update</span>(<span class="pl-v">X</span>, <span class="pl-s1">jo</span>.<span class="pl-s1">index</span>[<span class="pl-c1">0</span>], <span class="pl-s1">x</span>)
<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">zeros</span>(<span class="pl-c1">2</span>)
<span class="pl-v">X</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">zeros</span>(<span class="pl-c1">2</span>)

<span class="pl-c"># RuntimeError: Invalid argument: Updates tensor must be of rank 0; got 1.: </span>
<span class="pl-c"># This is a bug in JAX's shape-checking rules; please report it!</span>
<span class="pl-en">print</span>(<span class="pl-en">f</span>(<span class="pl-s1">x</span>, <span class="pl-v">X</span>))</pre></div>
<p dir="auto">The jaxpr for <code class="notranslate">f</code> in context of <code class="notranslate">x, X</code> is:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="{ 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] }"><pre class="notranslate"><code class="notranslate">{ 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] }
</code></pre></div>
<p dir="auto">Similarly to <a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="496460953" data-permission-text="Title is private" data-url="https://github.com/google/jax/issues/1379" data-hovercard-type="issue" data-hovercard-url="/google/jax/issues/1379/hovercard" href="https://github.com/google/jax/issues/1379">#1379</a>, shape checking should be done prior to execution with XLA:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as np

jax.ops.index_update(np.zeros(3), 0, np.ones(1))"><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">np</span>

<span class="pl-s1">jax</span>.<span class="pl-s1">ops</span>.<span class="pl-en">index_update</span>(<span class="pl-s1">np</span>.<span class="pl-en">zeros</span>(<span class="pl-c1">3</span>), <span class="pl-c1">0</span>, <span class="pl-s1">np</span>.<span class="pl-en">ones</span>(<span class="pl-c1">1</span>))</pre></div>
<p dir="auto">yields:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-2-d87d74931298&gt; in &lt;module&gt;()
      2 import jax.numpy as np
      3 
----&gt; 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   &quot;&quot;&quot;
--&gt; 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   )
---&gt; 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,
--&gt; 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:
--&gt; 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   &quot;&quot;&quot;Impl rule that compiles and runs a single primitive 'prim' using XLA.&quot;&quot;&quot;
    122   abstract_args = map(abstractify, args)
--&gt; 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))
--&gt; 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            &quot;This is a bug in JAX's shape-checking rules; please report it!\n&quot;
    167            &quot;https://github.com/google/jax/issues\n&quot;)
--&gt; 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"><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
&lt;ipython-input-2-d87d74931298&gt; in &lt;module&gt;()
      2 import jax.numpy as np
      3 
----&gt; 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   """
--&gt; 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   )
---&gt; 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,
--&gt; 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:
--&gt; 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)
--&gt; 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))
--&gt; 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")
--&gt; 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
</code></pre></div>