CWL now compatible with partially nondifferentiable body functions.

[patrick-kidger]

This is fairly finickity! So, the checkpointed while loop (CWL) needs
to figure out which cotangents to propagate backward. There are two
criteria for this:
(1) only some inputs actually need cotangents at all. Any parts of the
carry that are unrelated to those inputs can be skipped.
(2) some cotangents might be resolvable as symbolic zeros, and can thus
be skipped.

Previously, we handled (2) but did not handle (1). In practice this
was mostly fine. We just pretended all (inexact) inputs needed
cotangents, and let anything superfluous get DCE'd at the end.

However, this has one niche problem, which is that it is incompatible
with eqxi.nondifferentiable -- which is used to mark an inexact
input as nondifferentiable, and raise a trace-time error if we attempt
to differentiate it.

With this change, we're now a bit more careful: we run an additional
fixed-point iteration to check which inputs need cotangents, and then
skip the rest.

Note that this extra fixed-point iteration is actually done using a
JVP. This is designed to mirror what JAX does internally: perturbations
are tracked using a fixed-point iteration inside the JVP rule for
lax.{while_loop,scan}. I considered some other alternatives (using
VJPs or pe.dce_jaxpr) but I think those either fail in corner cases
or run into catch-22s.