jax
ce8ecbd1 - Add an extension mechanism to run_state that allows:

Commit
1 year ago
Add an extension mechanism to run_state that allows: * Uninitialized values * Custom ref aval construction This will allow us to replace `run_scoped` with `run_state`, and allow us to change the memory space of initialized values. Co-authored-by: Sharad Vikram <sharadmv@google.com> PiperOrigin-RevId: 688965089
Author
Parents
Loading