Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

when a tracer error happens in for_loop, should point to the user's body function #23637

Closed
mattjj opened this issue Sep 13, 2024 · 0 comments · Fixed by #25798
Closed

when a tracer error happens in for_loop, should point to the user's body function #23637

mattjj opened this issue Sep 13, 2024 · 0 comments · Fixed by #25798
Assignees
@gnecula
Labels
better_errors Improve the error reporting

Comments

@mattjj
Copy link
Collaborator

mattjj commented Sep 13, 2024 

import jax
def body(c, _): return bool(c), None
jax.lax.scan(body, 3., None, length=1)
Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/sholto8.py", line 3, in <module>
    jax.lax.scan(body, 3., None, length=1)
  File "/usr/local/google/home/mattjj/packages/jax/sholto8.py", line 2, in body
    def body(c, _): return bool(c), None
                           ^^^^^^^
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape float32[].
The error occurred while tracing the function body at /usr/local/google/home/mattjj/packages/jax/sholto8.py:2 for scan. This concrete value was not available in Python because it depends on the value of the argument c.

Notice especially this part that points to the user's function body:

The error occurred while tracing the function body at /usr/local/google/home/mattjj/packages/jax/sholto8.py:2 for scan.

But if we do the analogous thing with fori_loop:

import jax
def body(_, c): return bool(c)
jax.lax.fori_loop(0, 5, body, 3.)
Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/sholto8.py", line 3, in <module>
    jax.lax.fori_loop(0, 5, body, 3.)
  File "/usr/local/google/home/mattjj/packages/jax/sholto8.py", line 2, in body
    def body(_, c): return bool(c)
                           ^^^^^^^
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape float32[].
The error occurred while tracing the function scanned_fun at /usr/local/google/home/mattjj/packages/jax/jax/_src/lax/control_flow/loops.py:1945 for scan. This concrete value was not available in Python because it depends on the value of the argument loop_carry[1].
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

Oops, we point to an internal function name rather than the user's:

The error occurred while tracing the function scanned_fun at /usr/local/google/home/mattjj/packages/jax/jax/_src/lax/control_flow/loops.py:1945 for scan

We should probably improve that to point to the user function name, like we do with scan, rather than our internal wrapper.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Assignees

@gnecula gnecula

Labels
better_errors Improve the error reporting
Projects
None yet
Milestone
No milestone
Development

Successfully merging a pull request may close this issue.

[better_errors] Ensure that tracer errors in for_loop point to use code gnecula/jax
2 participants
@mattjj @gnecula