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

[mattjj]

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.