Issue 1:
Please:

Check for duplicate issues.

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

softmax will fail when using value_and_grads but succeed when calling it directly

import jax

def fn(arg_0):
return jax.nn.softmax(arg_0, axis=False).sum()

mykey = jax.random.PRNGKey(5386231)
array = jax.random.uniform(mykey, [10, 5], jax.numpy.float32, minval=-64, maxval=0)

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

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 2

If applicable, include full error messages/tracebacks.

---------------------------------------------------------------------------
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
--> 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)
---> 87 exec(code, run_globals)
88 return run_globals

~/Library/Python/3.9/lib/python/site-packages/ipykernel_launcher.py in <module>
15 from ipykernel import kernelapp as app
---> 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)
--> 845 app.start()
846

~/Library/Python/3.9/lib/python/site-packages/ipykernel/kernelapp.py in start(***failed resolving arguments***)
618 try:
--> 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)
--> 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:
--> 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:
-> 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:
---> 80 self._context.run(self._callback, *self._args)
81 except (SystemExit, KeyboardInterrupt):

~/Library/Python/3.9/lib/python/site-packages/tornado/ioloop.py in <lambda>(***failed resolving arguments***)
687 future.add_done_callback(
--> 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:
--> 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
--> 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:
--> 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
--> 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:
--> 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:
--> 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:
--> 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(
--> 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:
--> 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
--> 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
--> 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:
-> 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:
-> 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:
---> 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
-> 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)
-> 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:
-> 3441 exec(code_obj, self.user_global_ns, self.user_ns)
3442 finally:

<ipython-input-14-5a11ff1b18d5> in <module>
16 arg_0 = arg_0_array.clone()
---> 17 results['res_2'], _ = jax.value_and_grad(fn, (0))(arg_0)
18 except Exception as e:

<ipython-input-14-5a11ff1b18d5> in fn(***failed resolving arguments***)
4 axis = False
----> 5 return jax.nn.softmax(arg_0, axis=axis).sum()
6 try:

/usr/local/lib/python3.9/site-packages/jax/_src/nn/functions.py in softmax(***failed resolving arguments***)
333 unnormalized = jnp.exp(x - lax.stop_gradient(x_max))
--> 334 return unnormalized / jnp.sum(unnormalized, axis, where=where, keepdims=True)
335

/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):
--> 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):
--> 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:
---> 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 2

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)
<ipython-input-23-cba934583f87> in <module>
11
12 arg_0 = array.clone()
---> 13 jax.value_and_grad(fn, (0))(arg_0)

[... skipping hidden 23 frame]

<__array_function__ internals> 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,):
-> 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 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