-
Notifications
You must be signed in to change notification settings - Fork 581
/
estimation_process.py
147 lines (121 loc) · 5.68 KB
/
estimation_process.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
# Copyright 2020, The TensorFlow Federated Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Defines a template for a process that can compute an estimate."""
from typing import Optional
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.core.impl.computation import computation_base
from tensorflow_federated.python.core.impl.federated_context import federated_computation
from tensorflow_federated.python.core.templates import errors
from tensorflow_federated.python.core.templates import iterative_process
class EstimateNotAssignableError(TypeError):
"""`TypeError` for estimate not being assignable to expected type."""
pass
class EstimationProcess(iterative_process.IterativeProcess):
"""A stateful process that can compute an estimate of some value.
This class inherits the constraints documented by
`tff.templates.IterativeProcess`.
A `tff.templates.EstimationProcess` is an `tff.templates.IterativeProcess`
that in addition to the `initialize` and `next` functions, has a
`report` function that returns the result of some computation based on the
state of the process. The argument of `report` must be of the same type as the
state, that is, the type of object returned by `initialize`.
"""
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 report_fn_arg_type is None or not report_fn_arg_type.is_assignable_from(
self.state_type
):
raise errors.TemplateStateNotAssignableError(
'The state type of the process must be assignable to the '
'input argument of `report_fn`, but the state type is: '
f'{self.state_type}\n'
'and the argument of `report_fn` is:\n'
f'{report_fn_arg_type}'
)
self._report_fn = report_fn
@property
def report(self) -> computation_base.Computation:
"""A `tff.Computation` that computes the current estimate from `state`.
Given a `state` controlled by this process, computes and returns the most
recent estimate of the estimated value.
Returns:
A `tff.Computation`.
"""
return self._report_fn
def map(self, map_fn: computation_base.Computation):
"""Applies `map_fn` to the estimate function of the process.
This method will return a new instance of `EstimationProcess` with the same
`initailize` and `next` functions, and its `report` function replaced by
`map_fn(report(state))`.
Args:
map_fn: A `tff.Computation` to apply to the result of the `report`
function of the process. Must accept the return type of `report`.
Returns:
An `EstimationProcess`.
Raises:
EstimateNotAssignableError: If the return type of `report` is not
assignable to the expected input type of `map_fn`.
"""
py_typecheck.check_type(map_fn, computation_base.Computation)
estimate_type = self.report.type_signature.result
map_fn_arg_type = map_fn.type_signature.parameter
if map_fn_arg_type is None or not map_fn_arg_type.is_assignable_from(
estimate_type
):
raise EstimateNotAssignableError(
'The return type of `report` of this process must be '
'assignable to the input argument of `map_fn`, but '
f'`report` returns type:\n{estimate_type}\n'
f'and the argument of `map_fn` is:\n{map_fn_arg_type}'
)
transformed_report_fn = federated_computation.federated_computation(
lambda state: map_fn(self.report(state)), self.state_type
)
return EstimationProcess(
initialize_fn=self.initialize,
next_fn=self.next,
report_fn=transformed_report_fn,
)