Fix doctest and unbreak head

[IvyZX]

• Add functionality to suppress certain absl.logging messages during doctest. This will not break or unbreak anything (because logs will not break doctest), but can help clean up the printouts.
• Fix the Orbax upgrade guide to unbreak head.
• Add default required arguments for ContextManager.__exit__ in tensorboard.py

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 0% with 1 lines in your changes are missing coverage. Please review.

Project coverage is 58.43%. Comparing base (0c7d256) to head (3b23dc8).
Report is 1 commits behind head on main.

Files	Patch %	Lines
flax/metrics/tensorboard.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #3753   +/-   ##
=======================================
  Coverage   58.43%   58.43%           
=======================================
  Files         102      102           
  Lines       12365    12365           
=======================================
  Hits         7225     7225           
  Misses       5140     5140           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[chiamp]

Thanks!