def eager_user_op_builder(op_name): job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName() op_name = name_scope.GetJobNameScopePrefix(job_name) + op_name return UserOpConfBuilder(EagerUserOp, op_name, None)
def name_scope_stack_pop(): job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName() sess = session_context.GetDefaultSession() assert job_name in sess.job_name2name_scope_stack assert len(sess.job_name2name_scope_stack[job_name]) > 0 return sess.job_name2name_scope_stack[job_name].pop()
def eager_consistent_user_op_module_builder(op_type_name): job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName() op_name = name_scope.GetJobNameScopePrefix(job_name) + op_type_name return UserOpModuleBuilder(EagerConsistentUserOpModule, op_name, op_type_name)
def name_scope_stack_push(name): job_name = oneflow_api.JobBuildAndInferCtx_GetCurrentJobName() sess = session_context.GetDefaultSession() if job_name not in sess.job_name2name_scope_stack: sess.job_name2name_scope_stack[job_name] = [] sess.job_name2name_scope_stack[job_name].append(name)