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)