def embed_tensorflow_computation(comp, type_spec=None, device=None):
    'Embeds a TensorFlow computation for use in the eager context.\n\n  Args:\n    comp: An instance of `pb.Computation`.\n    type_spec: An optional `tff.Type` instance or something convertible to it.\n    device: An optional `tf.config.LogicalDevice`.\n\n  Returns:\n    Either a one-argument or a zero-argument callable that executes the\n    computation in eager mode.\n\n  Raises:\n    TypeError: If arguments are of the wrong types, e.g., in `comp` is not a\n      TensorFlow computation.\n  '
    py_typecheck.check_type(comp, pb.Computation)
    comp_type = type_serialization.deserialize_type(comp.type)
    type_spec = computation_types.to_type(type_spec)
    if (type_spec is not None):
        if (not type_spec.is_equivalent_to(comp_type)):
            raise TypeError('Expected a computation of type {}, got {}.'.format(type_spec, comp_type))
    else:
        type_spec = comp_type
    must_pin_function_to_cpu = type_analysis.contains(type_spec.result, (lambda t: t.is_sequence()))
    which_computation = comp.WhichOneof('computation')
    if (which_computation != 'tensorflow'):
        unexpected_building_block = building_blocks.ComputationBuildingBlock.from_proto(comp)
        raise TypeError('Expected a TensorFlow computation, found {}.'.format(unexpected_building_block))
    if type_spec.is_function():
        param_type = type_spec.parameter
        result_type = type_spec.result
    else:
        param_type = None
        result_type = type_spec
    wrapped_fn = _get_wrapped_function_from_comp(comp, must_pin_function_to_cpu, param_type, device)
    param_fns = []
    if (param_type is not None):
        for spec in structure.flatten(type_spec.parameter):
            if spec.is_tensor():
                param_fns.append((lambda x: x))
            else:
                py_typecheck.check_type(spec, computation_types.SequenceType)
                param_fns.append(tf.data.experimental.to_variant)
    result_fns = []
    for spec in structure.flatten(result_type):
        if spec.is_tensor():
            result_fns.append((lambda x: x))
        else:
            py_typecheck.check_type(spec, computation_types.SequenceType)
            tf_structure = type_conversions.type_to_tf_structure(spec.element)

            def fn(x, tf_structure=tf_structure):
                return tf.data.experimental.from_variant(x, tf_structure)
            result_fns.append(fn)

    def _fn_to_return(arg, param_fns, wrapped_fn):
        param_elements = []
        if (arg is not None):
            arg_parts = structure.flatten(arg)
            if (len(arg_parts) != len(param_fns)):
                raise RuntimeError('Expected {} arguments, found {}.'.format(len(param_fns), len(arg_parts)))
            for (arg_part, param_fn) in zip(arg_parts, param_fns):
                param_elements.append(param_fn(arg_part))
        result_parts = wrapped_fn(*param_elements)
        resources = []
        for op in wrapped_fn.graph.get_operations():
            if (op.type == 'VarHandleOp'):
                resources += op.outputs
        if resources:
            for resource in wrapped_fn.prune(feeds={}, fetches=resources)():
                tf.raw_ops.DestroyResourceOp(resource=resource)
        result_elements = []
        for (result_part, result_fn) in zip(result_parts, result_fns):
            result_elements.append(result_fn(result_part))
        return structure.pack_sequence_as(result_type, result_elements)
    fn_to_return = (lambda arg, p=param_fns, w=wrapped_fn: _fn_to_return(arg, p, w))
    if must_pin_function_to_cpu:
        old_fn_to_return = fn_to_return

        def fn_to_return(x):
            with tf.device('cpu'):
                return old_fn_to_return(x)
    elif (device is not None):
        old_fn_to_return = fn_to_return

        def fn_to_return(x):
            with tf.device(device.name):
                return old_fn_to_return(x)
    if (param_type is not None):
        return (lambda arg: fn_to_return(arg))
    else:
        return (lambda : fn_to_return(None))
Embeds a TensorFlow computation for use in the eager context.

Args:
  comp: An instance of `pb.Computation`.
  type_spec: An optional `tff.Type` instance or something convertible to it.
  device: An optional `tf.config.LogicalDevice`.

Returns:
  Either a one-argument or a zero-argument callable that executes the
  computation in eager mode.

Raises:
  TypeError: If arguments are of the wrong types, e.g., in `comp` is not a
    TensorFlow computation.
def embed_tensorflow_computation(comp, type_spec=None, device=None):
    'Embeds a TensorFlow computation for use in the eager context.\n\n  Args:\n    comp: An instance of `pb.Computation`.\n    type_spec: An optional `tff.Type` instance or something convertible to it.\n    device: An optional `tf.config.LogicalDevice`.\n\n  Returns:\n    Either a one-argument or a zero-argument callable that executes the\n    computation in eager mode.\n\n  Raises:\n    TypeError: If arguments are of the wrong types, e.g., in `comp` is not a\n      TensorFlow computation.\n  '
    py_typecheck.check_type(comp, pb.Computation)
    comp_type = type_serialization.deserialize_type(comp.type)
    type_spec = computation_types.to_type(type_spec)
    if (type_spec is not None):
        if (not type_spec.is_equivalent_to(comp_type)):
            raise TypeError('Expected a computation of type {}, got {}.'.format(type_spec, comp_type))
    else:
        type_spec = comp_type
    must_pin_function_to_cpu = type_analysis.contains(type_spec.result, (lambda t: t.is_sequence()))
    which_computation = comp.WhichOneof('computation')
    if (which_computation != 'tensorflow'):
        unexpected_building_block = building_blocks.ComputationBuildingBlock.from_proto(comp)
        raise TypeError('Expected a TensorFlow computation, found {}.'.format(unexpected_building_block))
    if type_spec.is_function():
        param_type = type_spec.parameter
        result_type = type_spec.result
    else:
        param_type = None
        result_type = type_spec
    wrapped_fn = _get_wrapped_function_from_comp(comp, must_pin_function_to_cpu, param_type, device)
    param_fns = []
    if (param_type is not None):
        for spec in structure.flatten(type_spec.parameter):
            if spec.is_tensor():
                param_fns.append((lambda x: x))
            else:
                py_typecheck.check_type(spec, computation_types.SequenceType)
                param_fns.append(tf.data.experimental.to_variant)
    result_fns = []
    for spec in structure.flatten(result_type):
        if spec.is_tensor():
            result_fns.append((lambda x: x))
        else:
            py_typecheck.check_type(spec, computation_types.SequenceType)
            tf_structure = type_conversions.type_to_tf_structure(spec.element)

            def fn(x, tf_structure=tf_structure):
                return tf.data.experimental.from_variant(x, tf_structure)
            result_fns.append(fn)

    def _fn_to_return(arg, param_fns, wrapped_fn):
        param_elements = []
        if (arg is not None):
            arg_parts = structure.flatten(arg)
            if (len(arg_parts) != len(param_fns)):
                raise RuntimeError('Expected {} arguments, found {}.'.format(len(param_fns), len(arg_parts)))
            for (arg_part, param_fn) in zip(arg_parts, param_fns):
                param_elements.append(param_fn(arg_part))
        result_parts = wrapped_fn(*param_elements)
        resources = []
        for op in wrapped_fn.graph.get_operations():
            if (op.type == 'VarHandleOp'):
                resources += op.outputs
        if resources:
            for resource in wrapped_fn.prune(feeds={}, fetches=resources)():
                tf.raw_ops.DestroyResourceOp(resource=resource)
        result_elements = []
        for (result_part, result_fn) in zip(result_parts, result_fns):
            result_elements.append(result_fn(result_part))
        return structure.pack_sequence_as(result_type, result_elements)
    fn_to_return = (lambda arg, p=param_fns, w=wrapped_fn: _fn_to_return(arg, p, w))
    if must_pin_function_to_cpu:
        old_fn_to_return = fn_to_return

        def fn_to_return(x):
            with tf.device('cpu'):
                return old_fn_to_return(x)
    elif (device is not None):
        old_fn_to_return = fn_to_return

        def fn_to_return(x):
            with tf.device(device.name):
                return old_fn_to_return(x)
    if (param_type is not None):
        return (lambda arg: fn_to_return(arg))
    else:
        return (lambda : fn_to_return(None))