tmoran6

[Bug] ReversibleSelect causes error when training model

3

### Description `ReversibleSelect` seems to mess up JAX's backtracing/JIT compilation. In the code provided below, we define a simple (non reversible) model which (1) splits the inputs; (2) does something...