Function wrapped_func

imperative/python/megengine/jit/partial_tracing.py:181–293

(*args, **kwargs)

            return hash(tuple([map_scalar_to_tuple(t._tuple_shape) for t in tensors]))

        def wrapped_func(*args, **kwargs):
            from ..traced_module.pytree import tree_flatten
            from ..core.autodiff.grad import Grad
            from ..autodiff.grad_manager import get_backwarding_grad_manager

            nonlocal traced
            nonlocal custom_autodiff
            nonlocal outdef
            nonlocal shape_hash
            nonlocal unexpected_shape_hash
            nonlocal inp_grad_maps, out_grad_maps
            trace_obj.convert_optimizer_state_to_tensor(*args, **kwargs)
            if not traced:
                traced = True
                fargs = trace_obj.flatten_inputs(*args, **kwargs)
                shape_hash = get_shape_hash(*fargs) if check_shape else None
                # we want to construct map like {inp_hid: inp_grad_hid, ...}. but the
                # backward() is not called now and the grad is unknown. so we record the
                # grad_slot here and then use grad_slot to get the grad hid when exiting
                # trace
                # **The grad grad_slot returned is the orginal grad, that means it is
                # not been processed by gradmanager attach callback, eg., the grad
                # before allreduce when distributed training**
                inp_grad_maps = {get_handle_id(t): get_grad_slot(t) for t in fargs}
                del fargs

                def exit_trace():
                    backward_trace_obj._trace.exit()
                    backward_trace_obj.unset_env()
                    for k, v in inp_grad_maps.items():
                        inp_grad_maps[k] = (
                            get_handle_id(v.grad) if v is not None else -1
                        )

                    # if we want to trace the gradmanager attach callback, we will replace
                    # the original grad by the grad processed the gradmanger attach callback
                    if trace_gmcallback:

                        def cb_wrapper(cb):
                            def wrapper(param, grad):
                                return (
                                    grad
                                    if grad._is_external_value()
                                    else cb(param, grad)
                                )

                            return wrapper

                        current_gm = get_backwarding_grad_manager()
                        for _attached_tensor_id, grad in current_gm._gradients.items():
                            spec = current_gm._attach_specs.get(_attached_tensor_id)
                            _attached_tensor = spec and spec.tensor()
                            if _attached_tensor is None:
                                continue
                            inp_hid = get_handle_id(_attached_tensor)
                            if inp_hid in inp_grad_maps:
                                spec.callbacks = list(map(cb_wrapper, spec.callbacks))
                                inp_grad_maps[inp_hid] = get_handle_id(grad)