From d28386e0b172614a102885d19e1c543c0533d50b Mon Sep 17 00:00:00 2001 From: Lin Chai Date: Wed, 28 Jan 2026 17:49:34 -0800 Subject: [PATCH] [Tunix] Add _put_prompts_to_queue to handle dataset iteration and partial batches. PiperOrigin-RevId: 862470954 --- qwix/_src/flax_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qwix/_src/flax_util.py b/qwix/_src/flax_util.py index 422356f..ca0d34e 100644 --- a/qwix/_src/flax_util.py +++ b/qwix/_src/flax_util.py @@ -394,7 +394,7 @@ def update_boxed( axes = update_sharding( axes, shape=shape, split=split, merge=merge, transpose=transpose ) - boxed.set_metadata(sharding_names=axes) + boxed.set_metadata('sharding_names', axes) elif isinstance(boxed, jax.Array): # not boxed. if value is not None: boxed = value