{"id":"https://openalex.org/W7162329053","doi":"https://doi.org/10.48550/arxiv.2605.23066","title":"Orbax: Distributed Checkpointing with JAX","display_name":"Orbax: Distributed Checkpointing with JAX","publication_year":2026,"publication_date":"2026-05-21","ids":{"openalex":"https://openalex.org/W7162329053","doi":"https://doi.org/10.48550/arxiv.2605.23066"},"language":null,"primary_location":{"id":"doi:10.48550/arxiv.2605.23066","is_oa":true,"landing_page_url":"https://doi.org/10.48550/arxiv.2605.23066","pdf_url":null,"source":{"id":"https://openalex.org/S4306400194","display_name":"arXiv (Cornell University)","issn_l":null,"issn":null,"is_oa":true,"is_in_doaj":false,"is_core":false,"host_organization":"https://openalex.org/I205783295","host_organization_name":"Cornell University","host_organization_lineage":["https://openalex.org/I205783295"],"host_organization_lineage_names":[],"type":"repository"},"license":"cc-by","license_id":"https://openalex.org/licenses/cc-by","version":null,"is_accepted":false,"is_published":false,"raw_source_name":null,"raw_type":"Preprint"},"type":"preprint","indexed_in":["datacite"],"open_access":{"is_oa":true,"oa_status":"green","oa_url":"https://doi.org/10.48550/arxiv.2605.23066","any_repository_has_fulltext":true},"authorships":[{"author_position":"first","author":{"id":"https://openalex.org/A5034219676","display_name":"Colin Gaffney","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Gaffney, Colin","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5136917130","display_name":"Shutong Li","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Li, Shutong","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5136942682","display_name":"Daniel Ng","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Ng, Daniel","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5093547843","display_name":"Anastasia Petrushkina","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Petrushkina, Anastasia","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5113932716","display_name":"Niket Kumar","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Kumar, Niket","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5136968285","display_name":"Adam Cogdell","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Cogdell, Adam","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5136928925","display_name":"Mridul Sahu","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Sahu, Mridul","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5136949268","display_name":"Yaning Liang","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Liang, Yaning","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5008689083","display_name":"Nikhil Bansal","orcid":"https://orcid.org/0000-0002-6290-0894"},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Bansal, Nikhil","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5136914650","display_name":"Justin Pan","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Pan, Justin","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5136964184","display_name":"Angel Mau","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Mau, Angel","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5136988373","display_name":"Abhishek Agrawal","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Agrawal, Abhishek","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5053982664","display_name":"Marco Berlot","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Berlot, Marco","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5042201790","display_name":"Ruoxin Sang","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Sang, Ruoxin","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"middle","author":{"id":"https://openalex.org/A5098979177","display_name":"Kiranbir Sodhia","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Sodhia, Kiranbir","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]},{"author_position":"last","author":{"id":"https://openalex.org/A5052374562","display_name":"Rakesh Iyer","orcid":null},"institutions":[],"countries":[],"is_corresponding":false,"raw_author_name":"Iyer, Rakesh","raw_affiliation_strings":[],"raw_orcid":null,"affiliations":[]}],"institutions":[],"countries_distinct_count":0,"institutions_distinct_count":0,"corresponding_author_ids":[],"corresponding_institution_ids":[],"apc_list":null,"apc_paid":null,"fwci":null,"has_fulltext":false,"cited_by_count":0,"citation_normalized_percentile":null,"cited_by_percentile_year":null,"biblio":{"volume":null,"issue":null,"first_page":null,"last_page":null},"is_retracted":false,"is_paratext":false,"is_xpac":false,"primary_topic":{"id":"https://openalex.org/T10772","display_name":"Distributed systems and fault tolerance","score":0.9724000096321106,"subfield":{"id":"https://openalex.org/subfields/1705","display_name":"Computer Networks and Communications"},"field":{"id":"https://openalex.org/fields/17","display_name":"Computer Science"},"domain":{"id":"https://openalex.org/domains/3","display_name":"Physical Sciences"}},"topics":[{"id":"https://openalex.org/T10772","display_name":"Distributed systems and fault tolerance","score":0.9724000096321106,"subfield":{"id":"https://openalex.org/subfields/1705","display_name":"Computer Networks and Communications"},"field":{"id":"https://openalex.org/fields/17","display_name":"Computer Science"},"domain":{"id":"https://openalex.org/domains/3","display_name":"Physical Sciences"}},{"id":"https://openalex.org/T10054","display_name":"Parallel Computing and Optimization Techniques","score":0.008799999952316284,"subfield":{"id":"https://openalex.org/subfields/1708","display_name":"Hardware and Architecture"},"field":{"id":"https://openalex.org/fields/17","display_name":"Computer Science"},"domain":{"id":"https://openalex.org/domains/3","display_name":"Physical Sciences"}},{"id":"https://openalex.org/T10126","display_name":"Logic, programming, and type systems","score":0.004699999932199717,"subfield":{"id":"https://openalex.org/subfields/1702","display_name":"Artificial Intelligence"},"field":{"id":"https://openalex.org/fields/17","display_name":"Computer Science"},"domain":{"id":"https://openalex.org/domains/3","display_name":"Physical Sciences"}}],"keywords":[{"id":"https://openalex.org/keywords/flexibility","display_name":"Flexibility (engineering)","score":0.6898999810218811},{"id":"https://openalex.org/keywords/modular-design","display_name":"Modular design","score":0.6660000085830688},{"id":"https://openalex.org/keywords/fault-tolerance","display_name":"Fault tolerance","score":0.3962000012397766},{"id":"https://openalex.org/keywords/distributed-database","display_name":"Distributed database","score":0.32280001044273376},{"id":"https://openalex.org/keywords/key","display_name":"Key (lock)","score":0.31380000710487366},{"id":"https://openalex.org/keywords/competitor-analysis","display_name":"Competitor analysis","score":0.3089999854564667}],"concepts":[{"id":"https://openalex.org/C41008148","wikidata":"https://www.wikidata.org/wiki/Q21198","display_name":"Computer science","level":0,"score":0.7448999881744385},{"id":"https://openalex.org/C2780598303","wikidata":"https://www.wikidata.org/wiki/Q65921492","display_name":"Flexibility (engineering)","level":2,"score":0.6898999810218811},{"id":"https://openalex.org/C101468663","wikidata":"https://www.wikidata.org/wiki/Q1620158","display_name":"Modular design","level":2,"score":0.6660000085830688},{"id":"https://openalex.org/C120314980","wikidata":"https://www.wikidata.org/wiki/Q180634","display_name":"Distributed computing","level":1,"score":0.48030000925064087},{"id":"https://openalex.org/C63540848","wikidata":"https://www.wikidata.org/wiki/Q3140932","display_name":"Fault tolerance","level":2,"score":0.3962000012397766},{"id":"https://openalex.org/C111919701","wikidata":"https://www.wikidata.org/wiki/Q9135","display_name":"Operating system","level":1,"score":0.36719998717308044},{"id":"https://openalex.org/C199360897","wikidata":"https://www.wikidata.org/wiki/Q9143","display_name":"Programming language","level":1,"score":0.358599990606308},{"id":"https://openalex.org/C115903868","wikidata":"https://www.wikidata.org/wiki/Q80993","display_name":"Software engineering","level":1,"score":0.351500004529953},{"id":"https://openalex.org/C70061542","wikidata":"https://www.wikidata.org/wiki/Q989016","display_name":"Distributed database","level":2,"score":0.32280001044273376},{"id":"https://openalex.org/C26517878","wikidata":"https://www.wikidata.org/wiki/Q228039","display_name":"Key (lock)","level":2,"score":0.31380000710487366},{"id":"https://openalex.org/C127576917","wikidata":"https://www.wikidata.org/wiki/Q624630","display_name":"Competitor analysis","level":2,"score":0.3089999854564667},{"id":"https://openalex.org/C168167062","wikidata":"https://www.wikidata.org/wiki/Q1117970","display_name":"Component (thermodynamics)","level":2,"score":0.3001999855041504},{"id":"https://openalex.org/C2777614354","wikidata":"https://www.wikidata.org/wiki/Q55648394","display_name":"Modular construction","level":3,"score":0.2757999897003174},{"id":"https://openalex.org/C65813073","wikidata":"https://www.wikidata.org/wiki/Q1622420","display_name":"High availability","level":2,"score":0.2727000117301941},{"id":"https://openalex.org/C2779478453","wikidata":"https://www.wikidata.org/wiki/Q6889748","display_name":"Modularity (biology)","level":2,"score":0.2556999921798706}],"mesh":[],"locations_count":1,"locations":[{"id":"doi:10.48550/arxiv.2605.23066","is_oa":true,"landing_page_url":"https://doi.org/10.48550/arxiv.2605.23066","pdf_url":null,"source":{"id":"https://openalex.org/S4306400194","display_name":"arXiv (Cornell University)","issn_l":null,"issn":null,"is_oa":true,"is_in_doaj":false,"is_core":false,"host_organization":"https://openalex.org/I205783295","host_organization_name":"Cornell University","host_organization_lineage":["https://openalex.org/I205783295"],"host_organization_lineage_names":[],"type":"repository"},"license":"cc-by","license_id":"https://openalex.org/licenses/cc-by","version":null,"is_accepted":false,"is_published":null,"raw_source_name":null,"raw_type":"Preprint"}],"best_oa_location":{"id":"doi:10.48550/arxiv.2605.23066","is_oa":true,"landing_page_url":"https://doi.org/10.48550/arxiv.2605.23066","pdf_url":null,"source":{"id":"https://openalex.org/S4306400194","display_name":"arXiv (Cornell University)","issn_l":null,"issn":null,"is_oa":true,"is_in_doaj":false,"is_core":false,"host_organization":"https://openalex.org/I205783295","host_organization_name":"Cornell University","host_organization_lineage":["https://openalex.org/I205783295"],"host_organization_lineage_names":[],"type":"repository"},"license":"cc-by","license_id":"https://openalex.org/licenses/cc-by","version":null,"is_accepted":false,"is_published":false,"raw_source_name":null,"raw_type":"Preprint"},"sustainable_development_goals":[],"awards":[],"funders":[],"has_content":{"pdf":false,"grobid_xml":false},"content_urls":null,"referenced_works_count":0,"referenced_works":[],"related_works":[],"abstract_inverted_index":{"In":[0,28],"a":[1,12,24,34],"landscape":[2],"of":[3,14,43],"high-performance":[4],"distributed":[5,44],"ML":[6,57],"systems,":[7],"JAX":[8],"has":[9],"emerged":[10],"as":[11],"framework":[13],"choice.":[15],"However,":[16],"JAX's":[17],"modular":[18],"design":[19],"philosophy":[20],"leaves":[21],"it":[22],"without":[23],"standardized":[25],"checkpointing":[26,37],"solution.":[27],"this":[29],"paper,":[30],"we":[31],"introduce":[32],"Orbax,":[33],"modular,":[35],"JAX-native":[36],"library":[38,78],"that":[39],"abstracts":[40],"the":[41,56],"complexities":[42],"accelerator":[45],"systems":[46],"while":[47],"also":[48],"providing":[49],"flexibility":[50],"for":[51,71,75],"user-friendly":[52],"checkpoint":[53],"manipulations":[54],"throughout":[55],"model":[58],"lifecycle.":[59],"We":[60],"demonstrate":[61],"performance":[62],"exceeding":[63],"comparable":[64],"PyTorch":[65],"competitors":[66],"by":[67],"up":[68],"to":[69],"3.5$\\times$":[70],"saving":[72],"and":[73],"2$\\times$":[74],"loading.":[76],"The":[77],"is":[79],"available":[80],"at":[81],"https://github.com/google/orbax.":[82]},"counts_by_year":[],"updated_date":"2026-07-01T06:00:48.157686","created_date":"2026-05-26T00:00:00"}