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
• 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.outputs
    • ValidationReport.latency
    • ValidationReport.xprof_url
    • ValidationReport.metadata
    • ValidationReport.status
    • ValidationReport.__init__()
• ValidationReportOption
  • ValidationReportOption
    • ValidationReportOption.floating_atol
    • ValidationReportOption.floating_rtol
    • ValidationReportOption.max_non_floating_mismatch_ratio
    • ValidationReportOption.output_report_path
    • ValidationReportOption.print_debug_info
    • ValidationReportOption.__post_init__()
    • ValidationReportOption.__eq__()
    • ValidationReportOption.__hash__
    • ValidationReportOption.__init__()