Traceback doesn't say where the function breaks

[GeoffNN]

Hello! 😃

I just changed the input to a complex neural network training loop. I'm getting a strange error on the second call of a (complicated) jitted function:

Exception has occurred: ValueError       (note: full exception trace is shown but execution is paused at: _run_module_as_main)
The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
  File "/home/miniconda3/envs/jaxsel/lib/python3.10/site-packages/jax/_src/array.py", line 226, in __bool__
    return bool(self._value)
  File "/home/myhome/google-research/jaxsel/jaxsel/examples/train.py", line 471, in training_loop
    (loss, (preds, logits, q, dense_submat, label_loss, curiosity_loss, entropy_loss, exploration_bonus)), model_grad = value_grad_loss_fn(
  File "/home/myhome/google-research/jaxsel/jaxsel/examples/train.py", line 737, in main
    training_loop(

Unfortunately the traceback doesn't say where this happens in the code at all. Usually jax does a pretty good job giving the full traceback (even with traced arrays), but here I get nothing but the very last frame of the traceback, which doesn't help me at all. In particular, I wouldn't know where to put a jax.debug.print or breakpoint since I don't know where this happens :(
When I print self in the last frame of the Traceback, I get

print(self)
> Array([ True,  True], dtype=bool)
type(self)
> <class 'jaxlib.xla_extension.Array'>

There is no error when the function is not jitted.

Any help or pointers would be great
Thanks!

I'm using jax 0.4.2, and cudnn 8.2.

[hawkinsp]

Try setting

jax.config.update('jax_traceback_filtering', 'off')

Does that help?

(It might be a bug in the traceback filtering, since we should only filter JAX-internal frames.)

[GeoffNN]

Here's the full traceback, in case it helps:

Exception has occurred: ValueError       (note: full exception trace is shown but execution is paused at: _run_module_as_main)
The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
  File "/.../miniconda3/envs/jaxsel/lib/python3.10/site-packages/jax/_src/array.py", line 226, in __bool__
    return bool(self._value)
  File "/.../google-research/jaxsel/jaxsel/examples/train.py", line 473, in training_loop
    (loss, (preds, logits, q, dense_submat, label_loss, curiosity_loss, entropy_loss, exploration_bonus)), model_grad = value_grad_loss_fn(
  File "/.../google-research/jaxsel/jaxsel/examples/train.py", line 740, in main
    training_loop(
  File "/.../miniconda3/envs/jaxsel/lib/python3.10/site-packages/absl/app.py", line 254, in _run_main
    sys.exit(main(argv))
  File "/.../miniconda3/envs/jaxsel/lib/python3.10/site-packages/absl/app.py", line 308, in run
    _run_main(main, args)
  File "/.../google-research/jaxsel/jaxsel/examples/train.py", line 782, in <module>
    app.run(main)
  File "/.../miniconda3/envs/jaxsel/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/.../miniconda3/envs/jaxsel/lib/python3.10/runpy.py", line 196, in _run_module_as_main (Current frame)
    return _run_code(code, main_globals, None,
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

[hawkinsp]

I don't suppose you can share a repro?

[GeoffNN]

I can, but it's definitely not minimal -- let me try to minimize it a bit.

[GeoffNN]

The strange thing is that the non-jitted function is fine, including when using grad. It just breaks on the second call of the jitted function with this weird traceback.

[hawkinsp]

I think that means that the C++ fast path for jit is eating part of the traceback, but it would be helpful to have a repro of the problem so I know what to fix...