def apply(transform_fn: computation_base.Computation,
arg_process: EstimationProcess):
"""Builds an `EstimationProcess` by applying `transform_fn` to `arg_process`.
Args:
transform_fn: A `computation_base.Computation` to apply to the estimate of
the arg_process.
arg_process: An `EstimationProcess` to which the transformation will be
applied.
Returns:
An estimation process that applies `transform_fn` to the result of calling
`arg_process.get_estimate`.
"""
py_typecheck.check_type(transform_fn, computation_base.Computation)
py_typecheck.check_type(arg_process, EstimationProcess)
arg_process_estimate_type = arg_process.get_estimate.type_signature.result
transform_fn_arg_type = transform_fn.type_signature.parameter
if not transform_fn_arg_type.is_assignable_from(arg_process_estimate_type):
raise errors.TemplateStateNotAssignableError(
f'The return type of `get_estimate` of `arg_process` must be '
f'assignable to the input argument of `transform_fn`, but '
f'`get_estimate` returns type:\n{arg_process_estimate_type}\n'
f'and the argument of `transform_fn` is:\n'
f'{transform_fn_arg_type}')
transformed_estimate_fn = computations.tf_computation(
lambda state: transform_fn(arg_process.get_estimate(state)),
arg_process.state_type)
return EstimationProcess(initialize_fn=arg_process.initialize,
next_fn=arg_process.next,
get_estimate_fn=transformed_estimate_fn)
def __init__(self,
initialize_fn: computation_base.Computation,
next_fn: computation_base.Computation,
next_is_multi_arg: Optional[bool] = None):
"""Creates a `tff.templates.IterativeProcess`.
Args:
initialize_fn: A no-arg `tff.Computation` that returns the initial state
of the iterative process. Let the type of this state be called `S`.
next_fn: A `tff.Computation` that represents the iterated function. The
first or only argument must match the state type `S`. The first or only
return value must also match state type `S`.
next_is_multi_arg: An optional boolean indicating that `next_fn` will
receive more than just the state argument (if `True`) or only the state
argument (if `False`). This parameter is primarily used to provide
better error messages.
Raises:
TypeError: If `initialize_fn` and `next_fn` are not instances of
`tff.Computation`.
TemplateInitFnParamNotEmptyError: If `initialize_fn` has any input
arguments.
TemplateStateNotAssignableError: If the `state` returned by either
`initialize_fn` or `next_fn` is not assignable to the first input
argument of `next_fn`.
"""
py_typecheck.check_type(initialize_fn, computation_base.Computation)
if initialize_fn.type_signature.parameter is not None:
raise errors.TemplateInitFnParamNotEmptyError(
f'Provided `initialize_fn` must be a no-arg function, but found '
f'input argument(s) {initialize_fn.type_signature.parameter}.')
initialize_result_type = initialize_fn.type_signature.result
py_typecheck.check_type(next_fn, computation_base.Computation)
next_parameter_type = next_fn.type_signature.parameter
state_type = _infer_state_type(initialize_result_type,
next_parameter_type, next_is_multi_arg)
next_result_type = next_fn.type_signature.result
if state_type.is_assignable_from(next_result_type):
# The whole return value is the state type
pass
elif (_is_nonempty_struct(next_result_type)
and state_type.is_assignable_from(next_result_type[0])):
# The first return value is state type
pass
else:
raise errors.TemplateStateNotAssignableError(
f'The first return argument of `next_fn` must be '
f'assignable to its first input argument, but found\n'
f'`next_fn` which returns type:\n{next_result_type}\n'
f'which does not match its first input argument:\n{state_type}'
)
self._state_type = state_type
self._initialize_fn = initialize_fn
self._next_fn = next_fn
def _infer_state_type(initialize_result_type, next_parameter_type,
next_is_multi_arg):
"""Infers the state type from the `initialize` and `next` types."""
if next_is_multi_arg is None:
# `state_type` may be `next_parameter_type` or
# `next_parameter_type[0]`, depending on which one was assignable from
# `initialize_result_type`.
if next_parameter_type.is_assignable_from(initialize_result_type):
return next_parameter_type
if (_is_nonempty_struct(next_parameter_type)
and next_parameter_type[0].is_assignable_from(
initialize_result_type)):
return next_parameter_type[0]
raise errors.TemplateStateNotAssignableError(
'The return type of `initialize_fn` must be assignable to either\n'
'the whole argument to `next_fn` or the first argument to `next_fn`,\n'
'but found `initialize_fn` return type:\n'
f'{initialize_result_type}\n'
'and `next_fn` with whole argument type:\n'
f'{next_parameter_type}')
elif next_is_multi_arg:
if not _is_nonempty_struct(next_parameter_type):
raise errors.TemplateNextFnNumArgsError(
'Expected `next_parameter_type` to be a structure type of at least '
f'length one, but found type:\n{next_parameter_type}')
if next_parameter_type[0].is_assignable_from(initialize_result_type):
return next_parameter_type[0]
raise errors.TemplateStateNotAssignableError(
'The return type of `initialize_fn` must be assignable to the first\n'
'argument to `next_fn`, but found `initialize_fn` return type:\n'
f'{initialize_result_type}\n'
'and `next_fn` whose first argument type is:\n'
f'{next_parameter_type}')
else:
# `next_is_multi_arg` is `False`
if next_parameter_type.is_assignable_from(initialize_result_type):
return next_parameter_type
raise errors.TemplateStateNotAssignableError(
'The return type of `initialize_fn` must be assignable to the whole\n'
'argument to `next_fn`, but found `initialize_fn` return type:\n'
f'{initialize_result_type}\n'
'and `next_fn` whose first argument type is:\n'
f'{next_parameter_type}')
def __init__(self, initialize_fn: computation_base.Computation,
next_fn: computation_base.Computation):
"""Creates a `tff.templates.MeasuredProcess`.
Args:
initialize_fn: A no-arg `tff.Computation` that creates the initial state
of the measured process.
next_fn: A `tff.Computation` that defines an iterated function. If
`initialize_fn` returns a non-federated type `S`, then `next_fn` must
return a `MeasuredProcessOutput` where the `state` attribute matches the
non-federated type `S`, and accept either a single argument of
non-federated type `S` or multiple arguments where the first argument
must be of non-federated type `S`.
Raises:
TypeError: If `initialize_fn` and `next_fn` are not instances of
`tff.Computation`.
TemplateInitFnParamNotEmptyError: If `initialize_fn` has any input
arguments.
TemplateStateNotAssignableError: If the `state` returned by either
`initialize_fn` or `next_fn` is not assignable to the first input
argument of `next_fn`.
TemplateNotMeasuredProcessOutputError: If `next_fn` does not return a
`MeasuredProcessOutput`.
"""
super().__init__(initialize_fn, next_fn)
next_result_type = next_fn.type_signature.result
if not (isinstance(next_result_type,
computation_types.StructWithPythonType) and
next_result_type.python_container is MeasuredProcessOutput):
raise errors.TemplateNotMeasuredProcessOutputError(
f'The `next_fn` of a `MeasuredProcess` must return a '
f'`MeasuredProcessOutput` object, but returns {next_result_type!r}'
)
# Perform a more strict type check on state than the base class. Base class
# ensures that state returned by initialize_fn is accepted as input argument
# of next_fn, and that this is in the returned structure. For
# MeasuredProcess, this explicitly needs to be in the state attribute. See
# `test_measured_process_output_as_state_raises` for an example.
if next_fn.type_signature.parameter.is_assignable_from(
initialize_fn.type_signature.result):
state_type = next_fn.type_signature.parameter
else:
state_type = next_fn.type_signature.parameter[0]
if not state_type.is_assignable_from(
next_fn.type_signature.result.state):
raise errors.TemplateStateNotAssignableError(
f'The state attrubute of returned MeasuredProcessOutput must be '
f'assignable to its first input argument, but found\n'
f'`next_fn` which returns MeasuredProcessOutput with state attribute '
f'of type:\n{next_result_type}\n'
f'which does not match its first input argument:\n{state_type}'
)
def __init__(self,
initialize_fn: computation_base.Computation,
next_fn: computation_base.Computation,
next_is_multi_arg: Optional[bool] = None):
"""Creates a `tff.templates.MeasuredProcess`.
Args:
initialize_fn: A no-arg `tff.Computation` that returns the initial state
of the measured process. Let the type of this state be called `S`.
next_fn: A `tff.Computation` that represents the iterated function. The
first or only argument must match the state type `S`. The return value
must be a `MeasuredProcessOutput` whose `state` member matches the
state type `S`.
next_is_multi_arg: An optional boolean indicating that `next_fn` will
receive more than just the state argument (if `True`) or only the state
argument (if `False`). This parameter is primarily used to provide
better error messages.
Raises:
TypeError: If `initialize_fn` and `next_fn` are not instances of
`tff.Computation`.
TemplateInitFnParamNotEmptyError: If `initialize_fn` has any input
arguments.
TemplateStateNotAssignableError: If the `state` returned by either
`initialize_fn` or `next_fn` is not assignable to the first input
argument of `next_fn`.
TemplateNotMeasuredProcessOutputError: If `next_fn` does not return a
`MeasuredProcessOutput`.
"""
super().__init__(initialize_fn, next_fn, next_is_multi_arg)
next_result_type = next_fn.type_signature.result
if not (isinstance(next_result_type,
computation_types.StructWithPythonType) and
next_result_type.python_container is MeasuredProcessOutput):
raise errors.TemplateNotMeasuredProcessOutputError(
f'The `next_fn` of a `MeasuredProcess` must return a '
f'`MeasuredProcessOutput` object, but returns {next_result_type!r}'
)
# Perform a more strict type check on state than the base class. Base class
# ensures that state returned by initialize_fn is accepted as input argument
# of next_fn, and that this is in the returned structure. For
# MeasuredProcess, this explicitly needs to be in the state attribute. See
# `test_measured_process_output_as_state_raises` for an example.
state_type = self.state_type
if not state_type.is_assignable_from(
next_fn.type_signature.result.state):
raise errors.TemplateStateNotAssignableError(
f'The state attrubute of returned MeasuredProcessOutput must be '
f'assignable to its first input argument, but found\n'
f'`next_fn` which returns MeasuredProcessOutput with state attribute '
f'of type:\n{next_result_type}\n'
f'which does not match its first input argument:\n{state_type}'
)
def __init__(self, initialize_fn: computation_base.Computation,
next_fn: computation_base.Computation,
get_estimate_fn: computation_base.Computation):
super().__init__(initialize_fn, next_fn)
py_typecheck.check_type(get_estimate_fn, computation_base.Computation)
estimate_fn_arg_type = get_estimate_fn.type_signature.parameter
if not estimate_fn_arg_type.is_assignable_from(self.state_type):
raise errors.TemplateStateNotAssignableError(
f'The state type of the process must be assignable to the '
f'input argument of `get_estimate_fn`, but the state type is: '
f'{self.state_type}\n'
f'and the argument of `get_estimate_fn` is:\n'
f'{estimate_fn_arg_type}')
self._get_estimate_fn = get_estimate_fn
def __init__(self,
initialize_fn: computation_base.Computation,
next_fn: computation_base.Computation,
report_fn: computation_base.Computation,
next_is_multi_arg: Optional[bool] = None):
"""Creates a `tff.templates.EstimationProcess`.
Args:
initialize_fn: A no-arg `tff.Computation` that returns the initial state
of the estimation process. Let the type of this state be called `S`.
next_fn: A `tff.Computation` that represents the iterated function. The
first or only argument must match the state type `S`. The first or only
return value must also match state type `S`.
report_fn: A `tff.Computation` that represents the estimation based on
state. Its input argument must match the state type `S`.
next_is_multi_arg: An optional boolean indicating that `next_fn` will
receive more than just the state argument (if `True`) or only the state
argument (if `False`). This parameter is primarily used to provide
better error messages.
Raises:
TypeError: If `initialize_fn`, `next_fn` and `report_fn` are not
instances of `tff.Computation`.
TemplateInitFnParamNotEmptyError: If `initialize_fn` has any input
arguments.
TemplateStateNotAssignableError: If the `state` returned by either
`initialize_fn` or `next_fn` is not assignable to the first input
argument of `next_fn` and `report_fn`.
"""
super().__init__(initialize_fn,
next_fn,
next_is_multi_arg=next_is_multi_arg)
py_typecheck.check_type(report_fn, computation_base.Computation)
report_fn_arg_type = report_fn.type_signature.parameter
if not report_fn_arg_type.is_assignable_from(self.state_type):
raise errors.TemplateStateNotAssignableError(
f'The state type of the process must be assignable to the '
f'input argument of `report_fn`, but the state type is: '
f'{self.state_type}\n'
f'and the argument of `report_fn` is:\n'
f'{report_fn_arg_type}')
self._report_fn = report_fn
def __init__(self, initialize_fn: computation_base.Computation,
next_fn: computation_base.Computation,
report_fn: computation_base.Computation):
"""Creates a `tff.templates.EstimationProcess`.
Args:
initialize_fn: A no-arg `tff.Computation` that creates the initial state
of the computation.
next_fn: A `tff.Computation` that represents the iterated function. If
`initialize_fn` returns a type `T`, then `next_fn` must either return a
type `U` which is compatible with `T` or multiple values where the first
type is `U`, and accept either a single argument of type `U` or multiple
arguments where the first argument must be of type `U`.
report_fn: A `tff.Computation` that represents the estimation based on
state. Its input argument must be assignable from return type of
`initialize_fn`.
Raises:
TypeError: If `initialize_fn`, `next_fn` and `report_fn` are not
instances of `tff.Computation`.
TemplateInitFnParamNotEmptyError: If `initialize_fn` has any input
arguments.
TemplateStateNotAssignableError: If the `state` returned by either
`initialize_fn` or `next_fn` is not assignable to the first input
argument of `next_fn` and `report_fn`.
"""
super().__init__(initialize_fn, next_fn)
py_typecheck.check_type(report_fn, computation_base.Computation)
report_fn_arg_type = report_fn.type_signature.parameter
if not report_fn_arg_type.is_assignable_from(self.state_type):
raise errors.TemplateStateNotAssignableError(
f'The state type of the process must be assignable to the '
f'input argument of `report_fn`, but the state type is: '
f'{self.state_type}\n'
f'and the argument of `report_fn` is:\n'
f'{report_fn_arg_type}')
self._report_fn = report_fn
def __init__(self, initialize_fn: computation_base.Computation,
next_fn: computation_base.Computation):
"""Creates a `tff.templates.IterativeProcess`.
Args:
initialize_fn: A no-arg `tff.Computation` that creates the initial state
of the computation.
next_fn: A `tff.Computation` that represents the iterated function. If
`initialize_fn` returns a type `T`, then `next_fn` must either return a
type `U` which is compatible with `T` or multiple values where the first
type is `U`, and accept either a single argument of type `U` or multiple
arguments where the first argument must be of type `U`.
Raises:
TypeError: If `initialize_fn` and `next_fn` are not instances of
`tff.Computation`.
TemplateInitFnParamNotEmptyError: If `initialize_fn` has any input
arguments.
TemplateStateNotAssignableError: If the `state` returned by either
`initialize_fn` or `next_fn` is not assignable to the first input
argument of `next_fn`.
"""
py_typecheck.check_type(initialize_fn, computation_base.Computation)
if initialize_fn.type_signature.parameter is not None:
raise errors.TemplateInitFnParamNotEmptyError(
f'Provided `initialize_fn` must be a no-arg function, but found '
f'input argument(s) {initialize_fn.type_signature.parameter}.')
initialize_result_type = initialize_fn.type_signature.result
py_typecheck.check_type(next_fn, computation_base.Computation)
next_parameter_type = next_fn.type_signature.parameter
# `next_first_parameter_type` may be `next_parameter_type` or
# `next_parameter_type[0]`, depending on which one was assignable from
# `initialize_result_type`.
if next_parameter_type.is_assignable_from(initialize_result_type):
# The only argument is the state type
state_type = next_parameter_type
elif (next_parameter_type.is_struct() and next_parameter_type and
next_parameter_type[0].is_assignable_from(initialize_result_type)):
# The first argument is the state type
state_type = next_parameter_type[0]
else:
raise errors.TemplateStateNotAssignableError(
f'The return type of `initialize_fn` must be assignable to '
f'the first input argument of `next_fn`, but:\n'
f'`initialize_fn` returned type:\n{initialize_result_type}\n'
f'and the first input argument of `next_fn` is:\n'
f'{next_parameter_type}')
next_result_type = next_fn.type_signature.result
if state_type.is_assignable_from(next_result_type):
# The whole return value is the state type
pass
elif (next_result_type.is_struct() and next_result_type and
state_type.is_assignable_from(next_result_type[0])):
# The first return value is state type
pass
else:
raise errors.TemplateStateNotAssignableError(
f'The first return argument of `next_fn` must be '
f'assignable to its first input argument, but found\n'
f'`next_fn` which returns type:\n{next_result_type}\n'
f'which does not match its first input argument:\n{state_type}')
self._state_type = state_type
self._initialize_fn = initialize_fn
self._next_fn = next_fn