Add contextual information to bitwidth structure error from federated_secure_sum
# limitations under the License.
"""A factory of intrinsics for use in composing federated computations."""
+import contextlib
+
from tensorflow_federated.python.common_libs import py_typecheck
from tensorflow_federated.python.core.api import computation_types
from tensorflow_federated.python.core.api import placements
@@ -28,6 +30,15 @@ from tensorflow_federated.python.core.impl.compiler import type_factory
from tensorflow_federated.python.core.impl.context_stack import context_stack_base
+@contextlib.contextmanager
+def err_ctx(err_ty, context: str):
+ """A context manager which wraps type errors in a string describing their context."""
+ try:
+ yield
+ except err_ty as err:
+ raise err_ty(context) from err
+
+
class IntrinsicFactory(object):
"""A factory that can constructs intrinsics over the given context stack."""
@@ -332,8 +343,13 @@ class IntrinsicFactory(object):
'value to be summed')
type_utils.check_is_structure_of_integers(value.type_signature)
bitwidth = value_impl.to_value(bitwidth, None, self._context_stack)
- type_utils.check_equivalent_types(value.type_signature.member,
- bitwidth.type_signature)
+ value_member_ty = value.type_signature.member
+ bitwidth_ty = bitwidth.type_signature
+ if not type_utils.are_equivalent_types(value_member_ty, bitwidth_ty):
+ raise TypeError(
+ 'Expected `federated_secure_sum` parameters `value` and `bitwidth` '
+ 'to have the same structure. Found `value` of `{}` and `bitwidth` of `{}`'
+ .format(value_member_ty, bitwidth_ty))
value = value_impl.ValueImpl.get_comp(value)
bitwidth = value_impl.ValueImpl.get_comp(bitwidth)
comp = building_block_factory.create_federated_secure_sum(value, bitwidth)
