def create_lambda_identity(
type_spec: computation_types.Type) -> pb.Computation:
"""Returns a lambda computation representing an identity function.
Has the type signature:
(T -> T)
Args:
type_spec: A `computation_types.Type`.
Returns:
An instance of `pb.Computation`.
"""
type_signature = type_factory.unary_op(type_spec)
result = pb.Computation(type=type_serialization.serialize_type(type_spec),
reference=pb.Reference(name='a'))
fn = pb.Lambda(parameter_name='a', result=result)
# We are unpacking the lambda argument here because `lambda` is a reserved
# keyword in Python, but it is also the name of the parameter for a
# `pb.Computation`.
# https://developers.google.com/protocol-buffers/docs/reference/python-generated#keyword-conflicts
return pb.Computation(
type=type_serialization.serialize_type(type_signature),
**{'lambda': fn}) # pytype: disable=wrong-keyword-args
def proto(self):
return pb.Computation(
type=type_serialization.serialize_type(self.type_signature),
**{
'lambda':
pb.Lambda(parameter_name=self._parameter_name,
result=self._result.proto)
})
def _create_lambda_identity_comp(type_spec):
"""Returns a `pb.Computation` representing an identity function."""
py_typecheck.check_type(type_spec, computation_types.Type)
type_signature = type_serialization.serialize_type(
type_factory.unary_op(type_spec))
result = pb.Computation(type=type_serialization.serialize_type(type_spec),
reference=pb.Reference(name='x'))
fn = pb.Lambda(parameter_name='x', result=result)
# We are unpacking the lambda argument here because `lambda` is a reserved
# keyword in Python, but it is also the name of the parameter for a
# `pb.Computation`.
# https://developers.google.com/protocol-buffers/docs/reference/python-generated#keyword-conflicts
return pb.Computation(type=type_signature, **{'lambda': fn}) # pytype: disable=wrong-keyword-args
def _create_lambda_identity_comp(type_spec):
py_typecheck.check_type(type_spec, computation_types.Type)
return pb.Computation(
**{
'type':
type_serialization.serialize_type(type_factory.unary_op(
type_spec)),
'lambda':
pb.Lambda(parameter_name='x',
result=pb.Computation(
type=type_serialization.serialize_type(type_spec),
reference=pb.Reference(name='x')))
})
def create_dummy_computation_lambda_identity():
"""Returns a lambda computation and type `(float32 -> float32)`."""
type_signature = type_factory.unary_op(tf.float32)
result = pb.Computation(
type=type_serialization.serialize_type(tf.float32),
reference=pb.Reference(name='a'))
fn = pb.Lambda(parameter_name='a', result=result)
# We are unpacking the lambda argument here because `lambda` is a reserved
# keyword in Python, but it is also the name of the parameter for a
# `pb.Computation`.
# https://developers.google.com/protocol-buffers/docs/reference/python-generated#keyword-conflicts
value = pb.Computation(
type=type_serialization.serialize_type(type_signature), **{'lambda': fn}) # pytype: disable=wrong-keyword-args
return value, type_signature
def create_dummy_computation_lambda_empty():
"""Returns a lambda computation and type `( -> <>)`."""
result_type = computation_types.NamedTupleType([])
type_signature = computation_types.FunctionType(None, result_type)
result = pb.Computation(
type=type_serialization.serialize_type(result_type),
tuple=pb.Tuple(element=[]))
fn = pb.Lambda(parameter_name=None, result=result)
# We are unpacking the lambda argument here because `lambda` is a reserved
# keyword in Python, but it is also the name of the parameter for a
# `pb.Computation`.
# https://developers.google.com/protocol-buffers/docs/reference/python-generated#keyword-conflicts
value = pb.Computation(
type=type_serialization.serialize_type(type_signature), **{'lambda': fn}) # pytype: disable=wrong-keyword-args
return value, type_signature
def create_lambda_empty_struct() -> pb.Computation:
"""Returns a lambda computation returning an empty struct.
Has the type signature:
( -> <>)
Returns:
An instance of `pb.Computation`.
"""
result_type = computation_types.StructType([])
type_signature = computation_types.FunctionType(None, result_type)
result = pb.Computation(
type=type_serialization.serialize_type(result_type),
struct=pb.Struct(element=[]))
fn = pb.Lambda(parameter_name=None, result=result)
# We are unpacking the lambda argument here because `lambda` is a reserved
# keyword in Python, but it is also the name of the parameter for a
# `pb.Computation`.
# https://developers.google.com/protocol-buffers/docs/reference/python-generated#keyword-conflicts
return pb.Computation(
type=type_serialization.serialize_type(type_signature), **{'lambda': fn}) # pytype: disable=wrong-keyword-args
def create_dummy_identity_lambda_computation(type_spec=tf.int32):
"""Returns a `pb.Computation` representing an identity lambda.
The type signature of this `pb.Computation` is:
(int32 -> int32)
Args:
type_spec: A type signature.
Returns:
A `pb.Computation`.
"""
type_signature = type_serialization.serialize_type(
type_factory.unary_op(type_spec))
result = pb.Computation(type=type_serialization.serialize_type(type_spec),
reference=pb.Reference(name='a'))
fn = pb.Lambda(parameter_name='a', result=result)
# We are unpacking the lambda argument here because `lambda` is a reserved
# keyword in Python, but it is also the name of the parameter for a
# `pb.Computation`.
# https://developers.google.com/protocol-buffers/docs/reference/python-generated#keyword-conflicts
return pb.Computation(type=type_signature, **{'lambda': fn}) # pytype: disable=wrong-keyword-args
