Changes the use of dict to OrderedDict in gather_encoder.py and simple_encoder.py.

PiperOrigin-RevId: 421564691

Author: PraChetit

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/gather_encoder.py

@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import core_encoder
@@ -147,7 +148,6 @@ def from_encoder(cls, encoder, tensorspec):
     if not tensorspec.shape.is_fully_defined():
       raise TypeError('The shape of provided tensorspec must be fully defined.')
 
-    tensorspec = tensorspec
     commuting_structure = encoder.commuting_structure
     state_update_aggregation_modes = tf.nest.flatten(
         encoder.state_update_aggregation_modes)
@@ -187,8 +187,8 @@ def from_encoder(cls, encoder, tensorspec):
     # of the tensor_encoding tool do not even need to be aware of it. This
     # argument is well supported for instance in the book of John Ousterhout,
     # "A Philosophy of Software Design".
-    internal_structure = {}
-    internal_py_values = {}
+    internal_structure = collections.OrderedDict()
+    internal_py_values = collections.OrderedDict()
 
     def _add_to_structure(key, value):
       if key not in internal_structure:
@@ -226,10 +226,10 @@ def get_params_fn(flat_state):
       _, input_shapes_after_sum = (
           core_encoder.split_shapes_by_commuting_structure(
               input_shapes, commuting_structure))
-      decode_after_sum_params = {
-          _PARAMS: decode_after_sum_params,
-          _SHAPES: input_shapes_after_sum
-      }
+      decode_after_sum_params = collections.OrderedDict([
+          (_PARAMS, decode_after_sum_params),
+          (_SHAPES, input_shapes_after_sum),
+      ])
 
       encode_params_py, encode_params_tf = py_utils.split_dict_py_tf(
           encode_params)
@@ -274,18 +274,18 @@ def encode_fn(x, params):
           core_encoder.split_shapes_by_commuting_structure(
               input_shapes, commuting_structure))
 
-      encoded_structure = {
-          _TENSORS: encoded_x,
-          _SHAPES: input_shapes_before_sum
-      }
+      encoded_structure = collections.OrderedDict([
+          (_TENSORS, encoded_x),
+          (_SHAPES, input_shapes_before_sum),
+      ])
       encoded_structure_py, encoded_structure_tf = py_utils.split_dict_py_tf(
           encoded_structure)
 
       _add_to_structure('encoded_structure', encoded_structure_tf)
       _add_to_structure('state_update_tensors', state_update_tensors)
       _add_to_py_values('encoded_structure', encoded_structure_py)
 
-      return (dict(
+      return (collections.OrderedDict(
           py_utils.flatten_with_joined_string_paths(encoded_structure_tf)),
               tuple(tf.nest.flatten(state_update_tensors)))
 
@@ -316,7 +316,7 @@ def decode_before_sum_fn(encoded_structure, params):
 
       _add_to_structure('part_decoded_structure', part_decoded_structure)
       if isinstance(part_decoded_structure, dict):
-        return dict(
+        return collections.OrderedDict(
             py_utils.flatten_with_joined_string_paths(part_decoded_structure))
       else:
         return part_decoded_structure

tensorflow_model_optimization/python/core/internal/tensor_encoding/core/simple_encoder.py

@@ -17,6 +17,7 @@
 from __future__ import division
 from __future__ import print_function
 
+import collections
 import tensorflow as tf
 
 from tensorflow_model_optimization.python.core.internal.tensor_encoding.core import core_encoder
@@ -79,8 +80,8 @@ def __init__(self, encoder, tensorspec):
     # methods, to be used in encode_fn and decode_fn methods, respectively.
     # Decorated by tf.function, their necessary side effects are realized during
     # call to get_concrete_function().
-    state_py_structure = {}
-    encoded_py_structure = {}
+    state_py_structure = collections.OrderedDict()
+    encoded_py_structure = collections.OrderedDict()
 
     @tf.function
     def initial_state_fn():
@@ -108,12 +109,12 @@ def encode_fn(x, flat_state):
       # The following code converts the nested structres necessary for the
       # underlying encoder, to a single flat dictionary, which is simpler to
       # manipulate by the users of SimpleEncoder.
-      full_encoded_structure = {
-          _TENSORS: encoded_x,
-          _PARAMS: decode_params,
-          _SHAPES: input_shapes
-      }
-      flat_encoded_structure = dict(
+      full_encoded_structure = collections.OrderedDict([
+          (_TENSORS, encoded_x),
+          (_PARAMS, decode_params),
+          (_SHAPES, input_shapes),
+      ])
+      flat_encoded_structure = collections.OrderedDict(
           py_utils.flatten_with_joined_string_paths(full_encoded_structure))
       flat_encoded_py_structure, flat_encoded_tf_structure = (
           py_utils.split_dict_py_tf(flat_encoded_structure))