Issue 1:
Please:

Check for duplicate issues.

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

psum will fail when using value_and_grads but it will succeed if call it directly

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]).

If applicable, include full error messages/tracebacks.

---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-22-c6dfa3d16c05> in <module>
11
12 arg_0 = array.clone()
---> 13 jax.value_and_grad(fn, (0))(arg_0)

[... skipping hidden 8 frame]

<ipython-input-22-c6dfa3d16c05> in fn(arg_0)
2 results = dict()
3 def fn(arg_0):
----> 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 <= a < operand.ndim for a in axes):
-> 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]).

Issue 2:
Please:

Check for duplicate issues.

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

cummax, cummin, cumprod will succeed when using value_and_grads but fail if directly call

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

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

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

If applicable, include full error messages/tracebacks.

---------------------------------------------------------------------------
ValueError Traceback (most recent call last)
<ipython-input-9-b034345b2caa> in <module>
15
16 arg_0 = array.clone()
---> 17 res1 = fn(arg_0)

<ipython-input-9-b034345b2caa> in fn(arg_0)
3 def fn(arg_0):
4 axis = -1
----> 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) -> Array:
2870 """Computes a cumulative minimum along `axis`."""
-> 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
--> 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):
--> 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):
--> 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."""
---> 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:
--> 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):
--> 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,
--> 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):
--> 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):
--> 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"):
--> 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):
--> 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():
-> 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]
-> 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:
--> 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):
--> 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
--> 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):
--> 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)
-> 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]
-> 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):
--> 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:
---> 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 < 0 or axis >= x.ndim:
-> 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)

version: 0.3.13