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.

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

```python #1

Open
ifan0924 opened this issue Sep 15, 2024 · 0 comments
Open

```python #1

ifan0924 opened this issue Sep 15, 2024 · 0 comments

Comments

@ifan0924
Copy link
Owner

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.

Originally posted by @mattjj in google/jax#23637

@ifan0924 ifan0924 mentioned this issue Sep 17, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant