From 790825c44702310cb98ba456b3af89f7462229aa Mon Sep 17 00:00:00 2001 From: Orbax Authors Date: Mon, 31 Aug 2026 11:16:58 -0700 Subject: [PATCH] group the tensors to serialize in linear rather than quadratic time PiperOrigin-RevId: 973983399 --- .../handlers/base_pytree_checkpoint_handler.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py b/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py index b3498546d..be4878bb3 100644 --- a/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py +++ b/checkpoint/orbax/checkpoint/_src/handlers/base_pytree_checkpoint_handler.py @@ -151,13 +151,13 @@ def _group_value( if handler not in grouped: grouped[handler] = BatchRequest(handler, [], [], [], []) request = grouped[handler] - grouped[handler] = dataclasses.replace( - request, - keys=request.keys + [tuple_key], - values=request.values + [value], - infos=request.infos + [info], - args=request.args + [arg], - ) + # `BatchRequest` is mutable, so the entry is appended in place: rebuilding + # it per leaf copies all four lists and makes this loop quadratic in the + # number of tensors, which dominates for a tree with many small leaves. + request.keys.append(tuple_key) + request.values.append(value) + request.infos.append(info) + request.args.append(arg) jax.tree_util.tree_map_with_path( _group_value,