Migrate calls to tff.structure.to_elements to use the API on federated_language.StructType or federated_language.framework.Struct.

PiperOrigin-RevId: 746173864

Author: michaelreneer

tensorflow_federated/python/core/backends/mapreduce/form_utils.py

@@ -1329,7 +1329,7 @@ def _has_placement(type_spec):
     # Add a reference to the injected broadcast call in the result so that it
     # does not get pruned by various tree transformations. We will remove this
     # additional element in the result after the first split operation.
-    revised_block_result = structure.to_elements(comp_tree.result.result) + [
+    revised_block_result = list(comp_tree.result.result.items()) + [
         federated_language.framework.Reference(
             'injected_broadcast_ref',
             injected_broadcast.type_signature,
@@ -1438,7 +1438,7 @@ def _unnest_lambda_parameter(comp):
         federated_language.framework.Block(
             after_broadcast.result.locals,
             federated_language.framework.Struct(
-                structure.to_elements(after_broadcast.result.result)[:-1]
+                list(after_broadcast.result.result.items())[:-1]
             ),
         ),
     )

tensorflow_federated/python/core/environments/jax_frontend/jax_serialization.py

@@ -134,7 +134,7 @@ def _make(
       return obj, next_unused_tensor_index
     elif isinstance(type_spec, federated_language.StructType):
       elements = []
-      for k, v in structure.to_elements(type_spec):
+      for k, v in type_spec.items():
         obj, next_unused_tensor_index = _make(v, next_unused_tensor_index)
         elements.append((k, obj))
       obj = _XlaSerializerStructArg(type_spec, elements)

tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_computation_test_utils.py

@@ -60,7 +60,7 @@ def _stamp_value_into_graph(
     if isinstance(value, (list, dict)):
       value = structure.from_container(value)
     stamped_elements = []
-    named_type_signatures = structure.to_elements(type_signature)
+    named_type_signatures = type_signature.items()
     for (name, type_signature), element in zip(named_type_signatures, value):
       stamped_element = _stamp_value_into_graph(element, type_signature, graph)
       stamped_elements.append((name, stamped_element))

tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_utils.py

@@ -94,7 +94,8 @@ def stamp_parameter_in_graph(parameter_name, parameter_type, graph):
   elif isinstance(parameter_type, federated_language.StructType):
     # The parameter_type could be a StructTypeWithPyContainer, however, we
     # ignore that for now. Instead, the proper containers will be inserted at
-    # call time by federated_language.framework.wrap_as_zero_or_one_arg_callable.
+    # call time by federated_language.framework.wrap_as_zero_or_one_arg
+    # _callable.
     if not parameter_type:
       # Stamps whimsy element to "populate" graph, as TensorFlow does not
       # support empty graphs.
@@ -533,7 +534,7 @@ def _assemble_result_from_graph(type_spec, binding, output_map):
           'Expected a struct binding, found {}.'.format(binding_oneof)
       )
     else:
-      type_elements = structure.to_elements(type_spec)
+      type_elements = type_spec.items()
       if len(binding.struct.element) != len(type_elements):
         raise ValueError(
             'Mismatching tuple sizes in type ({}) and binding ({}).'.format(
@@ -601,7 +602,7 @@ def _make_empty_list_structure_for_element_type_spec(type_spec):
   if isinstance(type_spec, federated_language.TensorType):
     return []
   elif isinstance(type_spec, federated_language.StructType):
-    elements = structure.to_elements(type_spec)
+    elements = type_spec.items()
     if all(k is not None for k, _ in elements):
       return collections.OrderedDict([
           (k, _make_empty_list_structure_for_element_type_spec(v))
@@ -684,7 +685,7 @@ def _handle_none_dimension(x):
       return np.empty(whimsy_shape, dtype=np.str_)
     return np.zeros(whimsy_shape, type_spec.dtype)
   elif isinstance(type_spec, federated_language.StructType):
-    elements = structure.to_elements(type_spec)
+    elements = type_spec.items()
     elem_list = []
     for _, elem_type in elements:
       elem_list.append(_make_whimsy_element_for_type_spec(elem_type))
@@ -735,7 +736,7 @@ def _append_to_list_structure_for_element_type_spec(nested, value, type_spec):
     # tf.data.Dataset.from_tensor_slices.
     nested.append(tf.convert_to_tensor(value, type_spec.dtype))
   elif isinstance(type_spec, federated_language.StructType):
-    elements = structure.to_elements(type_spec)
+    elements = type_spec.items()
     if isinstance(nested, collections.OrderedDict):
       if isinstance(value, py_typecheck.SupportsNamedTuple):
         # In Python 3.8 and later `_asdict` no longer return OrdereDict, rather
@@ -814,7 +815,7 @@ def _replace_empty_leaf_lists_with_numpy_arrays(lists, type_spec):
     else:
       return np.array([], dtype=type_spec.dtype)
   elif isinstance(type_spec, federated_language.StructType):
-    elements = structure.to_elements(type_spec)
+    elements = type_spec.items()
     if isinstance(lists, collections.OrderedDict):
       to_return = []
       for elem_name, elem_type in elements:
@@ -1074,7 +1075,7 @@ def _to_representative_value(type_spec, elements):
           return py_type(*values)
         return py_type(values)  # pylint: disable=too-many-function-args
     elif isinstance(type_spec, federated_language.StructType):
-      field_types = structure.to_elements(type_spec)
+      field_types = type_spec.items()
       is_all_named = all([name is not None for name, _ in field_types])
       if is_all_named:
         if isinstance(elements, py_typecheck.SupportsNamedTuple):

tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_utils_test.py

@@ -70,7 +70,7 @@ def _assert_binding_matches_type_and_value(
       if not isinstance(val, (list, tuple, structure.Struct)):
         self.assertIsInstance(val, dict)
         val = list(val.values())
-      for idx, e in enumerate(structure.to_elements(type_spec)):
+      for idx, e in enumerate(type_spec.items()):
         self._assert_binding_matches_type_and_value(
             binding.struct.element[idx], e[1], val[idx], graph, is_output
         )

tensorflow_federated/python/core/environments/tensorflow_backend/type_conversions.py

@@ -131,7 +131,7 @@ def _type_to_tf_dtypes_and_shapes(type_spec: federated_language.Type):
     shape = tf.TensorShape(type_spec.shape)
     return (type_spec.dtype, shape)
   elif isinstance(type_spec, federated_language.StructType):
-    elements = structure.to_elements(type_spec)
+    elements = type_spec.items()
     if not elements:
       output_dtypes = []
       output_shapes = []
@@ -238,7 +238,7 @@ def type_to_tf_structure(type_spec: federated_language.Type):
   if isinstance(type_spec, federated_language.TensorType):
     return tf.TensorSpec(type_spec.shape, type_spec.dtype)
   elif isinstance(type_spec, federated_language.StructType):
-    elements = structure.to_elements(type_spec)
+    elements = type_spec.items()
     if not elements:
       return ()
     element_outputs = [(k, type_to_tf_structure(v)) for k, v in elements]

tensorflow_federated/python/core/impl/compiler/BUILD

@@ -51,7 +51,6 @@ py_test(
         ":building_block_test_utils",
         ":transformations",
         ":tree_transformations",
-        "//tensorflow_federated/python/common_libs:structure",
         "@federated_language//federated_language",
         "@federated_language//federated_language/proto:computation_py_pb2",
     ],

tensorflow_federated/python/core/impl/compiler/transformations.py

@@ -1444,9 +1444,7 @@ def divisive_force_align_and_split_by_intrinsics(
   ]
 
   # Update the before comp result to produce the extended intermediate state.
-  before_result_elements = structure.to_elements(
-      preliminary_before_comp.result.result
-  )
+  before_result_elements = list(preliminary_before_comp.result.result.items())
   intermediate_state_index_in_before_result = 1
   intermediate_state_name, intermediate_state_vals = before_result_elements[
       intermediate_state_index_in_before_result
@@ -1461,8 +1459,7 @@ def divisive_force_align_and_split_by_intrinsics(
       for local_name, local_value in duplicated_locals
   ]
   extended_intermediate_state_vals = federated_language.framework.Struct(
-      structure.to_elements(intermediate_state_vals)
-      + duplicate_intermediate_state_vals
+      list(intermediate_state_vals.items()) + duplicate_intermediate_state_vals
   )
   before_result_elements[intermediate_state_index_in_before_result] = (
       intermediate_state_name,

tensorflow_federated/python/core/impl/compiler/transformations_test.py

@@ -19,7 +19,6 @@
 from federated_language.proto import computation_pb2
 import numpy as np
 
-from tensorflow_federated.python.common_libs import structure
 from tensorflow_federated.python.core.impl.compiler import building_block_test_utils
 from tensorflow_federated.python.core.impl.compiler import transformations
 from tensorflow_federated.python.core.impl.compiler import tree_transformations
@@ -751,7 +750,7 @@ def _check_transformed_comp_validity(
         transformed_comp.parameter_type, len(original_comp.parameter_type) + 1
     )
     self.assertEqual(
-        structure.to_elements(transformed_comp.parameter_type)[
+        transformed_comp.parameter_type.items()[
             len(transformed_comp.parameter_type) - 1
         ][0],
         lambda_parameter_extension_name,
@@ -1011,7 +1010,7 @@ def check_split_signatures(self, original_comp, before, intrinsic, after):
         before.type_signature.result, federated_language.StructType
     )
     self.assertEqual(
-        [x for x, _ in structure.to_elements(before.type_signature.result)],
+        [x for x, _ in before.type_signature.result.items()],
         ['intrinsic_args_from_before_comp', 'intermediate_state'],
     )
     self.assertIsInstance(
@@ -1025,7 +1024,7 @@ def check_split_signatures(self, original_comp, before, intrinsic, after):
         len(intrinsic.type_signature.parameter) - 1
     )
     intrinsic_arg_names = [
-        x for x, _ in structure.to_elements(intrinsic.type_signature.parameter)
+        x for x, _ in intrinsic.type_signature.parameter.items()
     ]
     self.assertEqual(
         intrinsic_arg_names[intrinsic_args_from_before_comp_index],
@@ -1047,7 +1046,7 @@ def check_split_signatures(self, original_comp, before, intrinsic, after):
     intrinsic_results_index = len(after.type_signature.parameter) - 2
     intermediate_state_index = len(after.type_signature.parameter) - 1
     after_comp_parameter_names = [
-        x for x, _ in structure.to_elements(after.type_signature.parameter)
+        x for x, _ in after.type_signature.parameter.items()
     ]
     self.assertEqual(
         after_comp_parameter_names[intrinsic_results_index], 'intrinsic_results'

tensorflow_federated/python/learning/models/model_weights_test.py

@@ -271,11 +271,9 @@ def test_converts_struct(self):
         lambda item: self.assertIsInstance(item, np.ndarray),
         converted.non_trainable,
     )
+    self.assertAllEqual(converted.trainable, [(None, np.array([1.0]))])
     self.assertAllEqual(
-        structure.to_elements(converted.trainable), [(None, np.array([1.0]))]
-    )
-    self.assertAllEqual(
-        structure.to_elements(converted.non_trainable),
+        converted.non_trainable,
         [(None, np.array([2.0])), (None, np.array([3.0]))],
     )
 
@@ -296,7 +294,7 @@ def test_converts_heterogeneous_struct(self):
         converted.trainable,
     )
     self.assertAllEqual(
-        structure.to_elements(converted.trainable),
+        converted.trainable,
         [
             ('a', 1),
             ('b', 2.0),