Recommendation for saving auxilary data

Often times we compute auxiliary variables each timestep within vector_field that we would like to save as part of processing for future reference/sanity checking. We would like to avoid having to manually recalculate these within postprocessing (for confident debugging). What is the best practice way of achieving this? One option is to have a custom stepper where the auxilary data is not "stepped"/incremented but instead updated/replaced but doesn't feel ideal. Ideally, something like a has_aux like in the standard jax.vjp could be supported. Any other suggestions? Thanks 🙂

Start by examining how vector_field data is produced and how a custom stepper currently handles state. Compare the requested auxiliary-data behavior with the jax.vjp has_aux pattern. Done means there is an agreed design for retaining per-timestep auxiliary values without requiring users to recompute them in postprocessing.