`debug.print` inside `scan` gets lost during AD

[dfm]

Description

import jax

def f(c, _):
  jax.debug.print("c = {c}", c=c, ordered=True)
  return c + 1, None

def g(x):
  return jax.lax.scan(f, x, length=2)[0]

jax.make_jaxpr(jax.value_and_grad(g))(1.0)

prints

{ lambda ; a:f32[]. let
    b:f32[] = scan[
      _split_transpose=False
      jaxpr={ lambda ; c:f32[]. let d:f32[] = add c 1.0:f32[] in (d,) }
      length=2
      linear=(False,)
      num_carry=1
      num_consts=0
      reverse=False
      unroll=1
    ] a
    e:f32[] = scan[
      _split_transpose=False
      jaxpr={ lambda ; f:f32[]. let  in (f,) }
      length=2
      linear=(True,)
      num_carry=1
      num_consts=0
      reverse=True
      unroll=1
    ] 1.0:f32[]
  in (b, e) }

Where the debug.print has been dropped from the forward pass.

I narrowed it down to getting lost here:

jax/jax/_src/lax/control_flow/loops.py

Lines 739 to 743 in af66ca9

	jaxpr_known_, invar_pvals_out, jaxpr_known_consts = pe.trace_to_jaxpr_nounits(
	lu.wrap_init(core.jaxpr_as_fun(jaxpr_known),
	debug_info=jaxpr_known.jaxpr.debug_info),
	const_pvals + other_pvals,
	instantiate=[True] * (len(out_uk) - sum(out_uk)) + [False] * num_res)

Because tracers_to_jaxpr won't capture these eventful eqns.

AFAICT, this hasn't ever worked (I tried it back until 0.4.6 which is the oldest jaxlib on PyPI), but it should!

System info (python version, jaxlib version, accelerator, etc.)

...

[dfm]

@dougalm points out that this will probably be best solved by use_direct_linearize. It doesn't yet fix the problem because we don't have the custom linearization rule for scan implemented, so I'll work on that!

[mattjj]

It might work if you put it under a remat, even with a save-everything policy. (I’m not sure if the whole scan has to be under there or if just doing it on the body is enough.)

[dfm]

It does work in the sense that the print isn't dropped, but it does get moved to the remat step. Adding

@jax.remat
def f(c, _):
  ...

to the example above, prints

{ lambda ; a:f32[]. let
    b:i32[2] = iota[dimension=0 dtype=int32 shape=(2,) sharding=None] 
    c:f32[] d:f32[2] = scan[
      _split_transpose=False
      jaxpr={ lambda ; e:f32[] f:i32[]. let
          g:f32[] = add e 1.0:f32[]
        in (g, e) }
      length=2
      linear=(False, False)
      num_carry=1
      num_consts=0
      reverse=False
      unroll=1
    ] a b
    h:f32[] = scan[
      _split_transpose=False
      jaxpr={ lambda ; i:f32[] j:f32[]. let
          k:f32[] = remat2[
            differentiated=True
            jaxpr={ lambda ; l:f32[] m:f32[]. let
                debug_callback[
                  callback=<function debug_callback.<locals>._flat_callback at 0x1088d1800>
                  effect=OrderedDebug
                  partitioned=False
                ] l
              in (m,) }
            policy=None
            prevent_cse=True
          ] j i
        in (k,) }
      length=2
      linear=(True, False)
      num_carry=1
      num_consts=0
      reverse=True
      unroll=1
    ] 1.0:f32[] d
  in (c, h) }

[patrick-kidger]

Possibly related: #16637

[dfm]

@patrick-kidger — Absolutely! Thanks for the reference. This is exactly the same issue :D! #28748 also addresses that one when under jax_use_direct_linearize is set to True, but that's not the default yet.