skip remat test that fails with autodiff
The purpose of jax.checkpoint / jax.remat is to control autodiff behavior. But there is a flax test for remat, namely test_remat_kwargs, which would fail when used with autodiff. The issue is that when used with autodiff jax.remat will abstract its arguments, which doesn't happen when no autodiff is involved. (But for jax.remat to do anything, autodiff has to be involved!)
This PR skips the test and adds a line showing the failure.
We noticed this issue because the new jax.remat implementation errors sooner, namely without waiting for autodiff.
We might be able to adapt this test code. Interestingly, no Google code seems to follow the same pattern as this test: this is the only failure which crops up when switching on the new jax.remat implementation. So an updated version of this test should probably look a bit different.
Co-authored-by: Roy Frostig <frostig@google.com>
