orbax.export API Reference

• DTensor utilities for multi-device/host export
  • initialize_dtensor()
  • dtensor_initialized()
  • shutdown_dtensor()
  • jax_mesh_to_dtensor_mesh()
  • maybe_enable_dtensor_export_on()
  • get_current_dtensor_mesh()
• ExportManager
  • ExportManager
    • ExportManager
• ExportManagerBase
  • ExportManagerBase
    • ExportManagerBase
• JaxModule
  • JaxModule
    • JaxModule
• ServingConfig
  • ServingConfig
    • ServingConfig
• General Utilities
  • TensorSpecWithDefault
    • TensorSpecWithDefault
  • CallableSignatures
    • CallableSignatures
  • Utility functions
    • with_default_args()
    • make_auto_batching_function()
• ValidationManager
  • ValidationManager
    • ValidationManager
• ValidationReport
  • ValidationReport
    • ValidationReport.__init__()
• ValidationReportOption
  • ValidationReportOption
    • ValidationReportOption.__post_init__()
    • ValidationReportOption.__eq__()
    • ValidationReportOption.__hash__
    • ValidationReportOption.__init__()