code: def _inf_fan_adjust_kaiming(scale, tensor, mode):
      fan_out, fan_in = tensor.infshape[:2]
      if tensor.infshape.ninf() == 0:
          return scale
      elif tensor.infshape.ninf() == 1:
          # should have fixed scale
          if mode == "fan_in" and fan_in.isinf():
              scale *= fan_in.width_mult() ** 0.5
          elif mode == "fan_out" and fan_out.isinf():
              scale *= fan_out.width_mult() ** 0.5
      elif tensor.infshape.ninf() == 2:
          # should scale like fanin
          assert fan_out.isinf() and fan_in.isinf()
          if mode == "fan_out":
              scale *= math.sqrt(
                  fan_out.width_mult() / fan_in.width_mult()
              )
      else:
          raise NotImplementedError(
              "can only handle <=2 inf dimensions currently"
          )
      return scale
  
file_path: NeMo/nemo/collections/nlp/modules/common/megatron/mup/init.py
imports: import math
  import warnings
  import torch
  from torch.nn.init import _calculate_correct_fan, _calculate_fan_in_and_fan_out, _no_grad_fill_, _no_grad_normal_, _no_grad_uniform_, calculate_gain
label_desc: Use image node_id 9 for calling a global function with example usage: _inf_fan_adjust_kaiming(scale, tensor, mode) and returns: scale, scale