code: def get_init_fn(
      method: str, input_dim: int, init_depth: tp.Optional[int] = None
  ):
      """LM layer initialization.
      Inspired from xlformers: https://github.com/fairinternal/xlformers
  
      Args:
          method (str): Method name for init function. Valid options are:
              'gaussian', 'uniform'.
          input_dim (int): Input dimension of the initialized module.
          init_depth (int, optional): Optional init depth value used to rescale
              the standard deviation if defined.
      """
      # Compute std
      std = 1 / math.sqrt(input_dim)
      # Rescale with depth
      if init_depth is not None:
          std = std / math.sqrt(2 * init_depth)
  
      if method == "gaussian":
          return partial(
              torch.nn.init.trunc_normal_,
              mean=0.0,
              std=std,
              a=-3 * std,
              b=3 * std,
          )
      elif method == "uniform":
          bound = (
              math.sqrt(3) * std
          )  # ensure the standard deviation is `std`
          return partial(torch.nn.init.uniform_, a=-bound, b=bound)
      else:
          raise ValueError("Unsupported layer initialization method")
  
imports: from dataclasses import dataclass
  from functools import partial
  import logging
  import math
  import typing
  import torch
  from torch import nn
  from ..utils import utils
  from ..modules.streaming import StreamingModule, State
  from ..modules.transformer import StreamingTransformer, create_norm_fn
  from ..modules.conditioners import ConditionFuser, ClassifierFreeGuidanceDropout, AttributeDropout, ConditioningProvider, ConditioningAttributes, ConditionType
  from ..modules.codebooks_patterns import CodebooksPatternProvider
  from ..modules.activations import get_activation_fn
label_desc: Use image node_id 1 for calling a global function with example usage: get_init_fn(method, input_dim, init_depth) and returns: partial, partial