From 575b94922539b7226dfbf3e3a98da2e93ec633d0 Mon Sep 17 00:00:00 2001 From: Ricardo Torres Date: Thu, 14 May 2026 20:59:00 -0700 Subject: [PATCH] The change ensures that the elements are packed back into the original structure in the correct order by using keys derived from the structure's flattened paths. PiperOrigin-RevId: 915766611 --- .../core/internal/tensor_encoding/core/simple_encoder.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/simple_encoder.py b/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/simple_encoder.py index 039665692..09ff5844a 100644 --- a/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/simple_encoder.py +++ b/tensorflow_model_optimization/python/core/internal/tensor_encoding/core/simple_encoder.py @@ -134,7 +134,12 @@ def decode_fn(encoded_structure): encoded_structure = py_utils.merge_dicts(encoded_structure, encoded_py_structure['flat_py']) encoded_structure = tf.nest.pack_sequence_as( - encoded_py_structure['full'], tf.nest.flatten(encoded_structure)) + encoded_py_structure['full'], + [ + encoded_structure[k] for k, _ in + py_utils.flatten_with_joined_string_paths( + encoded_py_structure['full']) + ]) return encoder.decode(encoded_structure[_TENSORS], encoded_structure[_PARAMS], encoded_structure[_SHAPES])