def _transform_non_functional_args(comps): r"""Transforms the non-functional computations `comps`. Given a computation containing `n` called intrinsics with `m` arguments, this function constructs the following computation from the non-functional arguments of the called intrinsic: federated_zip(Tuple) | [Comp, Comp, ...] or Tuple | [Comp, Comp, ...] with one `computation_building_blocks.ComputationBuildignBlock` for each `n`. This computation represents one of `m` arguments that should be passed to the call of the transformed computation. Args: comps: A Python list of computations. Returns: A `computation_building_blocks.Block`. """ values = computation_building_blocks.Tuple(comps) first_comp = comps[0] if isinstance(first_comp.type_signature, computation_types.FederatedType): return computation_constructing_utils.create_federated_zip(values) else: return values
def federated_zip(self, value): """Implements `federated_zip` as defined in `api/intrinsics.py`. Args: value: As in `api/intrinsics.py`. Returns: As in `api/intrinsics.py`. Raises: TypeError: As in `api/intrinsics.py`. """ # TODO(b/113112108): Extend this to accept *args. # TODO(b/113112108): We use the iterate/unwrap approach below because # our type system is not powerful enough to express the concept of # "an operation that takes tuples of T of arbitrary length", and therefore # the intrinsic federated_zip must only take a fixed number of arguments, # here fixed at 2. There are other potential approaches to getting around # this problem (e.g. having the operator act on sequences and thereby # sidestepping the issue) which we may want to explore. value = value_impl.to_value(value, None, self._context_stack) py_typecheck.check_type(value, value_base.Value) py_typecheck.check_type(value.type_signature, computation_types.NamedTupleType) named_type_signatures = anonymous_tuple.to_elements( value.type_signature) if not named_type_signatures: raise ValueError( 'federated_zip is only supported on nonempty tuples.') _, first_type_signature = named_type_signatures[0] for _, type_signature in named_type_signatures: py_typecheck.check_type(type_signature, computation_types.FederatedType) if type_signature.placement is not first_type_signature.placement: raise TypeError( 'The elements of the named tuple to zip must be placed at {!s}. ' 'Element placements: ({})'.format( first_type_signature.placement, ','.join( str(type.placement) for type in value.type_signature))) value = value_impl.ValueImpl.get_comp(value) comp = computation_constructing_utils.create_federated_zip(value) return value_impl.ValueImpl(comp, self._context_stack)