From d26e8f67d373fedbc0abd50e40c3f420db3752c1 Mon Sep 17 00:00:00 2001 From: Jake VanderPlas Date: Wed, 7 Oct 2026 06:25:42 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 995067105 --- learned_optimization/tree_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/learned_optimization/tree_utils.py b/learned_optimization/tree_utils.py index 5467c60f..d346f029 100644 --- a/learned_optimization/tree_utils.py +++ b/learned_optimization/tree_utils.py @@ -243,6 +243,6 @@ def partition_unflatten(unflattener: PartitionUnflatten, to_fill = [None for _ in keys] for name, part in zip(names, part_values): for n, p in zip(name, part): - to_fill[unmap[n]] = p + to_fill[unmap[n]] = p # pyrefly: ignore[unsupported-operation] return jax.tree_util.tree_unflatten(struct, to_fill)