def test_federated_zip_named_n_tuple(self, n, element_type):
fed_type = computation_types.at_clients(element_type)
initial_tuple_type = computation_types.to_type([fed_type] * n)
initial_tuple = _mock_data_of_type(initial_tuple_type)
naming_fn = str
named_result = intrinsics.federated_zip(
collections.OrderedDict(
(naming_fn(i), initial_tuple[i]) for i in range(n)))
self.assertIsInstance(named_result, value_impl.Value)
expected = computation_types.at_clients(
collections.OrderedDict(
(naming_fn(i), element_type) for i in range(n)))
type_test_utils.assert_types_identical(named_result.type_signature,
expected)
naming_fn = lambda i: str(i) if i % 2 == 0 else None
mixed_result = intrinsics.federated_zip(
structure.Struct(
(naming_fn(i), initial_tuple[i]) for i in range(n)))
self.assertIsInstance(mixed_result, value_impl.Value)
expected = computation_types.at_clients(
computation_types.StructType([(naming_fn(i), element_type)
for i in range(n)]))
type_test_utils.assert_types_identical(mixed_result.type_signature,
expected)
def test_zero_preserves_aggregated_dtype_with_mixed_float(
self, norm_order, type_spec):
factory = _zeroed_sum(norm_order=norm_order)
mixed_float = computation_types.to_type(type_spec)
process = factory.create(mixed_float)
type_test_utils.assert_types_identical(
mixed_float, process.next.type_signature.result.result.member)
def test_federated_zip_n_tuple(self, n):
fed_type = computation_types.at_clients(tf.int32)
x = _mock_data_of_type([fed_type] * n)
val = intrinsics.federated_zip(x)
self.assertIsInstance(val, value_impl.Value)
expected = computation_types.at_clients([tf.int32] * n)
type_test_utils.assert_types_identical(val.type_signature, expected)
def test_scalar_random_seed(self):
example_type = TensorType(tf.int32)
sample_computation = sampling._build_sample_value_computation(
example_type, sample_size=1)
reservoir_type = sampling._build_reservoir_type(example_type)
expected_type = FunctionType(parameter=collections.OrderedDict(
reservoir=reservoir_type, sample=example_type),
result=reservoir_type)
type_test_utils.assert_types_identical(
sample_computation.type_signature, expected_type)
# Get the sentinel seed so that the first call initializes based on
# timestamp.
reservoir = sampling._build_initial_sample_reservoir(example_type)
self.assertAllEqual(reservoir['random_seed'],
[sampling.SEED_SENTINEL, sampling.SEED_SENTINEL])
reservoir = sample_computation(reservoir, 1)
# The first value of the seed was the timestamp, it should be greater than
# 1_600_000_000_000 (September 2020) and within 60 seconds of now.
self.assertGreater(reservoir['random_seed'][0], 1_600_000_000_000)
self.assertLess(
tf.cast(tf.timestamp() * 1000.0, tf.int64) -
reservoir['random_seed'][0], 60)
# The second value should we a random number. We assert its not the
# sentinel, though it ccould be with probability 1 / 2**32.
self.assertNotEqual(reservoir['random_seed'][1],
sampling.SEED_SENTINEL)
def test_scalar_fixed_seed(self):
example_type = TensorType(tf.int32)
sample_computation = sampling._build_sample_value_computation(
example_type, sample_size=1)
reservoir_type = sampling._build_reservoir_type(example_type)
expected_type = FunctionType(parameter=collections.OrderedDict(
reservoir=reservoir_type, sample=example_type),
result=reservoir_type)
type_test_utils.assert_types_identical(
sample_computation.type_signature, expected_type)
reservoir = sampling._build_initial_sample_reservoir(example_type,
seed=TEST_SEED)
reservoir = sample_computation(reservoir, 1)
self.assertAllEqual(
reservoir,
collections.OrderedDict(random_seed=(TEST_SEED, 100565241),
random_values=[100565241],
samples=[1]))
# New value was not sampled, its random value was too low, but it
# changes the seed for the next iteration.
reservoir = sample_computation(reservoir, 2)
self.assertAllEqual(
reservoir,
collections.OrderedDict(random_seed=(TEST_SEED, -1479562987),
random_values=[100565241],
samples=[1]))
# The PRNG doesn't generate a number for sampling until 5th example.
for i in range(3, 6):
reservoir = sample_computation(reservoir, i)
self.assertAllEqual(
reservoir,
collections.OrderedDict(random_seed=(TEST_SEED, 756274747),
random_values=[756274747],
samples=[5]))
def test_federated_secure_select(self):
uri = intrinsic_defs.FEDERATED_SECURE_SELECT.uri
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType(
[
computation_types.at_clients(tf.int32), # client_keys
computation_types.at_server(tf.int32), # max_key
computation_types.at_server(tf.float32), # server_state
computation_types.FunctionType([tf.float32, tf.int32],
tf.float32) # select_fn
],
computation_types.at_clients(
computation_types.SequenceType(tf.float32))))
self.assertGreater(_count_intrinsics(comp, uri), 0)
# First without secure intrinsics shouldn't modify anything.
reduced, modified = tree_transformations.replace_intrinsics_with_bodies(
comp)
self.assertFalse(modified)
self.assertGreater(_count_intrinsics(comp, uri), 0)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
# Now replace bodies including secure intrinsics.
reduced, modified = tree_transformations.replace_secure_intrinsics_with_insecure_bodies(
comp)
self.assertTrue(modified)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
self.assertGreater(
_count_intrinsics(reduced, intrinsic_defs.FEDERATED_SELECT.uri), 0)
def test_federated_secure_modular_sum(self, value_dtype, modulus_type):
uri = intrinsic_defs.FEDERATED_SECURE_MODULAR_SUM.uri
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType(
parameter=[
computation_types.at_clients(value_dtype),
computation_types.to_type(modulus_type)
],
result=computation_types.at_server(value_dtype)))
# First without secure intrinsics shouldn't modify anything.
reduced, modified = tree_transformations.replace_intrinsics_with_bodies(
comp)
self.assertFalse(modified)
self.assertGreater(_count_intrinsics(comp, uri), 0)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
# Now replace bodies including secure intrinsics.
reduced, modified = tree_transformations.replace_secure_intrinsics_with_insecure_bodies(
comp)
self.assertTrue(modified)
# Inserting tensorflow, as we do here, does not preserve python containers
# currently.
type_test_utils.assert_types_equivalent(comp.type_signature,
reduced.type_signature)
self.assertGreater(
_count_intrinsics(reduced, intrinsic_defs.FEDERATED_SUM.uri), 0)
def test_federated_secure_sum_bitwidth(self, value_dtype, bitwidth_type):
uri = intrinsic_defs.FEDERATED_SECURE_SUM_BITWIDTH.uri
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType(
parameter=[
computation_types.at_clients(value_dtype),
computation_types.to_type(bitwidth_type)
],
result=computation_types.at_server(value_dtype)))
# First without secure intrinsics shouldn't modify anything.
reduced, modified = tree_transformations.replace_intrinsics_with_bodies(
comp)
self.assertFalse(modified)
self.assertGreater(_count_intrinsics(comp, uri), 0)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
# Now replace bodies including secure intrinsics.
reduced, modified = tree_transformations.replace_secure_intrinsics_with_insecure_bodies(
comp)
self.assertTrue(modified)
type_test_utils.assert_types_identical(comp.type_signature,
reduced.type_signature)
self.assertGreater(
_count_intrinsics(reduced, intrinsic_defs.FEDERATED_AGGREGATE.uri),
0)
def test_deserialize_federated_value_promotes_types(self):
x = [10]
smaller_type = computation_types.StructType([
(None, computation_types.to_type(tf.int32))
])
smaller_type_member_proto, _ = value_serialization.serialize_value(
x, smaller_type)
larger_type = computation_types.StructType([
('a', computation_types.to_type(tf.int32))
])
larger_type_member_proto, _ = value_serialization.serialize_value(
x, larger_type)
type_at_clients = type_serialization.serialize_type(
computation_types.at_clients(tf.int32))
unspecified_member_federated_type = computation_pb2.FederatedType(
placement=type_at_clients.federated.placement, all_equal=False)
federated_proto = executor_pb2.Value.Federated(
type=unspecified_member_federated_type,
value=[larger_type_member_proto, smaller_type_member_proto])
federated_value_proto = executor_pb2.Value(federated=federated_proto)
_, deserialized_type_spec = value_serialization.deserialize_value(
federated_value_proto)
type_test_utils.assert_types_identical(
deserialized_type_spec, computation_types.at_clients(larger_type))
def test_create_tuple_of_value_sequence(self):
datasets = (tf.data.Dataset.range(5), tf.data.Dataset.range(5))
executor = executor_bindings.create_tensorflow_executor()
struct_of_sequence_type = StructType([
(None, SequenceType(datasets[0].element_spec)),
(None, SequenceType(datasets[0].element_spec))
])
arg_value_pb, _ = value_serialization.serialize_value(
datasets, struct_of_sequence_type)
arg = executor.create_value(arg_value_pb)
@tensorflow_computation.tf_computation(struct_of_sequence_type)
def preprocess(datasets):
def double_value(x):
return 2 * x
@tf.function
def add_preprocessing(ds1, ds2):
return ds1.map(double_value), ds2.map(double_value)
return add_preprocessing(*datasets)
comp_pb = executor_pb2.Value(
computation=preprocess.get_proto(preprocess))
comp = executor.create_value(comp_pb)
result = executor.create_call(comp.ref, arg.ref)
output_pb = executor.materialize(result.ref)
result, result_type_spec = value_serialization.deserialize_value(
output_pb, type_hint=struct_of_sequence_type)
type_test_utils.assert_types_identical(result_type_spec,
struct_of_sequence_type)
def test_serialize_deserialize_nested_tuple_value_without_names(self): x = (10, 20) x_type = computation_types.to_type((tf.int32, tf.int32)) value_proto, value_type = value_serialization.serialize_value(x, x_type) type_test_utils.assert_types_identical(value_type, x_type) y, type_spec = value_serialization.deserialize_value(value_proto) type_test_utils.assert_types_equivalent(type_spec, x_type) self.assertEqual(y, structure.from_container((10, 20)))
def test_serialize_deserialize_variable_as_tensor_value(self): x = tf.Variable(10.0) type_spec = TensorType(tf.as_dtype(x.dtype), x.shape) value_proto, value_type = value_serialization.serialize_value(x, type_spec) type_test_utils.assert_types_identical(value_type, TensorType(tf.float32)) y, type_spec = value_serialization.deserialize_value(value_proto) type_test_utils.assert_types_identical(type_spec, TensorType(tf.float32)) self.assertAllEqual(x, y)
def test_serialize_deserialize_tensor_value_with_different_dtype(self):
x = tf.constant(10)
value_proto, value_type = value_serialization.serialize_value(
x, TensorType(tf.float32))
type_test_utils.assert_types_identical(value_type, TensorType(tf.float32))
y, type_spec = value_serialization.deserialize_value(value_proto)
type_test_utils.assert_types_identical(type_spec, TensorType(tf.float32))
self.assertEqual(y, 10.0)
def test_serialize_deserialize_tensor_value_without_hint(
self, x, serialize_type_spec):
value_proto, value_type = value_serialization.serialize_value(
x, serialize_type_spec)
type_test_utils.assert_types_identical(value_type, serialize_type_spec)
y, type_spec = value_serialization.deserialize_value(value_proto)
type_test_utils.assert_types_identical(type_spec, serialize_type_spec)
self.assertEqual(y.dtype, serialize_type_spec.dtype.as_numpy_dtype)
self.assertAllEqual(x, y)
def test_serialize_deserialize_federated_at_server(self):
x = 10
x_type = computation_types.at_server(tf.int32)
value_proto, value_type = value_serialization.serialize_value(x, x_type)
type_test_utils.assert_types_identical(
value_type, computation_types.at_server(tf.int32))
y, type_spec = value_serialization.deserialize_value(value_proto)
type_test_utils.assert_types_identical(type_spec, x_type)
self.assertEqual(y, 10)
def test_preserves_python_containers_in_after_merge(self):
mergeable_form = mergeable_comp_compiler.compile_to_mergeable_comp_form(
return_list)
self.assertIsInstance(
mergeable_form, mergeable_comp_execution_context.MergeableCompForm)
type_test_utils.assert_types_identical(
mergeable_form.after_merge.type_signature.result,
return_list.type_signature.result)
def test_serialize_deserialize_tensor_value_with_nontrivial_shape(self):
x = tf.constant([10, 20, 30])
value_proto, value_type = value_serialization.serialize_value(
x, TensorType(tf.int32, [3]))
type_test_utils.assert_types_identical(value_type,
TensorType(tf.int32, [3]))
y, type_spec = value_serialization.deserialize_value(value_proto)
type_test_utils.assert_types_identical(type_spec, TensorType(tf.int32, [3]))
self.assertAllEqual(x, y)
def test_scalar(self):
example_type = TensorType(tf.int32)
merge_computation = sampling._build_merge_samples_computation(
example_type, sample_size=5)
reservoir_type = sampling._build_reservoir_type(example_type)
expected_type = FunctionType(parameter=collections.OrderedDict(
a=reservoir_type, b=reservoir_type),
result=reservoir_type)
type_test_utils.assert_types_identical(
merge_computation.type_signature, expected_type)
reservoir_a = sampling._build_initial_sample_reservoir(example_type,
seed=TEST_SEED)
reservoir_a['random_values'] = [1, 3, 5]
reservoir_a['samples'] = [3, 9, 15]
with self.subTest('downsample'):
reservoir_b = sampling._build_initial_sample_reservoir(
example_type, seed=TEST_SEED + 1)
reservoir_b['random_values'] = [2, 4, 6, 8]
reservoir_b['samples'] = [6, 12, 18, 24]
merged_reservoir = merge_computation(reservoir_a, reservoir_b)
self.assertAllEqual(
merged_reservoir,
collections.OrderedDict(
# Arbitrarily take seeds from `a`, discarded later.
random_seed=tf.convert_to_tensor((TEST_SEED, TEST_SEED)),
random_values=[3, 5, 4, 6, 8],
samples=[9, 15, 12, 18, 24]))
with self.subTest('keep_all'):
reservoir_b = sampling._build_initial_sample_reservoir(
example_type, seed=TEST_SEED + 1)
reservoir_b['random_values'] = [2]
reservoir_b['samples'] = [6]
# We select the value from reservoir_b because its random_value was
# higher.
merged_reservoir = merge_computation(reservoir_a, reservoir_b)
self.assertAllEqual(
merged_reservoir,
collections.OrderedDict(
# Arbitrarily take seeds from `a`, discarded later.
random_seed=tf.convert_to_tensor((TEST_SEED, TEST_SEED)),
random_values=[1, 3, 5, 2],
samples=[3, 9, 15, 6]))
with self.subTest('tie_breakers'):
# In case of tie, we take the as many values from `a` first.
reservoir_b = sampling._build_initial_sample_reservoir(
example_type, seed=TEST_SEED)
reservoir_b['random_values'] = [5, 5, 5, 5, 5] # all tied with `a`
reservoir_b['samples'] = [-1, -1, -1, -1, -1]
merged_reservoir = merge_computation(reservoir_a, reservoir_b)
self.assertAllEqual(
merged_reservoir,
collections.OrderedDict(random_seed=tf.convert_to_tensor(
(TEST_SEED, TEST_SEED)),
random_values=[5, 5, 5, 5, 5],
samples=[15, -1, -1, -1, -1]))
def test_serialize_deserialize_sequence_of_scalars(self, dataset_fn):
ds = tf.data.Dataset.range(5).map(lambda x: x * 2)
ds_repr = dataset_fn(ds)
value_proto, value_type = value_serialization.serialize_value(
ds_repr, computation_types.SequenceType(tf.int64))
type_test_utils.assert_types_identical(
value_type, computation_types.SequenceType(tf.int64))
y, type_spec = value_serialization.deserialize_value(value_proto)
type_test_utils.assert_types_identical(
type_spec, computation_types.SequenceType(tf.int64))
self.assertAllEqual(list(y), [x * 2 for x in range(5)])
def test_serialize_deserialize_sequence_of_scalars_graph_mode(self):
with tf.Graph().as_default():
ds = tf.data.Dataset.range(5).map(lambda x: x * 2)
value_proto, value_type = value_serialization.serialize_value(
ds, computation_types.SequenceType(tf.int64))
type_test_utils.assert_types_identical(
value_type, computation_types.SequenceType(tf.int64))
y, type_spec = value_serialization.deserialize_value(value_proto)
type_test_utils.assert_types_identical(
type_spec, computation_types.SequenceType(tf.int64))
self.assertAllEqual(list(y), [x * 2 for x in range(5)])
def test_serialize_deserialize_sequence_of_ragged_tensors(self, dataset_fn):
self.skipTest('b/235492749')
ds = tf.data.Dataset.from_tensor_slices(tf.strings.split(['a b c', 'd e']))
ds_repr = dataset_fn(ds)
value_proto, value_type = value_serialization.serialize_value(
ds_repr, computation_types.SequenceType(element=ds.element_spec))
expected_type = computation_types.SequenceType(ds.element_spec)
type_test_utils.assert_types_identical(value_type, expected_type)
_, type_spec = value_serialization.deserialize_value(value_proto)
# Only checking for equivalence, we don't have the Python container
# after deserialization.
type_test_utils.assert_types_equivalent(type_spec, expected_type)
def test_create_value(self):
executor = executor_bindings.create_reference_resolving_executor(
executor_bindings.create_tensorflow_executor())
# 1. Test a simple tensor.
expected_type_spec = TensorType(shape=[3], dtype=tf.int64)
value_pb, _ = value_serialization.serialize_value([1, 2, 3],
expected_type_spec)
value = executor.create_value(value_pb)
self.assertIsInstance(value, executor_bindings.OwnedValueId)
self.assertEqual(value.ref, 0)
self.assertEqual(str(value), '0')
self.assertEqual(repr(value), r'<OwnedValueId: 0>')
materialized_value = executor.materialize(value.ref)
deserialized_value, type_spec = value_serialization.deserialize_value(
materialized_value)
type_test_utils.assert_types_identical(type_spec, expected_type_spec)
self.assertAllEqual(deserialized_value, [1, 2, 3])
# 2. Test a struct of tensors, ensure that we get a different ID.
expected_type_spec = StructType([
('a', TensorType(shape=[3], dtype=tf.int64)),
('b', TensorType(shape=[], dtype=tf.float32))
])
value_pb, _ = value_serialization.serialize_value(
collections.OrderedDict(a=tf.constant([1, 2, 3]),
b=tf.constant(42.0)), expected_type_spec)
value = executor.create_value(value_pb)
self.assertIsInstance(value, executor_bindings.OwnedValueId)
# Assert the value ID was incremented.
self.assertEqual(value.ref, 1)
self.assertEqual(str(value), '1')
self.assertEqual(repr(value), r'<OwnedValueId: 1>')
materialized_value = executor.materialize(value.ref)
deserialized_value, type_spec = value_serialization.deserialize_value(
materialized_value)
# Note: here we've lost the names `a` and `b` in the output. The output
# is a more _strict_ type.
self.assertTrue(expected_type_spec.is_assignable_from(type_spec))
deserialized_value = type_conversions.type_to_py_container(
deserialized_value, expected_type_spec)
self.assertAllClose(deserialized_value,
collections.OrderedDict(a=(1, 2, 3), b=42.0))
# 3. Test creating a value from a computation.
@tensorflow_computation.tf_computation(tf.int32, tf.int32)
def foo(a, b):
return tf.add(a, b)
value_pb, _ = value_serialization.serialize_value(foo)
value = executor.create_value(value_pb)
self.assertIsInstance(value, executor_bindings.OwnedValueId)
# Assert the value ID was incremented again.
self.assertEqual(value.ref, 2)
self.assertEqual(str(value), '2')
self.assertEqual(repr(value), '<OwnedValueId: 2>')
def test_serialize_deserialize_string_value(self, x):
tf_type = tf.as_dtype(x.dtype)
type_spec = TensorType(tf_type, x.shape)
value_proto, value_type = value_serialization.serialize_value(x, type_spec)
type_test_utils.assert_types_identical(value_type,
TensorType(tf_type, x.shape))
y, type_spec = value_serialization.deserialize_value(
value_proto, type_hint=type_spec)
type_test_utils.assert_types_identical(type_spec,
TensorType(tf_type, x.shape))
self.assertIsInstance(y, bytes)
self.assertAllEqual(x, y)
def test_unwrap_removes_federated_zips_at_clients(self):
list_type = computation_types.to_type([tf.int32, tf.float32] * 2)
clients_list_type = computation_types.at_server(list_type)
fed_tuple = building_blocks.Reference('tup', clients_list_type)
unzipped = building_block_factory.create_federated_unzip(fed_tuple)
before = building_block_factory.create_federated_zip(unzipped)
after, modified = tree_transformations.strip_placement(before)
self.assertTrue(modified)
self.assert_has_no_intrinsics_nor_federated_types(after)
type_test_utils.assert_types_identical(before.type_signature,
clients_list_type)
type_test_utils.assert_types_identical(after.type_signature, list_type)
def test_serialize_deserialize_tensor_value_without_hint_graph_mode(
self, x, serialize_type_spec):
# Test serializing non-eager tensors.
with tf.Graph().as_default():
tensor_x = tf.convert_to_tensor(x)
value_proto, value_type = value_serialization.serialize_value(
tensor_x, serialize_type_spec)
type_test_utils.assert_types_identical(value_type, serialize_type_spec)
y, deserialize_type_spec = value_serialization.deserialize_value(
value_proto)
type_test_utils.assert_types_identical(deserialize_type_spec,
serialize_type_spec)
self.assertEqual(y.dtype, serialize_type_spec.dtype.as_numpy_dtype)
self.assertAllEqual(x, y)
def test_federated_zip_n_tuple_mixed_args(self, n, m):
tuple_fed_type = computation_types.at_clients((tf.int32, tf.int32))
single_fed_type = computation_types.at_clients(tf.int32)
tuple_repeat = lambda v, n: (v, ) * n
initial_tuple_type = computation_types.to_type(
tuple_repeat(tuple_fed_type, n) + tuple_repeat(single_fed_type, m))
final_fed_type = computation_types.at_clients(
tuple_repeat((tf.int32, tf.int32), n) + tuple_repeat(tf.int32, m))
x = _mock_data_of_type(initial_tuple_type)
val = intrinsics.federated_zip(x)
self.assertIsInstance(val, value_impl.Value)
type_test_utils.assert_types_identical(val.type_signature,
final_fed_type)
def test_serialize_deserialize_nested_tuple_value_with_names(self):
x = collections.OrderedDict(
a=10, b=[20, 30], c=collections.OrderedDict(d=40))
x_type = computation_types.to_type(
collections.OrderedDict(
a=tf.int32,
b=[tf.int32, tf.int32],
c=collections.OrderedDict(d=tf.int32)))
value_proto, value_type = value_serialization.serialize_value(x, x_type)
type_test_utils.assert_types_identical(value_type, x_type)
y, type_spec = value_serialization.deserialize_value(value_proto)
# Don't assert on the Python container since it is lost in serialization.
type_test_utils.assert_types_equivalent(type_spec, x_type)
self.assertEqual(y, structure.from_container(x, recursive=True))
def test_scalar(self):
example_type = computation_types.to_type(TensorType(tf.int32))
finalize_computation = sampling._build_finalize_sample_computation(
example_type)
reservoir_type = sampling._build_reservoir_type(example_type)
expected_type = FunctionType(parameter=reservoir_type,
result=reservoir_type.samples)
type_test_utils.assert_types_identical(
finalize_computation.type_signature, expected_type)
reservoir = sampling._build_initial_sample_reservoir(example_type,
seed=TEST_SEED)
reservoir['random_values'] = [3, 5, 7]
test_samples = [3, 9, 27]
reservoir['samples'] = test_samples
self.assertAllEqual(finalize_computation(reservoir), test_samples)
def test_strip_placement_with_called_lambda(self):
int_type = computation_types.TensorType(tf.int32)
server_int_type = computation_types.at_server(int_type)
federated_ref = building_blocks.Reference('outer', server_int_type)
inner_federated_ref = building_blocks.Reference(
'inner', server_int_type)
identity_lambda = building_blocks.Lambda('inner', server_int_type,
inner_federated_ref)
before = building_blocks.Call(identity_lambda, federated_ref)
after, modified = tree_transformations.strip_placement(before)
self.assertTrue(modified)
self.assert_has_no_intrinsics_nor_federated_types(after)
type_test_utils.assert_types_identical(before.type_signature,
server_int_type)
type_test_utils.assert_types_identical(after.type_signature, int_type)
def test_serialize_deserialize_computation_value(self):
@tensorflow_computation.tf_computation
def comp():
return tf.constant(10)
value_proto, value_type = value_serialization.serialize_value(comp)
self.assertEqual(value_proto.WhichOneof('value'), 'computation')
type_test_utils.assert_types_identical(
value_type,
computation_types.FunctionType(parameter=None, result=tf.int32))
comp, type_spec = value_serialization.deserialize_value(value_proto)
# self.assertIsInstance(comp, computation_pb2.Computation)
type_test_utils.assert_types_identical(
type_spec,
computation_types.FunctionType(parameter=None, result=tf.int32))
