<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">psum</code> will fail when using <code class="notranslate">value_and_grads</code> but it will succeed if call it directly</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.lax.psum(arg_0, [-1]).sum()

mykey = jax.random.PRNGKey(21847582)
array = jax.random.uniform(mykey, [2], jax.numpy.float32, minval=-8, maxval=64)

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

arg_0 = array.clone()
jax.value_and_grad(fn, (0))(arg_0)
# ValueError: reduction axes (-1,) contains out-of-bounds indices for ShapedArray(float32[2])."><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">lax</span>.<span class="pl-en">psum</span>(<span class="pl-s1">arg_0</span>, [<span class="pl-c1">-</span><span class="pl-c1">1</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">21847582</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-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">-</span><span class="pl-c1">8</span>, <span class="pl-s1">maxval</span><span class="pl-c1">=</span><span class="pl-c1">64</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"># 25.462376</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: reduction axes (-1,) contains out-of-bounds indices for ShapedArray(float32[2]).</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-22-c6dfa3d16c05&gt; in &lt;module&gt;
     11 
     12 arg_0 = array.clone()
---&gt; 13 jax.value_and_grad(fn, (0))(arg_0)

    [... skipping hidden 8 frame]

&lt;ipython-input-22-c6dfa3d16c05&gt; in fn(arg_0)
      2 results = dict()
      3 def fn(arg_0):
----&gt; 4     return jax.lax.psum(arg_0, [-1]).sum()
      5 
      6 mykey = jax.random.PRNGKey(21847582)

    [... skipping hidden 14 frame]

/usr/local/lib/python3.9/site-packages/jax/_src/lax/lax.py in _reduce_op_shape_rule(***failed resolving arguments***)
   3488     raise ValueError(f&quot;duplicate value in 'axes' of reduction: {axes}&quot;)
   3489   if not all(0 &lt;= a &lt; operand.ndim for a in axes):
-&gt; 3490     raise ValueError(f&quot;reduction axes {axes} contains out-of-bounds indices for {operand}.&quot;)
   3491   axes = frozenset(axes)
   3492   return tuple(d for i, d in enumerate(operand.shape) if i not in axes)

ValueError: reduction axes (-1,) contains out-of-bounds indices for ShapedArray(float32[2])."><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
&lt;ipython-input-22-c6dfa3d16c05&gt; in &lt;module&gt;
     11 
     12 arg_0 = array.clone()
---&gt; 13 jax.value_and_grad(fn, (0))(arg_0)

    [... skipping hidden 8 frame]

&lt;ipython-input-22-c6dfa3d16c05&gt; in fn(arg_0)
      2 results = dict()
      3 def fn(arg_0):
----&gt; 4     return jax.lax.psum(arg_0, [-1]).sum()
      5 
      6 mykey = jax.random.PRNGKey(21847582)

    [... skipping hidden 14 frame]

/usr/local/lib/python3.9/site-packages/jax/_src/lax/lax.py in _reduce_op_shape_rule(***failed resolving arguments***)
   3488     raise ValueError(f"duplicate value in 'axes' of reduction: {axes}")
   3489   if not all(0 &lt;= a &lt; operand.ndim for a in axes):
-&gt; 3490     raise ValueError(f"reduction axes {axes} contains out-of-bounds indices for {operand}.")
   3491   axes = frozenset(axes)
   3492   return tuple(d for i, d in enumerate(operand.shape) if i not in axes)

ValueError: reduction axes (-1,) contains out-of-bounds indices for ShapedArray(float32[2]).
</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>