code: def getnorm(type):
      if type == "layernorm":
          return ReplicatedLayerNorm()
      if type == "layernorm-desync":
          return hk.LayerNorm(-1, True, True)
      elif type == "layernorm-nobias":
          return ReplicatedLayerNorm(offset=False)
      elif type == "rmsnorm":
          return RMSNorm(False, True)
      elif type == "scalenorm":
          return RMSNorm(False, False)
      elif type == "rmsnorm-bias":
          return RMSNorm(True, True)
      elif type == "scalenorm-bias":
          return RMSNorm(True, False)
      else:
          raise Exception("Not implemented")
  
imports: import haiku
  import jax
  import jax.numpy
  import numpy
  from einops import rearrange, repeat
  from mesh_transformer.util import f_psum, g_psum, maybe_shard, head_print
  from jax.experimental import PartitionSpec
  from jax.experimental.maps import thread_resources
label_desc: Use image node_id 1 for calling a global function with example usage: getnorm(type) and returns: ReplicatedLayerNorm, hk, ReplicatedLayerNorm, RMSNorm, RMSNorm, RMSNorm, RMSNorm