def _as_function_of_some_federated_subparameters(
bb: building_blocks.Lambda,
paths,
) -> building_blocks.Lambda:
"""Turns `x -> ...only uses parts of x...` into `parts_of_x -> ...`."""
tree_analysis.check_has_unique_names(bb)
bb = _prepare_for_rebinding(bb)
name_generator = building_block_factory.unique_name_generator(bb)
type_list = []
int_paths = []
for path in paths:
selected_type = bb.parameter_type
int_path = []
for index in path:
if not selected_type.is_struct():
raise _ParameterSelectionError(path, bb)
if isinstance(index, int):
if index > len(selected_type):
raise _ParameterSelectionError(path, bb)
int_path.append(index)
else:
py_typecheck.check_type(index, str)
if not structure.has_field(selected_type, index):
raise _ParameterSelectionError(path, bb)
int_path.append(
structure.name_to_index_map(selected_type)[index])
selected_type = selected_type[index]
if not selected_type.is_federated():
raise _NonFederatedSelectionError(
'Attempted to rebind references to parameter selection path '
f'{path} from type {bb.parameter_type}, but the value at that path '
f'was of non-federated type {selected_type}. Selections must all '
f'be of federated type. Original AST:\n{bb}')
int_paths.append(tuple(int_path))
type_list.append(selected_type)
placement = type_list[0].placement
if not all(x.placement is placement for x in type_list):
raise _MismatchedSelectionPlacementError(
'In order to zip the argument to the lower-level lambda together, all '
'selected arguments should be at the same placement. Your selections '
f'have resulted in the list of types:\n{type_list}')
zip_type = computation_types.FederatedType([x.member for x in type_list],
placement=placement)
ref_to_zip = building_blocks.Reference(next(name_generator), zip_type)
path_to_replacement = {}
for i, path in enumerate(int_paths):
path_to_replacement[path] = _construct_selection_from_federated_tuple(
ref_to_zip, i, name_generator)
new_lambda_body = _replace_selections(bb.result, bb.parameter_name,
path_to_replacement)
lambda_with_zipped_param = building_blocks.Lambda(
ref_to_zip.name, ref_to_zip.type_signature, new_lambda_body)
tree_analysis.check_contains_no_new_unbound_references(
bb, lambda_with_zipped_param)
return lambda_with_zipped_param
async def create_selection(self, source, index=None, name=None):
py_typecheck.check_type(source, RemoteValue)
py_typecheck.check_type(source.type_signature,
computation_types.StructType)
if index is None:
py_typecheck.check_type(name, str)
index = structure.name_to_index_map(source.type_signature)[name]
py_typecheck.check_type(index, int)
result_type = source.type_signature[index]
request = executor_pb2.CreateSelectionRequest(
source_ref=source.value_ref, index=index)
response = _request(self._stub.CreateSelection, request)
py_typecheck.check_type(response, executor_pb2.CreateSelectionResponse)
return RemoteValue(response.value_ref, result_type, self)
def _traverse_selection(comp, transform, context_tree, identifier_seq):
"""Helper function holding traversal logic for selection nodes."""
_ = next(identifier_seq)
source, source_modified = _transform_postorder_with_symbol_bindings_switch(
comp.source, transform, context_tree, identifier_seq)
if source_modified:
# Normalize selection to index based on the type signature of the
# original source. The new source may not have names present.
if comp.index is not None:
index = comp.index
else:
index = structure.name_to_index_map(
comp.source.type_signature)[comp.name]
comp = building_blocks.Selection(source, index=index)
comp, comp_modified = transform(comp, context_tree)
return comp, comp_modified or source_modified
def test_name_to_index_map_fully_named_struct(self):
partially_named_struct = structure.Struct([('b', 10), ('a', 20)])
name_to_index_dict = structure.name_to_index_map(partially_named_struct)
expected_name_to_index_map = {'b': 0, 'a': 1}
self.assertEqual(name_to_index_dict, expected_name_to_index_map)
def test_name_to_index_map_empty_unnamed_struct(self): unnamed_struct = structure.Struct([(None, 10), (None, 20)]) self.assertEmpty(structure.name_to_index_map(unnamed_struct))
def test_name_to_index_map_fully_named_struct(self):
partially_named_struct = structure.Struct.named(b=10, a=20)
name_to_index_dict = structure.name_to_index_map(partially_named_struct)
expected_name_to_index_map = {'b': 0, 'a': 1}
self.assertEqual(name_to_index_dict, expected_name_to_index_map)
def test_name_to_index_map_empty_unnamed_struct(self): unnamed_struct = structure.Struct.unnamed(10, 20) self.assertEmpty(structure.name_to_index_map(unnamed_struct))