def test_returns_computation(self, operator, type_signature, operands,
expected_result):
proto, _ = tensorflow_computation_factory.create_binary_operator_with_upcast(
type_signature, operator)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
self.assertIsInstance(actual_type, computation_types.FunctionType)
# Note: It is only useful to test the parameter type; the result type
# depends on the `operator` used, not the implemenation
# `create_binary_operator_with_upcast`.
expected_parameter_type = computation_types.StructType(type_signature)
self.assertEqual(actual_type.parameter, expected_parameter_type)
actual_result = test_utils.run_tensorflow(proto, operands)
self.assertEqual(actual_result, expected_result)
async def embed_tf_binary_operator_with_upcast(executor, type_spec, op):
"""Embeds a binary operator `op` on `type_spec`-typed values in `executor`.
Args:
executor: An instance of `tff.framework.Executor`.
type_spec: An instance of `tff.StructType` with two elements, the types of
the first and second argument to the binary operator.
op: An operator function (such as `tf.add` or `tf.multiply`) to apply to the
tensor-level constituents of the values, pointwise.
Returns:
An instance of `tff.framework.ExecutorValue` representing the operator in
a form embedded into the executor.
"""
proto, type_signature = tensorflow_computation_factory.create_binary_operator_with_upcast(
operator=op, type_signature=type_spec)
return await executor.create_value(proto, type_signature)
def test_returns_computation(self, operator, type_signature, operands,
expected_result):
# TODO(b/142795960): arguments in parameterized are called before test main.
# `tf.constant` will error out on GPU and TPU without proper initialization.
# A suggested workaround is to use numpy as argument and transform to TF
# tensor inside the function.
operands = tf.nest.map_structure(tf.constant, operands)
proto, _ = tensorflow_computation_factory.create_binary_operator_with_upcast(
type_signature, operator)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
self.assertIsInstance(actual_type, computation_types.FunctionType)
# Note: It is only useful to test the parameter type; the result type
# depends on the `operator` used, not the implemenation
# `create_binary_operator_with_upcast`.
expected_parameter_type = computation_types.StructType(type_signature)
self.assertEqual(actual_type.parameter, expected_parameter_type)
actual_result = test_utils.run_tensorflow(proto, operands)
self.assertEqual(actual_result, expected_result)
async def _compute_modulus(self, value, mask):
async def build_modulus_argument(value, mask):
# Create the mask at the same placement as value.
placed_mask = await self._executor.create_value(
await mask.compute(),
computation_types.FederatedType(mask.type_signature,
value.type_signature.placement,
all_equal=True))
arg_struct = await self._executor.create_struct(
[value, placed_mask])
if value.type_signature.placement == placements.SERVER:
return await self.compute_federated_zip_at_server(arg_struct)
elif value.type_signature.placement == placements.CLIENTS:
return await self.compute_federated_zip_at_clients(arg_struct)
else:
raise TypeError(
'Unknown placement [{p}], must be one of [CLIENTS, SERVER]'
.format(p=value.type_signature.placement))
modulus_comp_coro = self._executor.create_value(
*tensorflow_computation_factory.create_binary_operator_with_upcast(
computation_types.StructType([
value.type_signature.member, mask.type_signature
]), tf.bitwise.bitwise_and))
modulus_comp, modulus_comp_arg = await asyncio.gather(
modulus_comp_coro, build_modulus_argument(value, mask))
map_arg = federated_resolving_strategy.FederatedResolvingStrategyValue(
structure.Struct([
(None, modulus_comp.internal_representation),
(None, modulus_comp_arg.internal_representation),
]),
computation_types.StructType(
[modulus_comp.type_signature,
modulus_comp_arg.type_signature]))
if value.type_signature.all_equal:
return await self.compute_federated_map_all_equal(map_arg)
else:
return await self.compute_federated_map(map_arg)
def test_fails(self, operator, type_signature, operands):
operands = tf.nest.map_structure(tf.constant, operands)
with self.assertRaises(TypeError):
tensorflow_computation_factory.create_binary_operator_with_upcast(
type_signature, operator)
