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",".","''"]
docstr: Applies `fn` recursively to every submodule (as returned by `.children()`)
  as well as self. Typical use includes initializing the parameters of a model.
  
  NOTE:
      Calling `apply()` on an uninitialized lazy-module will result in an error.
      User is required to initialize a lazy-module (by doing a dummy forward pass)
      before calling `apply()` on the lazy-module.
  
  Args:
      fn (torch.nn.Module -> None): function to be applied to each submodule.
  
  Returns:
      torch.nn.Module: self
  
  Example::
  
      @torch.no_grad()
      def init_weights(m):
          print(m)
          if type(m) == torch.nn.LazyLinear:
              m.weight.fill_(1.0)
              print(m.weight)
  
      linear = torch.nn.LazyLinear(2)
      linear.apply(init_weights)  # this fails, because `linear` (a lazy-module) hasn't been initialized yet
  
      input = torch.randn(2, 10)
      linear(input)  # run a dummy forward pass to initialize the lazy-module
  
      linear.apply(init_weights)  # this works now
docstr_tok: ["Applies","`","fn","`","recursively","to","every","submodule","(","as","returned","by","`",".children","(",")","`",")","as","well","as","self",".","Typical","use","includes","initializing","the","parameters","of","a","model",".","NOTE",":","Calling","`","apply","(",")","`","on","an","uninitialized","lazy-module","will","result","in","an","error",".","User","is","required","to","initialize","a","lazy-module","(","by","doing","a","dummy","forward","pass",")","before","calling","`","apply","(",")","`","on","the","lazy-module",".","Args",":","fn","(","torch.nn.Module","-",">","None",")",":","function","to","be","applied","to","each","submodule",".","Returns",":","torch.nn.Module",":","self","Example",":",":","@","torch.no_grad","(",")","def","init_weights","(","m",")",":","print","(","m",")","if","type","(","m",")","==","torch.nn.LazyLinear",":","m.weight.fill_","(","1.0",")","print","(","m.weight",")","linear","=","torch.nn.LazyLinear","(","2",")","linear.apply","(","init_weights",")","#","this","fails",",","because","`","linear","`","(","a","lazy-module",")","has","n't","been","initialized","yet","input","=","torch.randn","(","2",",","10",")","linear","(","input",")","#","run","a","dummy","forward","pass","to","initialize","the","lazy-module","linear.apply","(","init_weights",")","#","this","works","now"]
code: def apply(
      self, fn: Callable[[torch.nn.Module], None]
  ) -> torch.nn.Module:
      """Applies `fn` recursively to every submodule (as returned by `.children()`)
      as well as self. Typical use includes initializing the parameters of a model.
  
      NOTE:
          Calling `apply()` on an uninitialized lazy-module will result in an error.
          User is required to initialize a lazy-module (by doing a dummy forward pass)
          before calling `apply()` on the lazy-module.
  
      Args:
          fn (torch.nn.Module -> None): function to be applied to each submodule.
  
      Returns:
          torch.nn.Module: self
  
      Example::
  
          @torch.no_grad()
          def init_weights(m):
              print(m)
              if type(m) == torch.nn.LazyLinear:
                  m.weight.fill_(1.0)
                  print(m.weight)
  
          linear = torch.nn.LazyLinear(2)
          linear.apply(init_weights)  # this fails, because `linear` (a lazy-module) hasn't been initialized yet
  
          input = torch.randn(2, 10)
          linear(input)  # run a dummy forward pass to initialize the lazy-module
  
          linear.apply(init_weights)  # this works now
      """
  
      if hasattr(self, "_initialize_hook"):
          raise RuntimeError(
              "Module {} has not been initialized. ".format(self)
              + "Please run a dummy forward pass on the model to initialize all modules, "
              + "or use torchrec.modules.lazy_extension.lazy_apply to attach a function "
              + "to this module which would be applied after this module is initialized."
          )
      # If the module is already initialized, call `super().apply(fn)` to
      # run the usual apply logic.
      # pyre-ignore[16]
      return super().apply(fn)
  
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 1 for calling the LazyModuleExtensionMixin obj's underlying member method code with example usage: obj.apply(fn) and returns: super