class_docstr: This is a temporary extension of `LazyModuleMixin` to support passing keyword
  arguments to lazy module's `forward` method.
  
  The long-term plan is to upstream this feature to `LazyModuleMixin`. Please see
  https://github.com/pytorch/pytorch/issues/59923 for details.
  
  Please see `TestLazyModuleExtensionMixin`, which contains unit tests that ensure:
    * `LazyModuleExtensionMixin._infer_parameters` has source code parity with
      torch.nn.modules.lazy.LazyModuleMixin._infer_parameters, except that the former
      can accept keyword arguments.
    * `LazyModuleExtensionMixin._call_impl` has source code parity with
      `torch.nn.Module._call_impl`, except that the former can pass keyword arguments
      to forward pre hooks."
class_docstr_tok: ["This","is","a","temporary","extension","of","`","LazyModuleMixin","`","to","support","passing","keyword","arguments","to","lazy","module","'s","`","forward","`","method",".","The","long-term","plan","is","to","upstream","this","feature","to","`","LazyModuleMixin","`",".","Please","see","https",":","\/\/github.com\/pytorch\/pytorch\/issues\/59923","for","details",".","Please","see","`","TestLazyModuleExtensionMixin","`",",","which","contains","unit","tests","that","ensure",":","*","`","LazyModuleExtensionMixin._infer_parameters","`","has","source","code","parity","with","torch.nn.modules.lazy.LazyModuleMixin._infer_parameters",",","except","that","the","former","can","accept","keyword","arguments",".","*","`","LazyModuleExtensionMixin._call_impl","`","has","source","code","parity","with","`","torch.nn.Module._call_impl","`",",","except","that","the","former","can","pass","keyword","arguments","to","forward","pre","hooks",".","''"]
code: def _call_impl(self, *input, **kwargs):  # noqa: C901
      # pyre-ignore[16]
      forward_call = (
          self._slow_forward
          if torch._C._get_tracing_state()
          else self.forward
      )
      # If we don't have any hooks, we want to skip the rest of the logic in
      # this function, and just call forward.
      # pyre-ignore[16]
      if not (
          self._backward_hooks
          or self._forward_hooks
          or self._forward_pre_hooks
          or _global_backward_hooks
          or _global_forward_hooks
          or _global_forward_pre_hooks
      ):
          return forward_call(*input, **kwargs)
      # Do not call functions when jit is used
      full_backward_hooks, non_full_backward_hooks = [], []
      if self._backward_hooks or _global_backward_hooks:
          # pyre-ignore[16]
          (
              full_backward_hooks,
              non_full_backward_hooks,
          ) = self._get_backward_hooks()
      if _global_forward_pre_hooks or self._forward_pre_hooks:
          # pyre-ignore[60]: Concatenation not yet support for multiple variadic
          #  tuples: `*torch.nn.modules.module._global_forward_pre_hooks.values(),
          #  *self._forward_pre_hooks.values()`.
          for hook in (
              *_global_forward_pre_hooks.values(),
              *self._forward_pre_hooks.values(),
          ):
              if len(inspect.signature(hook).parameters) == 3:
                  result = hook(self, input, kwargs)
              else:
                  result = hook(self, input)
              if result is not None:
                  if not isinstance(result, tuple):
                      result = (result,)
                  input = result
  
      bw_hook = None
      if full_backward_hooks:
          # pyre-fixme[20]: Argument `user_pre_hooks` expected.
          bw_hook = hooks.BackwardHook(self, full_backward_hooks)
          input = bw_hook.setup_input_hook(input)
  
      result = forward_call(*input, **kwargs)
      if _global_forward_hooks or self._forward_hooks:
          # pyre-ignore[60]: Concatenation not yet support for multiple variadic
          #  tuples: `*torch.nn.modules.module._global_forward_hooks.values(),
          #  *self._forward_hooks.values()`.
          for hook in (
              *_global_forward_hooks.values(),
              *self._forward_hooks.values(),
          ):
              hook_result = hook(self, input, result)
              if hook_result is not None:
                  result = hook_result
  
      if bw_hook:
          result = bw_hook.setup_output_hook(result)
  
      # Handle the non-full backward hooks
      if non_full_backward_hooks:
          var = result
          while not isinstance(var, torch.Tensor):
              if isinstance(var, dict):
                  var = next(
                      (
                          v
                          for v in var.values()
                          if isinstance(v, torch.Tensor)
                      )
                  )
              else:
                  var = var[0]
          grad_fn = var.grad_fn
          if grad_fn is not None:
              for hook in non_full_backward_hooks:
                  wrapper = functools.partial(hook, self)
                  functools.update_wrapper(wrapper, hook)
                  grad_fn.register_hook(wrapper)
              # pyre-ignore[16]
              self._maybe_warn_non_full_backward_hook(
                  input, result, grad_fn
              )
  
      return result
  
file_path: torchrec/torchrec/modules/lazy_extension.py
imports: import functools
  import inspect
  from typing import Any, Callable
  import torch
  import torch.utils.hooks
  from torch.nn.modules.lazy import _LazyProtocol, LazyModuleMixin
  from torch.nn.modules.module import _global_backward_hooks, _global_forward_hooks, _global_forward_pre_hooks
label_desc: Use image node_id 3 for calling the LazyModuleExtensionMixin obj's underlying member method code with example usage: obj._call_impl() and returns: result, forward_call