<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>
<p dir="auto"><code class="notranslate">average</code> will succeed if <code class="notranslate">axis=False</code> but fail when using <code class="notranslate">value_and_grads</code></p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax

def fn(arg_0):
    return jax.numpy.average(arg_0, axis=False).sum()

mykey = jax.random.PRNGKey(92127318)
array = jax.random.uniform(mykey, [2, 3, 4], jax.numpy.float32, minval=0, maxval=2)

arg_0 = array.clone()
print(fn(arg_0))
# 12.970601

arg_0 = array.clone()
jax.value_and_grad(fn, (0))(arg_0)
# ValueError: boolean array argument obj to delete must be one dimensional and match the axis length of 3"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-k">def</span> <span class="pl-en">fn</span>(<span class="pl-s1">arg_0</span>):
    <span class="pl-k">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-en">average</span>(<span class="pl-s1">arg_0</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-c1">False</span>).<span class="pl-en">sum</span>()

<span class="pl-s1">mykey</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">92127318</span>)
<span class="pl-s1">array</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">mykey</span>, [<span class="pl-c1">2</span>, <span class="pl-c1">3</span>, <span class="pl-c1">4</span>], <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-s1">float32</span>, <span class="pl-s1">minval</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">maxval</span><span class="pl-c1">=</span><span class="pl-c1">2</span>)

<span class="pl-s1">arg_0</span> <span class="pl-c1">=</span> <span class="pl-s1">array</span>.<span class="pl-en">clone</span>()
<span class="pl-en">print</span>(<span class="pl-en">fn</span>(<span class="pl-s1">arg_0</span>))
<span class="pl-c"># 12.970601</span>

<span class="pl-s1">arg_0</span> <span class="pl-c1">=</span> <span class="pl-s1">array</span>.<span class="pl-en">clone</span>()
<span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">fn</span>, (<span class="pl-c1">0</span>))(<span class="pl-s1">arg_0</span>)
<span class="pl-c"># ValueError: boolean array argument obj to delete must be one dimensional and match the axis length of 3</span></pre></div>
<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="---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation         Traceback (most recent call last)
/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/runpy.py in _run_module_as_main(***failed resolving arguments***)
    196         sys.argv[0] = mod_spec.origin
--&gt; 197     return _run_code(code, main_globals, None,
    198                      &quot;__main__&quot;, mod_spec)

/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/runpy.py in _run_code(***failed resolving arguments***)
     86                        __spec__ = mod_spec)
---&gt; 87     exec(code, run_globals)
     88     return run_globals

~/Library/Python/3.9/lib/python/site-packages/ipykernel_launcher.py in &lt;module&gt;
     15     from ipykernel import kernelapp as app
---&gt; 16     app.launch_new_instance()

~/Library/Python/3.9/lib/python/site-packages/traitlets/config/application.py in launch_instance(***failed resolving arguments***)
    844         app.initialize(argv)
--&gt; 845         app.start()
    846 

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelapp.py in start(***failed resolving arguments***)
    618             try:
--&gt; 619                 self.io_loop.start()
    620             except KeyboardInterrupt:

~/Library/Python/3.9/lib/python/site-packages/tornado/platform/asyncio.py in start(***failed resolving arguments***)
    198             asyncio.set_event_loop(self.asyncio_loop)
--&gt; 199             self.asyncio_loop.run_forever()
    200         finally:

/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/base_events.py in run_forever(***failed resolving arguments***)
    600             while True:
--&gt; 601                 self._run_once()
    602                 if self._stopping:

/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/base_events.py in _run_once(***failed resolving arguments***)
   1904             else:
-&gt; 1905                 handle._run()
   1906         handle = None  # Needed to break cycles when an exception occurs.

/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/events.py in _run(***failed resolving arguments***)
     79         try:
---&gt; 80             self._context.run(self._callback, *self._args)
     81         except (SystemExit, KeyboardInterrupt):

~/Library/Python/3.9/lib/python/site-packages/tornado/ioloop.py in &lt;lambda&gt;(***failed resolving arguments***)
    687             future.add_done_callback(
--&gt; 688                 lambda f: self._run_callback(functools.partial(callback, future))
    689             )

~/Library/Python/3.9/lib/python/site-packages/tornado/ioloop.py in _run_callback(***failed resolving arguments***)
    740         try:
--&gt; 741             ret = callback()
    742             if ret is not None:

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in inner(***failed resolving arguments***)
    813                 f = None  # noqa: F841
--&gt; 814                 self.ctx_run(self.run)
    815 

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in run(***failed resolving arguments***)
    774                     else:
--&gt; 775                         yielded = self.gen.send(value)
    776 

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelbase.py in process_one(***failed resolving arguments***)
    357                 return None
--&gt; 358         yield gen.maybe_future(dispatch(*args))
    359 

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--&gt; 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelbase.py in dispatch_shell(***failed resolving arguments***)
    260             try:
--&gt; 261                 yield gen.maybe_future(handler(stream, idents, msg))
    262             except Exception:

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--&gt; 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelbase.py in execute_request(***failed resolving arguments***)
    535         reply_content = yield gen.maybe_future(
--&gt; 536             self.do_execute(
    537                 code, silent, store_history,

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--&gt; 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/Library/Python/3.9/lib/python/site-packages/ipykernel/ipkernel.py in do_execute(***failed resolving arguments***)
    301                 # letting shell dispatch to loop runners
--&gt; 302                 res = shell.run_cell(code, store_history=store_history, silent=silent)
    303         finally:

~/Library/Python/3.9/lib/python/site-packages/ipykernel/zmqshell.py in run_cell(***failed resolving arguments***)
    538         self._last_traceback = None
--&gt; 539         return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
    540 

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in run_cell(***failed resolving arguments***)
   2897         try:
-&gt; 2898             result = self._run_cell(
   2899                 raw_cell, store_history, silent, shell_futures)

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in _run_cell(***failed resolving arguments***)
   2943         try:
-&gt; 2944             return runner(coro)
   2945         except BaseException as e:

~/Library/Python/3.9/lib/python/site-packages/IPython/core/async_helpers.py in _pseudo_sync_runner(***failed resolving arguments***)
     67     try:
---&gt; 68         coro.send(None)
     69     except StopIteration as exc:

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in run_cell_async(***failed resolving arguments***)
   3168 
-&gt; 3169                 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
   3170                        interactivity=interactivity, compiler=compiler, result=result)

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in run_ast_nodes(***failed resolving arguments***)
   3360                         asy = compare(code)
-&gt; 3361                     if (await self.run_code(code, result,  async_=asy)):
   3362                         return True

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in run_code(***failed resolving arguments***)
   3440                 else:
-&gt; 3441                     exec(code_obj, self.user_global_ns, self.user_ns)
   3442             finally:

&lt;ipython-input-16-d01aa5baed3f&gt; in &lt;module&gt;
     17     arg_0 = arg_0_array.clone()
---&gt; 18     results['res_2'], _ = jax.value_and_grad(fn, (0))(arg_0)
     19 except Exception as e:

&lt;ipython-input-16-d01aa5baed3f&gt; in fn(***failed resolving arguments***)
      5     returned = False
----&gt; 6     return jax.numpy.average(arg_0, arg_1, returned=returned).sum()
      7 try:

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in average(***failed resolving arguments***)
    284             returned=False):
--&gt; 285   return _average(a, _ensure_optional_axes(axis), weights, returned)
    286 

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in _average(***failed resolving arguments***)
    292   if weights is None: # Treat all weights as 1
--&gt; 293     avg = mean(a, axis=axis)
    294     if axis is None:

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in mean(***failed resolving arguments***)
    251          out=None, keepdims=False, *, where=None):
--&gt; 252   return _mean(a, _ensure_optional_axes(axis), dtype, out, keepdims,
    253                where=where)

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in _mean(***failed resolving arguments***)
    278   return lax.div(
--&gt; 279       sum(a, axis, dtype=dtype, keepdims=keepdims, where=where),
    280       lax.convert_element_type(normalizer, dtype))

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in sum(***failed resolving arguments***)
    164         out=None, keepdims=None, initial=None, where=None):
--&gt; 165   return _reduce_sum(a, axis=_ensure_optional_axes(axis), dtype=dtype, out=out,
    166                      keepdims=keepdims, initial=initial, where=where)

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in _reduce_sum(***failed resolving arguments***)
    156                 dtype=None, out=None, keepdims=None, initial=None, where=None):
--&gt; 157   return _reduction(a, &quot;sum&quot;, np.sum, lax.add, 0,
    158                     bool_op=lax.bitwise_or, upcast_f16_for_computation=True,

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in _reduction(***failed resolving arguments***)
     97   else:
---&gt; 98     result = lax.reduce(a, init_val, op, dims)
     99   if initial is not None:

JaxStackTraceBeforeTransformation: ValueError: boolean array argument obj to delete must be one dimensional and match the axis length of 3

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
&lt;ipython-input-25-7d77a6f1d189&gt; in &lt;module&gt;
     11 
     12 arg_0 = array.clone()
---&gt; 13 jax.value_and_grad(fn, (0))(arg_0)

    [... skipping hidden 37 frame]

&lt;__array_function__ internals&gt; in delete(*args, **kwargs)

/usr/local/lib/python3.9/site-packages/numpy/lib/function_base.py in delete(arr, obj, axis)
   4470         if obj.dtype == bool:
   4471             if obj.shape != (N,):
-&gt; 4472                 raise ValueError('boolean array argument obj to delete '
   4473                                  'must be one dimensional and match the axis '
   4474                                  'length of {}'.format(N))

ValueError: boolean array argument obj to delete must be one dimensional and match the axis length of 3"><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
JaxStackTraceBeforeTransformation         Traceback (most recent call last)
/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/runpy.py in _run_module_as_main(***failed resolving arguments***)
    196         sys.argv[0] = mod_spec.origin
--&gt; 197     return _run_code(code, main_globals, None,
    198                      "__main__", mod_spec)

/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/runpy.py in _run_code(***failed resolving arguments***)
     86                        __spec__ = mod_spec)
---&gt; 87     exec(code, run_globals)
     88     return run_globals

~/Library/Python/3.9/lib/python/site-packages/ipykernel_launcher.py in &lt;module&gt;
     15     from ipykernel import kernelapp as app
---&gt; 16     app.launch_new_instance()

~/Library/Python/3.9/lib/python/site-packages/traitlets/config/application.py in launch_instance(***failed resolving arguments***)
    844         app.initialize(argv)
--&gt; 845         app.start()
    846 

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelapp.py in start(***failed resolving arguments***)
    618             try:
--&gt; 619                 self.io_loop.start()
    620             except KeyboardInterrupt:

~/Library/Python/3.9/lib/python/site-packages/tornado/platform/asyncio.py in start(***failed resolving arguments***)
    198             asyncio.set_event_loop(self.asyncio_loop)
--&gt; 199             self.asyncio_loop.run_forever()
    200         finally:

/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/base_events.py in run_forever(***failed resolving arguments***)
    600             while True:
--&gt; 601                 self._run_once()
    602                 if self._stopping:

/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/base_events.py in _run_once(***failed resolving arguments***)
   1904             else:
-&gt; 1905                 handle._run()
   1906         handle = None  # Needed to break cycles when an exception occurs.

/usr/local/Cellar/python@3.9/3.9.12_1/Frameworks/Python.framework/Versions/3.9/lib/python3.9/asyncio/events.py in _run(***failed resolving arguments***)
     79         try:
---&gt; 80             self._context.run(self._callback, *self._args)
     81         except (SystemExit, KeyboardInterrupt):

~/Library/Python/3.9/lib/python/site-packages/tornado/ioloop.py in &lt;lambda&gt;(***failed resolving arguments***)
    687             future.add_done_callback(
--&gt; 688                 lambda f: self._run_callback(functools.partial(callback, future))
    689             )

~/Library/Python/3.9/lib/python/site-packages/tornado/ioloop.py in _run_callback(***failed resolving arguments***)
    740         try:
--&gt; 741             ret = callback()
    742             if ret is not None:

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in inner(***failed resolving arguments***)
    813                 f = None  # noqa: F841
--&gt; 814                 self.ctx_run(self.run)
    815 

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in run(***failed resolving arguments***)
    774                     else:
--&gt; 775                         yielded = self.gen.send(value)
    776 

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelbase.py in process_one(***failed resolving arguments***)
    357                 return None
--&gt; 358         yield gen.maybe_future(dispatch(*args))
    359 

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--&gt; 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelbase.py in dispatch_shell(***failed resolving arguments***)
    260             try:
--&gt; 261                 yield gen.maybe_future(handler(stream, idents, msg))
    262             except Exception:

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--&gt; 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelbase.py in execute_request(***failed resolving arguments***)
    535         reply_content = yield gen.maybe_future(
--&gt; 536             self.do_execute(
    537                 code, silent, store_history,

~/Library/Python/3.9/lib/python/site-packages/tornado/gen.py in wrapper(***failed resolving arguments***)
    233                 try:
--&gt; 234                     yielded = ctx_run(next, result)
    235                 except (StopIteration, Return) as e:

~/Library/Python/3.9/lib/python/site-packages/ipykernel/ipkernel.py in do_execute(***failed resolving arguments***)
    301                 # letting shell dispatch to loop runners
--&gt; 302                 res = shell.run_cell(code, store_history=store_history, silent=silent)
    303         finally:

~/Library/Python/3.9/lib/python/site-packages/ipykernel/zmqshell.py in run_cell(***failed resolving arguments***)
    538         self._last_traceback = None
--&gt; 539         return super(ZMQInteractiveShell, self).run_cell(*args, **kwargs)
    540 

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in run_cell(***failed resolving arguments***)
   2897         try:
-&gt; 2898             result = self._run_cell(
   2899                 raw_cell, store_history, silent, shell_futures)

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in _run_cell(***failed resolving arguments***)
   2943         try:
-&gt; 2944             return runner(coro)
   2945         except BaseException as e:

~/Library/Python/3.9/lib/python/site-packages/IPython/core/async_helpers.py in _pseudo_sync_runner(***failed resolving arguments***)
     67     try:
---&gt; 68         coro.send(None)
     69     except StopIteration as exc:

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in run_cell_async(***failed resolving arguments***)
   3168 
-&gt; 3169                 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
   3170                        interactivity=interactivity, compiler=compiler, result=result)

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in run_ast_nodes(***failed resolving arguments***)
   3360                         asy = compare(code)
-&gt; 3361                     if (await self.run_code(code, result,  async_=asy)):
   3362                         return True

~/Library/Python/3.9/lib/python/site-packages/IPython/core/interactiveshell.py in run_code(***failed resolving arguments***)
   3440                 else:
-&gt; 3441                     exec(code_obj, self.user_global_ns, self.user_ns)
   3442             finally:

&lt;ipython-input-16-d01aa5baed3f&gt; in &lt;module&gt;
     17     arg_0 = arg_0_array.clone()
---&gt; 18     results['res_2'], _ = jax.value_and_grad(fn, (0))(arg_0)
     19 except Exception as e:

&lt;ipython-input-16-d01aa5baed3f&gt; in fn(***failed resolving arguments***)
      5     returned = False
----&gt; 6     return jax.numpy.average(arg_0, arg_1, returned=returned).sum()
      7 try:

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in average(***failed resolving arguments***)
    284             returned=False):
--&gt; 285   return _average(a, _ensure_optional_axes(axis), weights, returned)
    286 

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in _average(***failed resolving arguments***)
    292   if weights is None: # Treat all weights as 1
--&gt; 293     avg = mean(a, axis=axis)
    294     if axis is None:

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in mean(***failed resolving arguments***)
    251          out=None, keepdims=False, *, where=None):
--&gt; 252   return _mean(a, _ensure_optional_axes(axis), dtype, out, keepdims,
    253                where=where)

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in _mean(***failed resolving arguments***)
    278   return lax.div(
--&gt; 279       sum(a, axis, dtype=dtype, keepdims=keepdims, where=where),
    280       lax.convert_element_type(normalizer, dtype))

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in sum(***failed resolving arguments***)
    164         out=None, keepdims=None, initial=None, where=None):
--&gt; 165   return _reduce_sum(a, axis=_ensure_optional_axes(axis), dtype=dtype, out=out,
    166                      keepdims=keepdims, initial=initial, where=where)

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in _reduce_sum(***failed resolving arguments***)
    156                 dtype=None, out=None, keepdims=None, initial=None, where=None):
--&gt; 157   return _reduction(a, "sum", np.sum, lax.add, 0,
    158                     bool_op=lax.bitwise_or, upcast_f16_for_computation=True,

/usr/local/lib/python3.9/site-packages/jax/_src/numpy/reductions.py in _reduction(***failed resolving arguments***)
     97   else:
---&gt; 98     result = lax.reduce(a, init_val, op, dims)
     99   if initial is not None:

JaxStackTraceBeforeTransformation: ValueError: boolean array argument obj to delete must be one dimensional and match the axis length of 3

The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.

--------------------

The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
&lt;ipython-input-25-7d77a6f1d189&gt; in &lt;module&gt;
     11 
     12 arg_0 = array.clone()
---&gt; 13 jax.value_and_grad(fn, (0))(arg_0)

    [... skipping hidden 37 frame]

&lt;__array_function__ internals&gt; in delete(*args, **kwargs)

/usr/local/lib/python3.9/site-packages/numpy/lib/function_base.py in delete(arr, obj, axis)
   4470         if obj.dtype == bool:
   4471             if obj.shape != (N,):
-&gt; 4472                 raise ValueError('boolean array argument obj to delete '
   4473                                  'must be one dimensional and match the axis '
   4474                                  'length of {}'.format(N))

ValueError: boolean array argument obj to delete must be one dimensional and match the axis length of 3
</code></pre></div>
<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>
<p dir="auto"><code class="notranslate">cummax, cummin, cumprod</code> will succeed when using <code class="notranslate">value_and_grads</code> but fail if directly call</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax

def fn(arg_0):
    axis = -1
    return jax.lax.cummax(arg_0, axis=axis).sum()

mykey = jax.random.PRNGKey(14830444)
array = jax.random.uniform(mykey, [2, 3], jax.numpy.float32, minval=0, maxval=128)

try:
    arg_0 = array.clone()
    res1 = fn(arg_0)
except Exception as e:
    print(e)

arg_0 = array.clone()
res2, _ = jax.value_and_grad(fn, (0))(arg_0)
print(res2)
# axis -1 is out of bounds for array of shape (2, 3)
# 676.8106"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-k">def</span> <span class="pl-en">fn</span>(<span class="pl-s1">arg_0</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">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">cummax</span>(<span class="pl-s1">arg_0</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-s1">axis</span>).<span class="pl-en">sum</span>()

<span class="pl-s1">mykey</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">14830444</span>)
<span class="pl-s1">array</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">mykey</span>, [<span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-s1">float32</span>, <span class="pl-s1">minval</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">maxval</span><span class="pl-c1">=</span><span class="pl-c1">128</span>)

<span class="pl-k">try</span>:
    <span class="pl-s1">arg_0</span> <span class="pl-c1">=</span> <span class="pl-s1">array</span>.<span class="pl-en">clone</span>()
    <span class="pl-s1">res1</span> <span class="pl-c1">=</span> <span class="pl-en">fn</span>(<span class="pl-s1">arg_0</span>)
<span class="pl-k">except</span> <span class="pl-v">Exception</span> <span class="pl-k">as</span> <span class="pl-s1">e</span>:
    <span class="pl-en">print</span>(<span class="pl-s1">e</span>)

<span class="pl-s1">arg_0</span> <span class="pl-c1">=</span> <span class="pl-s1">array</span>.<span class="pl-en">clone</span>()
<span class="pl-s1">res2</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">fn</span>, (<span class="pl-c1">0</span>))(<span class="pl-s1">arg_0</span>)
<span class="pl-en">print</span>(<span class="pl-s1">res2</span>)
<span class="pl-c"># axis -1 is out of bounds for array of shape (2, 3)</span>
<span class="pl-c"># 676.8106</span></pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax

def fn(arg_0):
    axis = -1
    return jax.lax.cummin(arg_0, axis=axis).sum()

mykey = jax.random.PRNGKey(14830444)
array = jax.random.uniform(mykey, [2, 3], jax.numpy.float32, minval=0, maxval=128)

try:
    arg_0 = array.clone()
    res1 = fn(arg_0)
except Exception as e:
    print(e)

arg_0 = array.clone()
res2, _ = jax.value_and_grad(fn, (0))(arg_0)
print(res2)
# axis -1 is out of bounds for array of shape (2, 3)
# 549.35925"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-k">def</span> <span class="pl-en">fn</span>(<span class="pl-s1">arg_0</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">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">cummin</span>(<span class="pl-s1">arg_0</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-s1">axis</span>).<span class="pl-en">sum</span>()

<span class="pl-s1">mykey</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">14830444</span>)
<span class="pl-s1">array</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">mykey</span>, [<span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-s1">float32</span>, <span class="pl-s1">minval</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">maxval</span><span class="pl-c1">=</span><span class="pl-c1">128</span>)

<span class="pl-k">try</span>:
    <span class="pl-s1">arg_0</span> <span class="pl-c1">=</span> <span class="pl-s1">array</span>.<span class="pl-en">clone</span>()
    <span class="pl-s1">res1</span> <span class="pl-c1">=</span> <span class="pl-en">fn</span>(<span class="pl-s1">arg_0</span>)
<span class="pl-k">except</span> <span class="pl-v">Exception</span> <span class="pl-k">as</span> <span class="pl-s1">e</span>:
    <span class="pl-en">print</span>(<span class="pl-s1">e</span>)

<span class="pl-s1">arg_0</span> <span class="pl-c1">=</span> <span class="pl-s1">array</span>.<span class="pl-en">clone</span>()
<span class="pl-s1">res2</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">fn</span>, (<span class="pl-c1">0</span>))(<span class="pl-s1">arg_0</span>)
<span class="pl-en">print</span>(<span class="pl-s1">res2</span>)
<span class="pl-c"># axis -1 is out of bounds for array of shape (2, 3)</span>
<span class="pl-c"># 549.35925</span></pre></div>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax

def fn(arg_0):
    axis = -1
    return jax.lax.cumprod(arg_0, axis=axis).sum()

mykey = jax.random.PRNGKey(14830444)
array = jax.random.uniform(mykey, [2, 3], jax.numpy.float32, minval=0, maxval=128)

try:
    arg_0 = array.clone()
    res1 = fn(arg_0)
except Exception as e:
    print(e)

arg_0 = array.clone()
res2, _ = jax.value_and_grad(fn, (0))(arg_0)
print(res2)
# axis -1 is out of bounds for array of shape (2, 3)
# 1533018.6"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>

<span class="pl-k">def</span> <span class="pl-en">fn</span>(<span class="pl-s1">arg_0</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">return</span> <span class="pl-s1">jax</span>.<span class="pl-s1">lax</span>.<span class="pl-en">cumprod</span>(<span class="pl-s1">arg_0</span>, <span class="pl-s1">axis</span><span class="pl-c1">=</span><span class="pl-s1">axis</span>).<span class="pl-en">sum</span>()

<span class="pl-s1">mykey</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-v">PRNGKey</span>(<span class="pl-c1">14830444</span>)
<span class="pl-s1">array</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-s1">random</span>.<span class="pl-en">uniform</span>(<span class="pl-s1">mykey</span>, [<span class="pl-c1">2</span>, <span class="pl-c1">3</span>], <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-s1">float32</span>, <span class="pl-s1">minval</span><span class="pl-c1">=</span><span class="pl-c1">0</span>, <span class="pl-s1">maxval</span><span class="pl-c1">=</span><span class="pl-c1">128</span>)

<span class="pl-k">try</span>:
    <span class="pl-s1">arg_0</span> <span class="pl-c1">=</span> <span class="pl-s1">array</span>.<span class="pl-en">clone</span>()
    <span class="pl-s1">res1</span> <span class="pl-c1">=</span> <span class="pl-en">fn</span>(<span class="pl-s1">arg_0</span>)
<span class="pl-k">except</span> <span class="pl-v">Exception</span> <span class="pl-k">as</span> <span class="pl-s1">e</span>:
    <span class="pl-en">print</span>(<span class="pl-s1">e</span>)

<span class="pl-s1">arg_0</span> <span class="pl-c1">=</span> <span class="pl-s1">array</span>.<span class="pl-en">clone</span>()
<span class="pl-s1">res2</span>, <span class="pl-s1">_</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">value_and_grad</span>(<span class="pl-s1">fn</span>, (<span class="pl-c1">0</span>))(<span class="pl-s1">arg_0</span>)
<span class="pl-en">print</span>(<span class="pl-s1">res2</span>)
<span class="pl-c"># axis -1 is out of bounds for array of shape (2, 3)</span>
<span class="pl-c"># 1533018.6</span></pre></div>
<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="---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
&lt;ipython-input-9-b034345b2caa&gt; in &lt;module&gt;
     15 
     16 arg_0 = array.clone()
---&gt; 17 res1 = fn(arg_0)

&lt;ipython-input-9-b034345b2caa&gt; in fn(arg_0)
      3 def fn(arg_0):
      4     axis = -1
----&gt; 5     return jax.lax.cummin(arg_0, axis=axis).sum()
      6 
      7 mykey = jax.random.PRNGKey(14830444)

/usr/local/lib/python3.9/site-packages/jax/_src/lax/control_flow.py in cummin(operand, axis, reverse)
   2869 def cummin(operand: Array, axis: int = 0, reverse: bool = False) -&gt; Array:
   2870   &quot;&quot;&quot;Computes a cumulative minimum along `axis`.&quot;&quot;&quot;
-&gt; 2871   return cummin_p.bind(operand, axis=int(axis), reverse=bool(reverse))
   2872 
   2873 def _cumred_shape_rule(x, *, axis: int, reverse: bool):

/usr/local/lib/python3.9/site-packages/jax/core.py in bind(self, *args, **params)
    321     assert (not config.jax_enable_checks or
    322             all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--&gt; 323     return self.bind_with_trace(find_top_trace(args), args, params)
    324 
    325   def bind_with_trace(self, trace, args, params):

/usr/local/lib/python3.9/site-packages/jax/core.py in bind_with_trace(self, trace, args, params)
    324 
    325   def bind_with_trace(self, trace, args, params):
--&gt; 326     out = trace.process_primitive(self, map(trace.full_raise, args), params)
    327     return map(full_lower, out) if self.multiple_results else full_lower(out)
    328 

/usr/local/lib/python3.9/site-packages/jax/core.py in process_primitive(self, primitive, tracers, params)
    673 
    674   def process_primitive(self, primitive, tracers, params):
--&gt; 675     return primitive.impl(*tracers, **params)
    676 
    677   def process_call(self, primitive, f, tracers, params):

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in apply_primitive(prim, *args, **params)
     96 def apply_primitive(prim, *args, **params):
     97   &quot;&quot;&quot;Impl rule that compiles and runs a single primitive 'prim' using XLA.&quot;&quot;&quot;
---&gt; 98   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args),
     99                                         **params)
    100   return compiled_fun(*args)

/usr/local/lib/python3.9/site-packages/jax/_src/util.py in wrapper(*args, **kwargs)
    217         return f(*args, **kwargs)
    218       else:
--&gt; 219         return cached(config._trace_context(), *args, **kwargs)
    220 
    221     wrapper.cache_clear = cached.cache_clear

/usr/local/lib/python3.9/site-packages/jax/_src/util.py in cached(_, *args, **kwargs)
    210     @functools.lru_cache(max_size)
    211     def cached(_, *args, **kwargs):
--&gt; 212       return f(*args, **kwargs)
    213 
    214     @functools.wraps(f)

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in xla_primitive_callable(prim, *arg_specs, **params)
    146     else:
    147       return out,
--&gt; 148   compiled = _xla_callable_uncached(lu.wrap_init(prim_fun), device, None,
    149                                     prim.name, donated_invars, False, *arg_specs)
    150   if not prim.multiple_results:

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in _xla_callable_uncached(fun, device, backend, name, donated_invars, keep_unused, *arg_specs)
    228 def _xla_callable_uncached(fun: lu.WrappedFun, device, backend, name,
    229                            donated_invars, keep_unused, *arg_specs):
--&gt; 230   return lower_xla_callable(fun, device, backend, name, donated_invars, False,
    231                             keep_unused, *arg_specs).compile().unsafe_call
    232 

/usr/local/lib/python3.9/site-packages/jax/_src/profiler.py in wrapper(*args, **kwargs)
    204   def wrapper(*args, **kwargs):
    205     with TraceAnnotation(name, **decorator_kwargs):
--&gt; 206       return func(*args, **kwargs)
    207     return wrapper
    208   return wrapper

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in lower_xla_callable(fun, device, backend, name, donated_invars, always_lower, keep_unused, *arg_specs)
    270   with log_elapsed_time(f&quot;Finished tracing + transforming {fun.__name__} &quot;
    271                         &quot;for jit in {elapsed_time} sec&quot;):
--&gt; 272     jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
    273         fun, abstract_args, pe.debug_info_final(fun, &quot;jit&quot;), which_explicit)
    274   if any(isinstance(c, core.Tracer) for c in consts):

/usr/local/lib/python3.9/site-packages/jax/_src/profiler.py in wrapper(*args, **kwargs)
    204   def wrapper(*args, **kwargs):
    205     with TraceAnnotation(name, **decorator_kwargs):
--&gt; 206       return func(*args, **kwargs)
    207     return wrapper
    208   return wrapper

/usr/local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_final(fun, in_avals, debug_info, keep_inputs)
   1891     main.jaxpr_stack = ()  # type: ignore
   1892     with core.new_sublevel():
-&gt; 1893       jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
   1894         fun, main, in_avals, keep_inputs=keep_inputs)
   1895     del fun, main

/usr/local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals, keep_inputs)
   1863     in_tracers = _input_type_to_tracers(trace, in_avals)
   1864     in_tracers_ = [t for t, keep in zip(in_tracers, keep_inputs) if keep]
-&gt; 1865     ans = fun.call_wrapped(*in_tracers_)
   1866     out_tracers = map(trace.full_raise, ans)
   1867     jaxpr, consts = frame.to_jaxpr(out_tracers)

/usr/local/lib/python3.9/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    166 
    167     try:
--&gt; 168       ans = self.f(*args, **dict(self.params, **kwargs))
    169     except:
    170       # Some transformations yield from inside context managers, so we have to

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in prim_fun(*args)
    141   device = _device_from_arg_devices(arg_devices)
    142   def prim_fun(*args):
--&gt; 143     out = prim.bind(*args, **params)
    144     if prim.multiple_results:
    145       return out

/usr/local/lib/python3.9/site-packages/jax/core.py in bind(self, *args, **params)
    321     assert (not config.jax_enable_checks or
    322             all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--&gt; 323     return self.bind_with_trace(find_top_trace(args), args, params)
    324 
    325   def bind_with_trace(self, trace, args, params):

/usr/local/lib/python3.9/site-packages/jax/core.py in bind_with_trace(self, trace, args, params)
    324 
    325   def bind_with_trace(self, trace, args, params):
--&gt; 326     out = trace.process_primitive(self, map(trace.full_raise, args), params)
    327     return map(full_lower, out) if self.multiple_results else full_lower(out)
    328 

/usr/local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in process_primitive(self, primitive, tracers, params)
   1558     if primitive in custom_staging_rules:
   1559       return custom_staging_rules[primitive](self, *tracers, **params)
-&gt; 1560     return self.default_process_primitive(primitive, tracers, params)
   1561 
   1562   def default_process_primitive(self, primitive, tracers, params):

/usr/local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in default_process_primitive(self, primitive, tracers, params)
   1562   def default_process_primitive(self, primitive, tracers, params):
   1563     avals = [t.aval for t in tracers]
-&gt; 1564     out_avals, effects = primitive.abstract_eval(*avals, **params)
   1565     out_avals = [out_avals] if not primitive.multiple_results else out_avals
   1566     source_info = source_info_util.current()

/usr/local/lib/python3.9/site-packages/jax/core.py in abstract_eval_(*args, **kwargs)
    357 def _effect_free_abstract_eval(abstract_eval):
    358   def abstract_eval_(*args, **kwargs):
--&gt; 359     return abstract_eval(*args, **kwargs), no_effects
    360   return abstract_eval_
    361 

/usr/local/lib/python3.9/site-packages/jax/_src/lax/utils.py in standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule, named_shape_rule, *avals, **kwargs)
     64     return core.ConcreteArray(out.dtype, out, weak_type=weak_type)
     65   elif least_specialized is core.ShapedArray:
---&gt; 66     return core.ShapedArray(shape_rule(*avals, **kwargs),
     67                             dtype_rule(*avals, **kwargs), weak_type=weak_type,
     68                             named_shape=named_shape_rule(*avals, **kwargs))

/usr/local/lib/python3.9/site-packages/jax/_src/lax/control_flow.py in _cumred_shape_rule(x, axis, reverse)
   2873 def _cumred_shape_rule(x, *, axis: int, reverse: bool):
   2874   if axis &lt; 0 or axis &gt;= x.ndim:
-&gt; 2875     raise ValueError(
   2876         &quot;axis {} is out of bounds for array of shape {}&quot;.format(axis, x.shape))
   2877   return x.shape

ValueError: axis -1 is out of bounds for array of shape (2, 3)"><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
&lt;ipython-input-9-b034345b2caa&gt; in &lt;module&gt;
     15 
     16 arg_0 = array.clone()
---&gt; 17 res1 = fn(arg_0)

&lt;ipython-input-9-b034345b2caa&gt; in fn(arg_0)
      3 def fn(arg_0):
      4     axis = -1
----&gt; 5     return jax.lax.cummin(arg_0, axis=axis).sum()
      6 
      7 mykey = jax.random.PRNGKey(14830444)

/usr/local/lib/python3.9/site-packages/jax/_src/lax/control_flow.py in cummin(operand, axis, reverse)
   2869 def cummin(operand: Array, axis: int = 0, reverse: bool = False) -&gt; Array:
   2870   """Computes a cumulative minimum along `axis`."""
-&gt; 2871   return cummin_p.bind(operand, axis=int(axis), reverse=bool(reverse))
   2872 
   2873 def _cumred_shape_rule(x, *, axis: int, reverse: bool):

/usr/local/lib/python3.9/site-packages/jax/core.py in bind(self, *args, **params)
    321     assert (not config.jax_enable_checks or
    322             all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--&gt; 323     return self.bind_with_trace(find_top_trace(args), args, params)
    324 
    325   def bind_with_trace(self, trace, args, params):

/usr/local/lib/python3.9/site-packages/jax/core.py in bind_with_trace(self, trace, args, params)
    324 
    325   def bind_with_trace(self, trace, args, params):
--&gt; 326     out = trace.process_primitive(self, map(trace.full_raise, args), params)
    327     return map(full_lower, out) if self.multiple_results else full_lower(out)
    328 

/usr/local/lib/python3.9/site-packages/jax/core.py in process_primitive(self, primitive, tracers, params)
    673 
    674   def process_primitive(self, primitive, tracers, params):
--&gt; 675     return primitive.impl(*tracers, **params)
    676 
    677   def process_call(self, primitive, f, tracers, params):

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in apply_primitive(prim, *args, **params)
     96 def apply_primitive(prim, *args, **params):
     97   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
---&gt; 98   compiled_fun = xla_primitive_callable(prim, *unsafe_map(arg_spec, args),
     99                                         **params)
    100   return compiled_fun(*args)

/usr/local/lib/python3.9/site-packages/jax/_src/util.py in wrapper(*args, **kwargs)
    217         return f(*args, **kwargs)
    218       else:
--&gt; 219         return cached(config._trace_context(), *args, **kwargs)
    220 
    221     wrapper.cache_clear = cached.cache_clear

/usr/local/lib/python3.9/site-packages/jax/_src/util.py in cached(_, *args, **kwargs)
    210     @functools.lru_cache(max_size)
    211     def cached(_, *args, **kwargs):
--&gt; 212       return f(*args, **kwargs)
    213 
    214     @functools.wraps(f)

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in xla_primitive_callable(prim, *arg_specs, **params)
    146     else:
    147       return out,
--&gt; 148   compiled = _xla_callable_uncached(lu.wrap_init(prim_fun), device, None,
    149                                     prim.name, donated_invars, False, *arg_specs)
    150   if not prim.multiple_results:

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in _xla_callable_uncached(fun, device, backend, name, donated_invars, keep_unused, *arg_specs)
    228 def _xla_callable_uncached(fun: lu.WrappedFun, device, backend, name,
    229                            donated_invars, keep_unused, *arg_specs):
--&gt; 230   return lower_xla_callable(fun, device, backend, name, donated_invars, False,
    231                             keep_unused, *arg_specs).compile().unsafe_call
    232 

/usr/local/lib/python3.9/site-packages/jax/_src/profiler.py in wrapper(*args, **kwargs)
    204   def wrapper(*args, **kwargs):
    205     with TraceAnnotation(name, **decorator_kwargs):
--&gt; 206       return func(*args, **kwargs)
    207     return wrapper
    208   return wrapper

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in lower_xla_callable(fun, device, backend, name, donated_invars, always_lower, keep_unused, *arg_specs)
    270   with log_elapsed_time(f"Finished tracing + transforming {fun.__name__} "
    271                         "for jit in {elapsed_time} sec"):
--&gt; 272     jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
    273         fun, abstract_args, pe.debug_info_final(fun, "jit"), which_explicit)
    274   if any(isinstance(c, core.Tracer) for c in consts):

/usr/local/lib/python3.9/site-packages/jax/_src/profiler.py in wrapper(*args, **kwargs)
    204   def wrapper(*args, **kwargs):
    205     with TraceAnnotation(name, **decorator_kwargs):
--&gt; 206       return func(*args, **kwargs)
    207     return wrapper
    208   return wrapper

/usr/local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr_final(fun, in_avals, debug_info, keep_inputs)
   1891     main.jaxpr_stack = ()  # type: ignore
   1892     with core.new_sublevel():
-&gt; 1893       jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
   1894         fun, main, in_avals, keep_inputs=keep_inputs)
   1895     del fun, main

/usr/local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in trace_to_subjaxpr_dynamic(fun, main, in_avals, keep_inputs)
   1863     in_tracers = _input_type_to_tracers(trace, in_avals)
   1864     in_tracers_ = [t for t, keep in zip(in_tracers, keep_inputs) if keep]
-&gt; 1865     ans = fun.call_wrapped(*in_tracers_)
   1866     out_tracers = map(trace.full_raise, ans)
   1867     jaxpr, consts = frame.to_jaxpr(out_tracers)

/usr/local/lib/python3.9/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    166 
    167     try:
--&gt; 168       ans = self.f(*args, **dict(self.params, **kwargs))
    169     except:
    170       # Some transformations yield from inside context managers, so we have to

/usr/local/lib/python3.9/site-packages/jax/_src/dispatch.py in prim_fun(*args)
    141   device = _device_from_arg_devices(arg_devices)
    142   def prim_fun(*args):
--&gt; 143     out = prim.bind(*args, **params)
    144     if prim.multiple_results:
    145       return out

/usr/local/lib/python3.9/site-packages/jax/core.py in bind(self, *args, **params)
    321     assert (not config.jax_enable_checks or
    322             all(isinstance(arg, Tracer) or valid_jaxtype(arg) for arg in args)), args
--&gt; 323     return self.bind_with_trace(find_top_trace(args), args, params)
    324 
    325   def bind_with_trace(self, trace, args, params):

/usr/local/lib/python3.9/site-packages/jax/core.py in bind_with_trace(self, trace, args, params)
    324 
    325   def bind_with_trace(self, trace, args, params):
--&gt; 326     out = trace.process_primitive(self, map(trace.full_raise, args), params)
    327     return map(full_lower, out) if self.multiple_results else full_lower(out)
    328 

/usr/local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in process_primitive(self, primitive, tracers, params)
   1558     if primitive in custom_staging_rules:
   1559       return custom_staging_rules[primitive](self, *tracers, **params)
-&gt; 1560     return self.default_process_primitive(primitive, tracers, params)
   1561 
   1562   def default_process_primitive(self, primitive, tracers, params):

/usr/local/lib/python3.9/site-packages/jax/interpreters/partial_eval.py in default_process_primitive(self, primitive, tracers, params)
   1562   def default_process_primitive(self, primitive, tracers, params):
   1563     avals = [t.aval for t in tracers]
-&gt; 1564     out_avals, effects = primitive.abstract_eval(*avals, **params)
   1565     out_avals = [out_avals] if not primitive.multiple_results else out_avals
   1566     source_info = source_info_util.current()

/usr/local/lib/python3.9/site-packages/jax/core.py in abstract_eval_(*args, **kwargs)
    357 def _effect_free_abstract_eval(abstract_eval):
    358   def abstract_eval_(*args, **kwargs):
--&gt; 359     return abstract_eval(*args, **kwargs), no_effects
    360   return abstract_eval_
    361 

/usr/local/lib/python3.9/site-packages/jax/_src/lax/utils.py in standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule, named_shape_rule, *avals, **kwargs)
     64     return core.ConcreteArray(out.dtype, out, weak_type=weak_type)
     65   elif least_specialized is core.ShapedArray:
---&gt; 66     return core.ShapedArray(shape_rule(*avals, **kwargs),
     67                             dtype_rule(*avals, **kwargs), weak_type=weak_type,
     68                             named_shape=named_shape_rule(*avals, **kwargs))

/usr/local/lib/python3.9/site-packages/jax/_src/lax/control_flow.py in _cumred_shape_rule(x, axis, reverse)
   2873 def _cumred_shape_rule(x, *, axis: int, reverse: bool):
   2874   if axis &lt; 0 or axis &gt;= x.ndim:
-&gt; 2875     raise ValueError(
   2876         "axis {} is out of bounds for array of shape {}".format(axis, x.shape))
   2877   return x.shape

ValueError: axis -1 is out of bounds for array of shape (2, 3)
</code></pre></div>
<p dir="auto">version: 0.3.13</p>