def record_intrinsic_calls(comp):
"""Identifies matching calls and adds them to `aggregation_calls`."""
if not comp.is_call():
return
# Called lambdas will only trigger aggregation if they themselves contain
# aggregation, which will be caught when the lambea itself is traversed.
if comp.function.is_lambda():
return
# Aggregation cannot be occurring if the output type is not federated
if not type_analysis.contains_federated_types(
comp.function.type_signature.result):
return
# We can't tell whether an arbitrary AST fragment results in an intrinsic
# with a given URI, so we report an error in this case.
if not comp.function.is_intrinsic():
raise ValueError(
'Cannot determine whether call contains aggregation: ' +
str(comp))
# Aggregation with inputs that don't contain any tensors isn't interesting.
#
# NOTE: this is only applicable to intrinsic calls. Users can write their
# own functions that internally materialize values at clients + aggregate
# without taking any input tensors.
#
# This means that this check *must* come after the check above ensuring
# that we're only talking about calls to `building_blocks.Intrinsic`s.
if comp.argument is None or not type_analysis.contains_tensor_types(
comp.argument.type_signature):
return
if comp.function.uri in aggregation_uris:
aggregation_calls.append(comp)
def build_broadcast_process(value_type: computation_types.Type):
"""Builds `DistributionProcess` directly broadcasting values.
The created process has empty state and reports no measurements.
Args:
value_type: A non-federated `tff.Type` of value to be broadcasted.
Returns:
A `DistributionProcess` for broadcasting `value_type`.
Raises:
TypeError: If `value_type` contains a `tff.types.FederatedType`.
"""
py_typecheck.check_type(
value_type, (computation_types.TensorType, computation_types.StructType))
if type_analysis.contains_federated_types(value_type):
raise TypeError(
f'Provided value_type must not contain any tff.types.FederatedType, '
f'but found: {value_type}')
@federated_computation.federated_computation
def init_fn():
return intrinsics.federated_value((), placements.SERVER)
@federated_computation.federated_computation(
init_fn.type_signature.result, computation_types.at_server(value_type))
def next_fn(state, value):
empty_measurements = intrinsics.federated_value((), placements.SERVER)
return measured_process.MeasuredProcessOutput(
state, intrinsics.federated_broadcast(value), empty_measurements)
return DistributionProcess(init_fn, next_fn)
def _is_local(comp):
cached = local_cache.get(comp, None)
if cached is not None:
return cached
if (comp.is_intrinsic() or comp.is_data() or comp.is_placement() or
type_analysis.contains_federated_types(comp.type_signature)):
local_cache[comp] = False
return False
if (comp.is_compiled_computation()
and comp.proto.WhichOneof('computation') == 'xla'):
local_cache[comp] = False
return False
for child in comp.children():
if not _is_local(child):
local_cache[comp] = False
return False
return True
def federated_reduce(self, value, zero, op):
"""Implements `federated_reduce` as defined in `api/intrinsics.py`."""
value = value_impl.to_value(value, None, self._context_stack)
value = value_utils.ensure_federated_value(value,
placement_literals.CLIENTS,
'value to be reduced')
zero = value_impl.to_value(zero, None, self._context_stack)
if type_analysis.contains_federated_types(zero.type_signature):
raise TypeError(
'`zero` may not contain a federated type, found type:\n' +
str(zero.type_signature))
op = value_impl.to_value(
op,
None,
self._context_stack,
parameter_type_hint=computation_types.StructType(
[zero.type_signature, value.type_signature.member]))
op.type_signature.check_function()
if not op.type_signature.result.is_assignable_from(
zero.type_signature):
raise TypeError(
'`zero` must be assignable to the result type from `op`:\n',
computation_types.type_mismatch_error_message(
zero.type_signature, op.type_signature.result,
computation_types.TypeRelation.ASSIGNABLE))
op_type_expected = type_factory.reduction_op(
op.type_signature.result, value.type_signature.member)
if not op_type_expected.is_assignable_from(op.type_signature):
raise TypeError('Expected an operator of type {}, got {}.'.format(
op_type_expected, op.type_signature))
value = value_impl.ValueImpl.get_comp(value)
zero = value_impl.ValueImpl.get_comp(zero)
op = value_impl.ValueImpl.get_comp(op)
comp = building_block_factory.create_federated_reduce(value, zero, op)
comp = self._bind_comp_as_reference(comp)
return value_impl.ValueImpl(comp, self._context_stack)
def build_apply_optimizer_finalizer(
optimizer: optimizer_base.Optimizer,
model_weights_type: computation_types.StructType):
"""Builds finalizer that applies a step of an optimizer.
The provided `model_weights_type` must be a non-federated `tff.Type` with the
`tff.learning.ModelWeights` container.
The 2nd input argument of the created `FinalizerProcess.next` expects a value
matching `model_weights_type` and its 3rd argument expects value matching
`model_weights_type.trainable`. The `optimizer` will be applied to the
trainable model weights only, leaving non_trainable weights unmodified.
The state of the process is the state of the `optimizer` and the process
returns empty measurements.
Args:
optimizer: A `tff.learning.optimizers.Optimizer`.
model_weights_type: A non-federated `tff.Type` of the model weights to be
optimized, which must have a `tff.learning.ModelWeights` container.
Returns:
A `FinalizerProcess` that applies the `optimizer`.
Raises:
TypeError: If `value_type` does not have a `tff.learning.ModelWeights`
Python container, or contains a `tff.types.FederatedType`.
"""
# TODO(b/190334722): Include support for Keras optimizers via
# tff.learning.optimizers.KerasOptimizer when ready.
py_typecheck.check_type(optimizer, optimizer_base.Optimizer)
if (not model_weights_type.is_struct_with_python()
or model_weights_type.python_container != model_utils.ModelWeights
or type_analysis.contains_federated_types(model_weights_type)):
raise TypeError(
f'Provided value_type must be a tff.types.StructType with its python '
f'container being tff.learning.ModelWeights, not containing a '
f'tff.types.FederatedType, but found: {model_weights_type}')
@computations.federated_computation
def init_fn():
tensor_specs = type_conversions.type_to_tf_tensor_specs(
model_weights_type.trainable)
return intrinsics.federated_eval(
computations.tf_computation(
lambda: optimizer.initialize(tensor_specs)), placements.SERVER)
@computations.federated_computation(
init_fn.type_signature.result,
computation_types.at_server(model_weights_type),
computation_types.at_server(model_weights_type.trainable))
def next_fn(state, weights, update):
optimizer_state, new_trainable_weights = intrinsics.federated_map(
computations.tf_computation(optimizer.next),
(state, weights.trainable, update))
new_weights = intrinsics.federated_zip(
model_utils.ModelWeights(new_trainable_weights,
weights.non_trainable))
empty_measurements = intrinsics.federated_value((), placements.SERVER)
return measured_process.MeasuredProcessOutput(optimizer_state,
new_weights,
empty_measurements)
return FinalizerProcess(init_fn, next_fn)
def build_apply_optimizer_finalizer(
optimizer_fn: Union[optimizer_base.Optimizer,
Callable[[], tf.keras.optimizers.Optimizer]],
model_weights_type: computation_types.StructType):
"""Builds finalizer that applies a step of an optimizer.
The provided `model_weights_type` must be a non-federated `tff.Type` with the
`tff.learning.ModelWeights` container.
The 2nd input argument of the created `FinalizerProcess.next` expects a value
matching `model_weights_type` and its 3rd argument expects value matching
`model_weights_type.trainable`. The `optimizer` will be applied to the
trainable model weights only, leaving non_trainable weights unmodified.
The state of the process is the state of the `optimizer` and the process
returns empty measurements.
Args:
optimizer_fn: A `tff.learning.optimizers.Optimizer` or a no-arg function
that returns a `tf.keras.optimizers.Optimizer`.
This optimizer is used to apply client updates to the server model.
model_weights_type: A non-federated `tff.Type` of the model weights to be
optimized, which must have a `tff.learning.ModelWeights` container.
Returns:
A `FinalizerProcess` that applies the `optimizer`.
Raises:
TypeError: If `value_type` does not have a `tff.learning.ModelWeights`
Python container, or contains a `tff.types.FederatedType`.
"""
if not isinstance(optimizer_fn, optimizer_base.Optimizer):
if not callable(optimizer_fn) or not isinstance(
optimizer_fn(), tf.keras.optimizers.Optimizer):
raise TypeError(
'The optimizer_fn must be a `tff.learning.optimizers.Optimizer`, or '
'a no-arg callable returning a `tf.keras.optimizers.Optimizer`.'
)
if (not model_weights_type.is_struct_with_python()
or model_weights_type.python_container != model_utils.ModelWeights
or type_analysis.contains_federated_types(model_weights_type)):
raise TypeError(
f'Provided value_type must be a tff.types.StructType with its python '
f'container being tff.learning.ModelWeights, not containing a '
f'tff.types.FederatedType, but found: {model_weights_type}')
if isinstance(optimizer_fn, optimizer_base.Optimizer):
init_tf, next_tf = _build_tff_optimizer_initialize_and_next(
model_weights_type, optimizer_fn)
else:
init_tf, next_tf = _build_keras_optimizer_initialize_and_next(
model_weights_type, optimizer_fn)
@federated_computation.federated_computation
def init_fn():
return intrinsics.federated_eval(init_tf, placements.SERVER)
@federated_computation.federated_computation(
init_fn.type_signature.result,
computation_types.at_server(model_weights_type),
computation_types.at_server(model_weights_type.trainable))
def next_fn(state, weights, update):
optimizer_state, new_trainable_weights = intrinsics.federated_map(
next_tf, (state, weights.trainable, update))
new_weights = intrinsics.federated_zip(
model_utils.ModelWeights(new_trainable_weights,
weights.non_trainable))
empty_measurements = intrinsics.federated_value((), placements.SERVER)
return measured_process.MeasuredProcessOutput(optimizer_state,
new_weights,
empty_measurements)
return FinalizerProcess(init_fn, next_fn)