handle test failure in diffrax under new remat
handle test failure in diffrax under new remat

This change is to unblock landing new remat. I don't fully understand how, but
diffrax's bounded_while_loop (a tricky bit of code!) is hitting a condition I
thought was impossible. So let's handle that case!

I didn't add a test for it because I don't want to encode it as supported
behavior, and because handling it is actually simple.

Really I just want to unblock new-remat-by-default...

google/jax,google/jax,google/jax,google/jax