def do_encode(self, tensor_spec_value, encode_fn):
encoded_tensor_spec = struct_pb2.StructuredValue()
encoded_tensor_spec.tensor_spec_value.CopyFrom(
struct_pb2.TensorSpecProto(
shape=encode_fn(tensor_spec_value.shape).tensor_shape_value,
dtype=encode_fn(tensor_spec_value.dtype).tensor_dtype_value,
name=tensor_spec_value.name))
return encoded_tensor_spec
def experimental_as_proto(self) -> struct_pb2.TensorSpecProto:
"""Returns a proto representation of the TensorSpec instance."""
return struct_pb2.TensorSpecProto(
shape=self.shape.experimental_as_proto(),
dtype=self.dtype.experimental_as_proto().datatype,
name=self.name)
