Problems with orbax export documentation

[reyammer]

I was trying to find info about how to export an orbax checkpoint to SavedModel, and found broken links or code examples in the docs with exception / stack traces:

• pypi page (https://pypi.org/project/orbax-export/) points to https://orbax.readthedocs.io/en/latest/orbax_export_101.html, which leads to 404

• https://orbax.readthedocs.io/en/latest/guides/export/orbax_export_101.html has code, but the "the output" part of the code just shows stacktraces / exceptions, e.g.,

ImportError                               Traceback (most recent call last)
Cell In[1], line 2
      1 # Import Orbax Export classes.
----> 2 from orbax.export import ExportManager
      3 from orbax.export import JaxModule
      4 from orbax.export import ServingConfig

[qdhack]

Thanks for raising the issue. We are working on a fix or rollback of the relevant proto pb2.

[andrewkchan]

Same issue here: https://orbax.readthedocs.io/en/latest/guides/checkpoint/transformations.html#restoring-a-checkpoint

Looks like the cells intended to contain outputs of various checkpoint restorations just show stuff like

---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
Cell In[11], line 1
----> 1 restored1

NameError: name 'restored1' is not defined