def compile(self, computation_to_compile):
"""Compiles `computation_to_compile`.
Args:
computation_to_compile: An instance of `computation_base.Computation` to
compile.
Returns:
An instance of `computation_base.Computation` that repeesents the result.
"""
py_typecheck.check_type(computation_to_compile,
computation_base.Computation)
computation_proto = computation_impl.ComputationImpl.get_proto(
computation_to_compile)
py_typecheck.check_type(computation_proto, pb.Computation)
comp = building_blocks.ComputationBuildingBlock.from_proto(
computation_proto)
# TODO(b/113123410): Add a compiler options argument that characterizes the
# desired form of the output. To be driven by what the specific backend the
# pipeline is targeting is able to understand. Pending a more fleshed out
# design of the backend API.
# Replace intrinsics with their bodies, for now manually in a fixed order.
# TODO(b/113123410): Replace this with a more automated implementation that
# does not rely on manual maintenance.
comp, _ = value_transformations.replace_intrinsics_with_bodies(
comp, self._context_stack)
comp, _ = tree_transformations.remove_duplicate_building_blocks(comp)
return computation_impl.ComputationImpl(comp.proto,
self._context_stack)
def _replace_intrinsics_with_bodies(comp):
"""Replaces intrinsics with their bodies as defined in `intrinsic_bodies.py`.
Args:
comp: Instance of `building_blocks.ComputationBuildingBlock` in which we
wish to replace all intrinsics with their bodies.
Returns:
An instance of `building_blocks.ComputationBuildingBlock` with
all intrinsics defined in `intrinsic_bodies.py` replaced with their bodies.
"""
py_typecheck.check_type(comp, building_blocks.ComputationBuildingBlock)
context_stack = context_stack_impl.context_stack
comp, _ = value_transformations.replace_intrinsics_with_bodies(
comp, context_stack)
return comp
def test_generic_divide_reduces(self):
uri = intrinsic_defs.GENERIC_DIVIDE.uri
context_stack = context_stack_impl.context_stack
comp = building_blocks.Intrinsic(
uri,
computation_types.FunctionType([tf.float32, tf.float32],
tf.float32))
count_before_reduction = _count_intrinsics(comp, uri)
reduced, modified = value_transformations.replace_intrinsics_with_bodies(
comp, context_stack)
count_after_reduction = _count_intrinsics(reduced, uri)
self.assertGreater(count_before_reduction, 0)
self.assertEqual(count_after_reduction, 0)
tree_analysis.check_intrinsics_whitelisted_for_reduction(reduced)
self.assertTrue(modified)
def test_federated_weighted_mean_reduces(self):
uri = intrinsic_defs.FEDERATED_WEIGHTED_MEAN.uri
context_stack = context_stack_impl.context_stack
@computations.federated_computation(
computation_types.FederatedType(tf.float32, placements.CLIENTS))
def foo(x):
return intrinsics.federated_mean(x, x)
foo_building_block = building_blocks.ComputationBuildingBlock.from_proto(
foo._computation_proto)
count_before_reduction = _count_intrinsics(foo_building_block, uri)
reduced, modified = value_transformations.replace_intrinsics_with_bodies(
foo_building_block, context_stack)
count_after_reduction = _count_intrinsics(reduced, uri)
self.assertGreater(count_before_reduction, 0)
self.assertEqual(count_after_reduction, 0)
self.assertTrue(modified)
def compile(self, computation_to_compile):
"""Compiles `computation_to_compile`.
Args:
computation_to_compile: An instance of `computation_base.Computation` to
compile.
Returns:
An instance of `computation_base.Computation` that repeesents the result.
"""
py_typecheck.check_type(computation_to_compile,
computation_base.Computation)
computation_proto = computation_impl.ComputationImpl.get_proto(
computation_to_compile)
py_typecheck.check_type(computation_proto, pb.Computation)
comp = building_blocks.ComputationBuildingBlock.from_proto(
computation_proto)
# TODO(b/113123410): Add a compiler options argument that characterizes the
# desired form of the output. To be driven by what the specific backend the
# pipeline is targeting is able to understand. Pending a more fleshed out
# design of the backend API.
# Replace intrinsics with their bodies, for now manually in a fixed order.
# TODO(b/113123410): Replace this with a more automated implementation that
# does not rely on manual maintenance.
comp, _ = value_transformations.replace_intrinsics_with_bodies(
comp, self._context_stack)
# Replaces called lambdas with LET constructs with a single local symbol.
comp, _ = transformations.replace_called_lambda_with_block(comp)
# Removes maped or applied identities.
comp, _ = transformations.remove_mapped_or_applied_identity(comp)
# Remove duplicate computations. This is important! otherwise the semantics
# non-deterministic computations (e.g. a `tff.tf_computation` depending on
# `tf.random`) will give unexpected behavior. Additionally, this may reduce
# the amount of calls into TF for some ASTs.
comp, _ = transformations.uniquify_reference_names(comp)
comp, _ = transformations.extract_computations(comp)
comp, _ = transformations.remove_duplicate_computations(comp)
return computation_impl.ComputationImpl(comp.proto, self._context_stack)
def test_raises_on_none(self):
context_stack = context_stack_impl.context_stack
with self.assertRaises(TypeError):
value_transformations.replace_intrinsics_with_bodies(
None, context_stack)
