def tf_buffer(data=None): """Context manager that creates and deletes TF_Buffer. Example usage: with tf_buffer() as buf: # get serialized graph def into buf ... proto_data = c_api.TF_GetBuffer(buf) graph_def.ParseFromString(compat.as_bytes(proto_data)) # buf has been deleted with tf_buffer(some_string) as buf: c_api.TF_SomeFunction(buf) # buf has been deleted Args: data: An optional `bytes`, `str`, or `unicode` object. If not None, the yielded buffer will contain this data. Yields: Created TF_Buffer """ if data: buf = c_api.TF_NewBufferFromString(compat.as_bytes(data)) else: buf = c_api.TF_NewBuffer() try: yield buf finally: c_api.TF_DeleteBuffer(buf)
def get_api_def(self, op_name): api_def_proto = api_def_pb2.ApiDef() buf = c_api.TF_ApiDefMapGet(self._api_def_map, op_name, len(op_name)) try: api_def_proto.ParseFromString(c_api.TF_GetBuffer(buf)) finally: c_api.TF_DeleteBuffer(buf) return api_def_proto
def __init__(self): op_def_proto = op_def_pb2.OpList() buf = c_api.TF_GetAllOpList() try: op_def_proto.ParseFromString(c_api.TF_GetBuffer(buf)) self._api_def_map = c_api.TF_NewApiDefMap(buf) finally: c_api.TF_DeleteBuffer(buf) self._op_per_name = {} for op in op_def_proto.op: self._op_per_name[op.name] = op
def __del__(self): c_api.TF_DeleteBuffer(self.buffer)