Switch to JAX's newer name-stack based named_call implementation by default.
levskaya opened this issue · comments
#2063 added initial testing support for using JAX's new dynamic name-stack approach to implementing named_call for labelled profiles.
We'd like to switch this approach to the default and remove the FLAX_PROFILE environment variable config in the near future.
This is a tracking bug for that issue.
#2349 finishes this - we've kept FLAX_PROFILE as a config option, though it defaults to true now. This was left as an option to help decrease the size of serialized HLOS.