Add a better error message for grad of while not supported

[gnecula]

The issue that I wanted to fix was that when running grad(lax.while_loop),
the error was a cryptic assertion failure (that all primals are known
after linearization, in ad.py:linearize). I could not figure out
how to detect before that assertion that we are doing a reverse AD
for while_loop. If I turn off that assertion the code fails a bit further down.

So, I implemented a simple form of partial evaluation, to allow the primals after linearization to be known, so that the code proceeds and can then fail gracefully when trying to transpose the while.

This is not a proper implementation of partial evaluation. The known outputs are computed early, properly. But the unknown outputs are computed by a whole computation of the loop, including the known parts. There are no residuals saved from the known computation. Perhaps one can do better, but it is not obvious. In any case, this code is only useful for a better error message :-(

Fixes issue: #2129

[mattjj]

I am behind in getting to this because I'm still trying to fix bugs from #2026. Sorry! I hope to get to it sometime this week. LMK if that doesn't work.

[gnecula]

That will work. Another one to look at is #2455, cjfj@ is waiting on that one also.

[gnecula]

@mattjj and @dougalm: I'd like your thoughts on this.

[mattjj]

Only after looking at and thinking about this change, I had an idea for an alternative, simpler way to get a good error message here.

There's a tag on JaxprTraces (from #1848) that indicates whether we're doing partial evaluation for the purpose of staging things out (in which case we likely want to stage out the whole while_loop primitive, see #2638) or for the purpose of reverse-mode autodiff. I think we can use that tag to raise the error we want here.

Concretely, I think we can do this:

def _while_loop_partial_eval(trace, *tracers, cond_nconsts, cond_jaxpr, body_nconsts, body_jaxpr):
  if trace.master.trace_type is pe.StagingJaxprTrace:
    params = dict(cond_nconsts=cond_nconsts, cond_jaxpr=cond_jaxpr, body_nconsts=body_nconsts, body_jaxpr=body_jaxpr)
    return trace.default_process_primitive(while_p, tracers, params)
  else:
    raise ValueError("Reverse-mode differentiation does not work for lax.while_loop.")

I haven't tried this yet though...

EDIT: actually, this precise idea won't work: I forgot that control flow primitives trace without that tag. So, the tag right now indicates whether the JaxprTrace is for (1) reverse-mode autodiff or control flow, versus (2) jit/pmap. We need to further distinguish between the first two sub-cases. Still, I think this tag-based direction, where we just tag JaxprTraces based on their purpos and check the tag, is probably a good one because it's simpler... WDYT?

[mattjj]

Nice! How about replacing "abstract" with "unknown"? Also perhaps remove "By default".

[gnecula]

We need to further distinguish between the first two sub-cases. Still, I think this tag-based direction, where we just tag JaxprTraces based on their purpos and check the tag, is probably a good one because it's simpler... WDYT?

I'd love this to work, this is what I was looking for. But how could it be enough to check the master trace type, don't we have to somehow inspect the entire trace stack?

[mattjj]

We'll end up checking all the partial eval traces on the trace stack because for each one we'll call into the _while_loop_partial_eval rule. If any one of them fails, i.e. if any one of them finds it's responsible for a reverse-mode autodiff partial eval, then we'll get an exception (but if all of them are instead either for outer control flow primitives or for jit/pmap staging then it'll go through).

We can actually attach any info to the master trace, like we can also just set trace.master.some_property = True anywhere. We don't need to do it via a type tag on master.trace_type. That just seemed like the smallest change at the time of #1848. But let's put our heads together and think about how best to attach this metadata for JaxprTraces!

Should we start with an enum that just indicates whether it's for reverse mode autodiff, control flow, or staging out of the system (i.e. jit/pmap)?

[gnecula]

I have tried something along the lines you suggested, but I think I still don't understand.

First, please check the comment that I wrote at the top of JaxprTrace. It is a refinement of a comment you had in JaxprTrace.process_custom_jvp_calll, adding the case for linearization.

I decided to try for now a simple change where we tag the linearization case, to separate it out of the
handling of control-flow. It seems easier to tag the former, because control-flow handling is quite
pervasive. I added LinearizeJaxprTrace, a sibling of StagingJaxprTrace.

The problem is with grad(cond(while)). By the time I get to partial eval of the nested while, the master that was created for linearize is not in the current trace.

The only case that does not work

[gnecula]

BTW, if I scan core.trace_stack I do see the LinearizeJaxprTrace. Can't a MasterTrace have parents?

[mattjj]

Hmm, what's a parent in this context?

[mattjj]

The comments LGTM!

[mattjj]

Tricky... I think the issue is that when we do partial evaluation of control flow primitives, with things like partial_eval_jaxpr, we need to propagate this tag information. That is, we're kicking off a new master when we call partial_eval_jaxpr inside the partial eval rule of cond; the fact that the cond partial eval is happening because of linearization is not being forwarded through partial_eval_jaxpr.

[mattjj]

Yes, that's the issue. The test passes after

1. adding a trace_type argument to partial_eval_jaxpr (or we could do the for_linearize / stage_out signature, though maybe a single argument is easier)
2. changing partial_eval_jaxpr by adding the line for_linearize = trace_type is LinearizeJaxprTrace at the top, then making the two calls to trace_to_jaxpr in its body pass that value

These initial style control flow primitives are the trickiest, at least the way we process them by round-tripping jaxprs back into traceables via jaxpr_as_fun/eval_jaxpr and tracing them. That's caused me some headaches earlier today with the PRNG reuse checker I'm prototyping. Maybe we should brainstorm how we might process jaxprs without involving eval_jaxpr. It could be as simple as writing a couple separate jaxpr interpreters. (In effect right now we are reusing the same eval_jaxpr interpreter, but transforming it with our usual Python transforms.) To be honest, though, I'm not sure if that would actually help much; it's just an alternative that would eliminate one level of tracing.

[gnecula]

Hmm, what's a parent in this context?

A better name is previous: the trace that was on top of the stack when you pushed this one. In essence, this would allow one to scan the trace stack from any particular tracer.

[gnecula]

I have implemented your suggestion. The one difference is that instead of passing for_linearize flag to trace_to_jaxpr I just pass trace_type. This now subsumes stage_out, but I left that alone.

BTW, the previous implementation had the minor advantage of actually implementing linearize for while_loop. Now we get the same error as for grad.

[gnecula]

Looks like random.gamma does not work, its VJP uses while (superficially, could be removed, see here) .

But the problem still stands, the current version give spurious errors if while_loop is in the VJP. The error is deserved if you want to do higher-order differentiation, but not for the first differentiation. Looks like I need to adjust this check further.

[mattjj]

A better name is previous: the trace that was on top of the stack when you pushed this one. In essence, this would allow one to scan the trace stack from any particular tracer.

Ah I see, thanks. Personally I try not to think in terms of the trace stack; every trace just has to worry about itself, and the core handles the fact that there are multiple traces going on.

BTW, the previous implementation had the minor advantage of actually implementing linearize for while_loop. Now we get the same error as for grad.

Oh, I didn't realize that. That would be a nice property to have. Should we revive the other solution then? Especially also given the gamma issue.

These refined "jaxpr trace type" tags are probably nice to keep around for other reasons, but making while_loop work with linearize is a big plus with your original approach.

[gnecula]

I have revived the partial evaluation of while_loop. I kept the passing of trace_type along, but I dropped the LinearizeJaxprTrace type tag.

I test that linearize works but grad errors gracefully.

[mattjj]

Awesome!

[cghawthorne]

Quick suggestion for the error message: I think it would be helpful if it suggested the user try lax.scan instead. This wasn't obvious to me when I ran into this same problem and had to ask for help from @jekbradbury