code: def _init_nest_weights(
      module: nn.Module, name: str = "", head_bias: float = 0.0
  ):
      """NesT weight initialization
      Can replicate Jax implementation. Otherwise follows vision_transformer.py
      """
      if isinstance(module, nn.Linear):
          if name.startswith("head"):
              trunc_normal_(module.weight, std=0.02, a=-2, b=2)
              nn.init.constant_(module.bias, head_bias)
          else:
              trunc_normal_(module.weight, std=0.02, a=-2, b=2)
              if module.bias is not None:
                  nn.init.zeros_(module.bias)
      elif isinstance(module, nn.Conv2d):
          trunc_normal_(module.weight, std=0.02, a=-2, b=2)
          if module.bias is not None:
              nn.init.zeros_(module.bias)
  
imports: import collections.abc
  import logging
  import math
  from functools import partial
  import torch
  import torch.nn.functional
  from torch import nn
  from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
  from timm.layers import PatchEmbed, Mlp, DropPath, create_classifier, trunc_normal_, _assert
  from timm.layers import create_conv2d, create_pool2d, to_ntuple, use_fused_attn, LayerNorm
  from ._builder import build_model_with_cfg
  from ._features_fx import register_notrace_function
  from ._manipulate import checkpoint_seq, named_apply
  from ._registry import register_model, generate_default_cfgs, register_model_deprecations
label_desc: Use image node_id 3 for calling a global function with example usage: _init_nest_weights(module, name, head_bias) without return types