Repository: google-deepmind/dm-haiku

Issue:
"Encountered an unexpected tracer" when scanning inside transformed function
Hi there. I'm doing something that requires me to iterate over my batches of data and send them to a haiku module. I was experiencing very slow compilation, so based on [this comment](https://www.gitmemory.com/issue/google/jax/2242/586772365), I thought I would try converting my `for batch in data` to a `jax.lax.scan` call. However, I ran into an "unexpected tracer" error doing this, despite the fact that I'm not (directly?) modifying global state.

Minimal reproduction code:
```
import haiku as hk
import jax
import jax.numpy as jnp

def build_forward_fn():
    def forward_fn(_,
                   is_training = True):
        # Define a model
        m = hk.Linear(300)

        # Prepare some data to scan over
        data = jnp.arange(0, 256).astype('float32')
        data = jnp.expand_dims(data, axis=0)
        data = jnp.stack([data, data])
        data = jnp.expand_dims(data, axis=0)

        def scanner_fn(carry, batch):
            out = m(batch)
            return carry, out

        jax.lax.scan(scanner_fn, 0, data)

    return forward_fn

if __name__ == '__main__':
    forward_fn = build_forward_fn()
    forward_fn = hk.transform(forward_fn)
    forward_fn_init = jax.jit(forward_fn.init)
    rng = jax.random.PRNGKey(100)
    forward_fn_init(rng, 1)
```
The traceback:
```
Traceback (most recent call last):
  File "test.py", line 30, in <module>
    forward_fn_init(rng, 1)
  File "/venvpath/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 183, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/venvpath/lib/python3.8/site-packages/jax/_src/api.py", line 399, in cache_miss
    out_flat = xla.xla_call(
  File "/venvpath/lib/python3.8/site-packages/jax/core.py", line 1561, in bind
    return call_bind(self, fun, *args, **params)
  File "/venvpath/lib/python3.8/site-packages/jax/core.py", line 1552, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "/venvpath/lib/python3.8/site-packages/jax/core.py", line 1564, in process
    return trace.process_call(self, fun, tracers, params)
  File "/venvpath/lib/python3.8/site-packages/jax/core.py", line 607, in process_call
    return primitive.impl(f, *tracers, **params)
  File "/venvpath/lib/python3.8/site-packages/jax/interpreters/xla.py", line 607, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "/venvpath/lib/python3.8/site-packages/jax/linear_util.py", line 262, in memoized_fun
    ans = call(fun, *args)
  File "/venvpath/lib/python3.8/site-packages/jax/interpreters/xla.py", line 682, in _xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(
  File "/venvpath/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1285, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "/venvpath/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 1263, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "/venvpath/lib/python3.8/site-packages/jax/linear_util.py", line 179, in call_wrapped
    ans = gen.send(ans)
  File "/venvpath/lib/python3.8/site-packages/jax/core.py", line 1523, in process_env_traces
    outs = map(trace.full_raise, outs)
  File "/venvpath/lib/python3.8/site-packages/jax/_src/util.py", line 41, in safe_map
    return list(map(f, *args))
  File "/venvpath/lib/python3.8/site-packages/jax/core.py", line 375, in full_raise
    val._assert_live()
  File "/venvpath/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 933, in _assert_live
    raise core.escaped_tracer_error(self, None)
jax._src.traceback_util.UnfilteredStackTrace: jax.core.UnexpectedTracerError: Encountered an unexpected tracer. Perhaps this tracer escaped through global state from a previously traced function.
The functions being transformed should not save traced values to global state.
The tracer that caused this error was created on line /venvpath/lib/python3.8/site-packages/haiku/_src/base.py:368 (run_creators). The tracer has shape (300,) and dtype float32.
When the tracer was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:
/venvpath/lib/python3.8/site-packages/haiku/_src/module.py:428 (wrapped)
/venvpath/lib/python3.8/site-packages/haiku/_src/module.py:279 (run_interceptors)
/venvpath/lib/python3.8/site-packages/haiku/_src/basic.py:183 (__call__)
/venvpath/lib/python3.8/site-packages/haiku/_src/base.py:310 (get_parameter)
/venvpath/lib/python3.8/site-packages/haiku/_src/base.py:368 (run_creators)
The function being traced when the tracer leaked was scanner_fn at test.py:17 traced for scan.
To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "test.py", line 30, in <module>
    forward_fn_init(rng, 1)
  File "/venvpath/lib/python3.8/site-packages/jax/interpreters/partial_eval.py", line 933, in _assert_live
    raise core.escaped_tracer_error(self, None)
jax.core.UnexpectedTracerError: Encountered an unexpected tracer. Perhaps this tracer escaped through global state from a previously traced function.
The functions being transformed should not save traced values to global state.
The tracer that caused this error was created on line /venvpath/lib/python3.8/site-packages/haiku/_src/base.py:368 (run_creators). The tracer has shape (300,) and dtype float32.
When the tracer was created, the final 5 stack frames (most recent last) excluding JAX-internal frames were:
/venvpath/lib/python3.8/site-packages/haiku/_src/module.py:428 (wrapped)
/venvpath/lib/python3.8/site-packages/haiku/_src/module.py:279 (run_interceptors)
/venvpath/lib/python3.8/site-packages/haiku/_src/basic.py:183 (__call__)
/venvpath/lib/python3.8/site-packages/haiku/_src/base.py:310 (get_parameter)
/venvpath/lib/python3.8/site-packages/haiku/_src/base.py:368 (run_creators)
The function being traced when the tracer leaked was scanner_fn at test.py:17 traced for scan.
To catch the leak earlier, try setting the environment variable JAX_CHECK_TRACER_LEAKS or using the `jax.checking_leaks` context manager.
```

I'm opening this in the haiku repo rather than the JAX repo because the error specifically comes up when I add the haiku module into the scan function - I haven't yet been able to reproduce this without haiku code. I haven't been able to come up with a reason why what I'm doing here wouldn't be supported, but I'm also not yet that familiar with the internals, so if this isn't a bug and there's something I'm missing, please let me know.


Discussion:
Of course I tried to debug this for an hour, opened the issue, then a minute later found the fix. Changing `jax.lax.scan` to `hk.scan` resolves the issue. It would be nice if haiku could somehow intercept these kinds of errors and mention that using the wrapped transformations may be necessary.

Resolving patch:
diff --git a/haiku/_src/transform.py b/haiku/_src/transform.py
--- a/haiku/_src/transform.py
+++ b/haiku/_src/transform.py
@@ -21,6 +21,7 @@
 from haiku._src import base
 from haiku._src import data_structures
 from haiku._src import typing
+import jax
 
 # If you are forking replace this with `import haiku as hk`.
 hk = types.ModuleType("haiku")
@@ -330,6 +331,14 @@ def transform_with_state(f) -> TransformedWithState:
   """
   analytics.log_once("transform_with_state")
 
+  unexpected_tracer_hint = (
+      "An UnexpectedTracerError was raised while inside a Haiku transformed "
+      "function (see error above).\n"
+      "Hint: are you using a JAX transform or JAX control-flow function "
+      "(jax.vmap/jax.scan/...) inside a Haiku transform? You might want to use "
+      "the Haiku version of the transform instead (hk.vmap/hk.scan/...).\n"
+      "See https://dm-haiku.readthedocs.io/en/latest/notebooks/transforms.html "
+      "on why you can't use JAX transforms inside a Haiku module.")
   def init_fn(
       rng: Optional[Union[PRNGKey, int]],
       *args,
@@ -338,7 +347,10 @@ def init_fn(
     """Initializes your function collecting parameters and state."""
     rng = to_prng_sequence(rng, err_msg=INIT_RNG_ERROR)
     with base.new_context(rng=rng) as ctx:
-      f(*args, **kwargs)
+      try:
+        f(*args, **kwargs)
+      except jax.errors.UnexpectedTracerError as e:
+        raise jax.errors.UnexpectedTracerError(unexpected_tracer_hint) from e
     return ctx.collect_params(), ctx.collect_initial_state()
 
   def apply_fn(
@@ -354,7 +366,10 @@ def apply_fn(
     rng = to_prng_sequence(
         rng, err_msg=(APPLY_RNG_STATE_ERROR if state else APPLY_RNG_ERROR))
     with base.new_context(params=params, state=state, rng=rng) as ctx:
-      out = f(*args, **kwargs)
+      try:
+        out = f(*args, **kwargs)
+      except jax.errors.UnexpectedTracerError as e:
+        raise jax.errors.UnexpectedTracerError(unexpected_tracer_hint) from e
     return out, ctx.collect_state()
 
   tie_in_original_fn(f, init_fn, apply_fn)


Test patch:
diff --git a/haiku/_src/transform_test.py b/haiku/_src/transform_test.py
--- a/haiku/_src/transform_test.py
+++ b/haiku/_src/transform_test.py
@@ -502,6 +502,16 @@ def f():
     self.assertEqual(type(state_in["~"]), cls)
     self.assertEqual(type(state_out["~"]), cls)
 
+  def test_unexpected_tracer_error_hint(self):
+    def leaks_and_uses_tracer():
+      jax.jit(base.next_rng_key)()
+      base.next_rng_key()
+
+    init, _ = transform.transform(leaks_and_uses_tracer)
+    with self.assertRaisesRegex(jax.errors.UnexpectedTracerError,
+                                "want to use the Haiku version"):
+      init(jax.random.PRNGKey(42))
+
 
 class ObjectWithTransform: