Add `discharge_state` to convert stateful jaxprs into pure jaxprs
Add `discharge_state` to convert stateful jaxprs into pure jaxprs

Co-authored-by: Matthew Johnson <e68678f07030b810d1440f7f1c1c423f5afe56fe@google.com>

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