def _WatcherHandler(handler_uuid, of_blob_ptr):
uuid2handler = session_ctx.GetDefaultSession().uuid2watch_handler
assert handler_uuid in uuid2handler
blob_watched, handler = uuid2handler[handler_uuid]
assert callable(handler)
ndarray_lists = ofblob.OfBlob(of_blob_ptr).CopyToNdarrayLists()
local_blob = local_blob_util.MakeLocalBlob(ndarray_lists, blob_watched)
handler(oft_util.TransformWatchedBlob(local_blob, handler))
def HandlerParallelIdAndLocalBlob(parallel_id, local_blob):
assert parallel_id not in parallel_id2consistent_local_blob
parallel_id2consistent_local_blob[parallel_id] = local_blob
if len(parallel_id2consistent_local_blob) != len_sub_remote_blobs:
return
local_blob_list = [
parallel_id2consistent_local_blob[parallel_id]
for i in range(len_sub_remote_blobs)
]
local_blob = local_blob_util.MergeLocalBlobs(local_blob_list,
blob_watched)
handler(oft_util.TransformWatchedBlob(local_blob, handler))
def HandlerParallelIdAndLocalBlob(parallel_id, local_blob):
assert parallel_id not in parallel_id2consistent_local_blob
parallel_id2consistent_local_blob[parallel_id] = local_blob
if len(parallel_id2consistent_local_blob) != len_sub_remote_blobs:
return
local_blob_list = [
parallel_id2consistent_local_blob[parallel_id]
for i in range(len_sub_remote_blobs)
]
local_numpy = local_blob_list[0].numpy()
if len(local_blob_list) > 1:
print("WARNING: watch return tensor list will concat as axis = 0.")
local_numpy_list = [x.numpy() for x in local_blob_list]
local_numpy = np.concatenate(local_numpy_list, axis=0)
local_blob = local_blob_util.LocalBlob(local_numpy,
blob_watched.is_dynamic)
handler(oft_util.TransformWatchedBlob(local_blob, handler))
def EagerWatch(blob_watched, handler_or_prompt=None):
handler = _CheckOrMakeHandler(blob_watched, handler_or_prompt)
local_blob = local_blob_util.MakeLocalBlob4EagerBlob(blob_watched)
handler(oft_util.TransformWatchedBlob(local_blob, handler))
