[Re-submit] Compute true loss Flax examples (#19504)
* Compute true loss
* fixup
* final
* final
* final
* Update examples/flax/language-modeling/run_bart_dlm_flax.py
Co-authored-by: Sanchit Gandhi <93869735+sanchit-gandhi@users.noreply.github.com>
* jax.tree_map => jax.tree_util.tree_map
* Compute true loss
* final
* fixup
* final
* final
* Update examples/flax/language-modeling/run_bart_dlm_flax.py
Co-authored-by: Sanchit Gandhi <93869735+sanchit-gandhi@users.noreply.github.com>
* jax.tree_map => jax.tree_util.tree_map
Co-authored-by: Sanchit Gandhi <93869735+sanchit-gandhi@users.noreply.github.com>