def _check_value_type(value_type):
"""Check value_type meets documented criteria."""
if not (value_type.is_tensor() or
(value_type.is_struct_with_python() and
type_analysis.is_structure_of_tensors(value_type))):
raise TypeError('Expected `value_type` to be `TensorType` or '
'`StructWithPythonType` containing only `TensorType`. '
f'Found type: {repr(value_type)}')
if not (type_analysis.is_structure_of_floats(value_type) or
type_analysis.is_structure_of_integers(value_type)):
raise TypeError('Component dtypes of `value_type` must be all integers or '
f'all floats. Found {value_type}.')
def _check_value_type_compatible_with_config_mode(self, value_type):
py_typecheck.check_type(value_type, factory.ValueType.__args__)
if self._config_mode == _Config.INT:
if not type_analysis.is_structure_of_integers(value_type):
raise TypeError(
f'The `SecureSumFactory` was configured to work with integer '
f'dtypes. All values in provided `value_type` hence must be of '
f'integer dtype. \nProvided value_type: {value_type}')
elif self._config_mode == _Config.FLOAT:
if not type_analysis.is_structure_of_floats(value_type):
raise TypeError(
f'The `SecureSumFactory` was configured to work with floating '
f'point dtypes. All values in provided `value_type` hence must be '
f'of floating point dtype. \nProvided value_type: {value_type}')
else:
raise ValueError(f'Unexpected internal config type: {self._config_mode}')
def create(self, value_type):
# Checks value_type and compute client data dimension.
if (value_type.is_struct_with_python()
and type_analysis.is_structure_of_tensors(value_type)):
num_elements_struct = type_conversions.structure_from_tensor_type_tree(
lambda x: x.shape.num_elements(), value_type)
self._client_dim = sum(tf.nest.flatten(num_elements_struct))
elif value_type.is_tensor():
self._client_dim = value_type.shape.num_elements()
else:
raise TypeError(
'Expected `value_type` to be `TensorType` or '
'`StructWithPythonType` containing only `TensorType`. '
f'Found type: {repr(value_type)}')
# Checks that all values are integers or floats.
if not (type_analysis.is_structure_of_floats(value_type)
or type_analysis.is_structure_of_integers(value_type)):
raise TypeError(
'Component dtypes of `value_type` must all be integers '
f'or floats. Found {repr(value_type)}.')
ddp_agg_process = self._build_aggregation_factory().create(value_type)
init_fn = ddp_agg_process.initialize
@federated_computation.federated_computation(
init_fn.type_signature.result,
computation_types.at_clients(value_type))
def next_fn(state, value):
agg_output = ddp_agg_process.next(state, value)
new_measurements = self._derive_measurements(
agg_output.state, agg_output.measurements)
new_state = agg_output.state
if self._auto_l2_clip:
new_state = self._autotune_component_states(agg_output.state)
return measured_process.MeasuredProcessOutput(
state=new_state,
result=agg_output.result,
measurements=new_measurements)
return aggregation_process.AggregationProcess(init_fn, next_fn)
def create(
self, value_type: factory.ValueType
) -> aggregation_process.AggregationProcess:
# Checks value_type and compute client data dimension.
if (value_type.is_struct()
and type_analysis.is_structure_of_tensors(value_type)):
num_elements_struct = type_conversions.structure_from_tensor_type_tree(
lambda x: x.shape.num_elements(), value_type)
client_dim = sum(tf.nest.flatten(num_elements_struct))
elif value_type.is_tensor():
client_dim = value_type.shape.num_elements()
else:
raise TypeError('Expected `value_type` to be `TensorType` or '
'`StructType` containing only `TensorType`. '
f'Found type: {repr(value_type)}')
# Checks that all values are integers.
if not type_analysis.is_structure_of_integers(value_type):
raise TypeError(
'Component dtypes of `value_type` must all be integers. '
f'Found {repr(value_type)}.')
# Checks that we have enough elements to estimate standard deviation.
if self._estimate_stddev:
if client_dim <= 1:
raise ValueError(
'The stddev estimation procedure expects more than '
'1 element from `value_type`. Found `value_type` of '
f'{value_type} with {client_dim} elements.')
elif client_dim <= 100:
warnings.warn(
f'`value_type` has only {client_dim} elements. The '
'estimated standard deviation may be noisy. Consider '
'setting `estimate_stddev` to True only if the input '
'tensor/structure have more than 100 elements.')
inner_agg_process = self._inner_agg_factory.create(value_type)
init_fn = inner_agg_process.initialize
next_fn = self._create_next_fn(inner_agg_process.next,
init_fn.type_signature.result,
value_type)
return aggregation_process.AggregationProcess(init_fn, next_fn)
def create(self, value_type):
type_args = typing.get_args(factory.ValueType)
py_typecheck.check_type(value_type, type_args)
if not type_analysis.is_structure_of_integers(value_type):
raise TypeError(
'Provided value_type must either be an integer type or'
f'a structure of integer types, 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_clients(value_type))
def next_fn(state, value):
if self._symmetric_range:
# Sum in [-M+1, M-1].
# Delegation to `federated_secure_modular_sum` with modulus 2*M-1 is
# equivalent to modular clip to range [-M+1, M-1]. Then, represent `x`
# in that range as `(x + 2*M-1) % 2*M-1` which is congruent with `x`
# under the desired modulus, thus compatible with secure aggreagtion.
# This is reverted after summation by modular clip to the initial range.
summed_value = intrinsics.federated_secure_modular_sum(
value, 2 * self._modulus - 1)
summed_value = intrinsics.federated_map(
tensorflow_computation.tf_computation(
self._mod_clip_after_symmetric_range_sum),
summed_value)
else:
summed_value = intrinsics.federated_secure_modular_sum(
value, self._modulus)
empty_measurements = intrinsics.federated_value((),
placements.SERVER)
return measured_process.MeasuredProcessOutput(
state, summed_value, empty_measurements)
return aggregation_process.AggregationProcess(init_fn, next_fn)
def _check_value_type_compatible_with_config_mode(self, value_type):
type_args = typing.get_args(factory.ValueType)
py_typecheck.check_type(value_type, type_args)
if not _is_structure_of_single_dtype(value_type):
raise TypeError(
f'Expected a type which is a structure containing the same dtypes, '
f'found {value_type}.')
if self._config_mode == _Config.INT:
if not type_analysis.is_structure_of_integers(value_type):
raise TypeError(
f'The `SecureSumFactory` was configured to work with integer '
f'dtypes. All values in provided `value_type` hence must be of '
f'integer dtype. \nProvided value_type: {value_type}')
elif self._config_mode == _Config.FLOAT:
if not type_analysis.is_structure_of_floats(value_type):
raise TypeError(
f'The `SecureSumFactory` was configured to work with floating '
f'point dtypes. All values in provided `value_type` hence must be '
f'of floating point dtype. \nProvided value_type: {value_type}'
)
else:
raise ValueError(
f'Unexpected internal config type: {self._config_mode}')
def test_returns_false(self, type_spec):
self.assertFalse(type_analysis.is_structure_of_integers(type_spec))
async def compute_federated_secure_sum(
self, arg: federated_resolving_strategy.FederatedResolvingStrategyValue
) -> federated_resolving_strategy.FederatedResolvingStrategyValue:
logging.warning(
'The implementation of the `tff.federated_secure_sum` intrinsic '
'provided by the `tff.backends.test` runtime uses no cryptography.'
)
py_typecheck.check_type(arg.internal_representation, structure.Struct)
py_typecheck.check_len(arg.internal_representation, 2)
summands, bitwidth = await asyncio.gather(
self.ingest_value(arg.internal_representation[0],
arg.type_signature[0]).compute(),
self.ingest_value(arg.internal_representation[1],
arg.type_signature[1]).compute())
summands_type = arg.type_signature[0].member
if not type_analysis.is_structure_of_integers(summands_type):
raise TypeError(
'Cannot compute `federated_secure_sum` on summands that are not '
'TensorType or StructType of TensorType. Got {t}'.format(
t=repr(summands_type)))
if (summands_type.is_struct()
and not structure.is_same_structure(summands_type, bitwidth)):
raise TypeError(
'Cannot compute `federated_secure_sum` if summands and bitwidth are '
'not the same structure. Got summands={s}, bitwidth={b}'.
format(s=repr(summands_type), b=repr(bitwidth.type_signature)))
num_additional_bits = await self._compute_extra_bits_for_secagg()
# Clamp to 64 bits, otherwise we can't represent the mask in TensorFlow.
extended_bitwidth = _map_numpy_or_structure(
bitwidth, fn=lambda b: min(b.numpy() + num_additional_bits, 64))
logging.debug('Emulated secure sum effective bitwidth: %s',
extended_bitwidth)
# Now we need to cast the summands into the integral type that is large
# enough to represent the sum and the mask.
summation_type_spec = _compute_summation_type_for_bitwidth(
extended_bitwidth, summands_type)
# `summands` is a list of all clients' summands. We map
# `_map_numpy_or_structure` to the list, applying it pointwise to clients.
summand_tensors = tf.nest.map_structure(_extract_numpy_arrays,
summands)
# Dtype conversion trick: pull the summand values out, and push them back
# into the executor using the new dtypes decided based on bitwidth.
casted_summands = await self._executor.create_value(
summand_tensors, computation_types.at_clients(summation_type_spec))
# To emulate SecAgg without the random masks, we must mask the summands to
# the effective bitwidth. This isn't strictly necessary because we also
# mask the sum result and modulus operator is distributive, but this more
# accurately reflects the system.
mask = await self._embed_tf_secure_sum_mask_value(
summation_type_spec, extended_bitwidth)
masked_summands = await self._compute_modulus(casted_summands, mask)
logging.debug('Computed masked modular summands as: %s', await
masked_summands.compute())
# Then perform the sum and modolulo operation (using powers of 2 bitmasking)
# on the sum, using the computed effective bitwidth.
sum_result = await self.compute_federated_sum(masked_summands)
modular_sums = await self._compute_modulus(sum_result, mask)
# Dtype conversion trick again, pull the modular sum values out, and push
# them back into the executor using the dypte from the summands.
modular_sum_values = _extract_numpy_arrays(await
modular_sums.compute())
logging.debug('Computed modular sums as: %s', modular_sum_values)
return await self._executor.create_value(
modular_sum_values, computation_types.at_server(summands_type))
def _check_is_integer_struct(value_type, label):
if not type_analysis.is_structure_of_integers(value_type):
raise TypeError(f'Component dtypes of `{label}` must all be integers. '
f'Found {repr(value_type)}.')
