code: def __init__(self, config, name=None):
      super().__init__(name=name)
      out_dim = config["n_vocab"]
  
      self.dim = out_dim
      self.norm = hk.LayerNorm(-1, True, True)
  
      self.proj = hk.Linear(self.dim)
  
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 to create a new Projection object from inherited base classes: hk with example: obj = Projection(config, name)