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)