def _AddAndInferMirroredOp(mirrored_lbn, op_conf, sub_consistent_blob_list):
compile_context.CurJobAddMirroredOp(op_conf)
job_name = c_api_util.JobBuildAndInferCtx_GetCurrentJobName()
num_sub_lbi = c_api_util.JobBuildAndInferCtx_MirroredBlobGetNumSubLbi(
job_name, mirrored_lbn)
for i in range(num_sub_lbi):
sub_lbi = c_api_util.JobBuildAndInferCtx_MirroredBlobGetSubLbi(
job_name, mirrored_lbn, i)
sub_consistent_blob_list.append(
remote_blob_util.ConsistentBlob(sub_lbi))
def __init__(self, lbi, **kw):
MirroredBlob.__init__(self, lbi, **kw)
self.job_name_ = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName()
self.sub_consistent_blob_list_ = []
lbn = self.logical_blob_name
num_sub_lbi = c_api_util.JobBuildAndInferCtx_MirroredBlobGetNumSubLbi(
self.job_name_, lbn)
for i in range(num_sub_lbi):
sub_lbi = c_api_util.JobBuildAndInferCtx_MirroredBlobGetSubLbi(
self.job_name_, lbn, i)
consistent_blob = LazyConsistentBlob(sub_lbi)
self.sub_consistent_blob_list_.append(consistent_blob)
def _AddAndInferMirroredOp(mirrored_lbn, op_conf, sub_consistent_blob_list):
compile_context.CurJobAddMirroredOp(op_conf)
job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName()
num_sub_lbi = c_api_util.JobBuildAndInferCtx_MirroredBlobGetNumSubLbi(
job_name, mirrored_lbn)
for i in range(num_sub_lbi):
sub_lbi = c_api_util.JobBuildAndInferCtx_MirroredBlobGetSubLbi(
job_name, mirrored_lbn, i)
lbi = lbi_util.LogicalBlobId()
lbi.set_op_name(sub_lbi.op_name)
lbi.set_blob_name(sub_lbi.blob_name)
sub_consistent_blob_list.append(
oneflow_api.ConsistentBlob(lbi, "", oneflow_api.distribute.auto()))