code: def __init__(self, config, name=None, init_scale=1.0):
      super().__init__(name=name)
      heads = config["n_heads"]
      dim = config["d_model"]
      shards = config["cores_per_replica"]
      norm = getnorm(config["norm"])
      self.is_rotary = config["pe"] == "rotary"
  
      assert dim % heads == 0
      assert heads % shards == 0
  
      self.dim = dim
      self.dim_per_head = dim // heads
      self.heads_per_shard = heads // shards
      self.dim_per_shard = dim // shards
      self.pe_rotary_dims = config.get(
          "pe_rotary_dims", self.dim_per_head
      )
  
      self.norm = norm
  
      self.q = hk.Linear(self.dim_per_shard, with_bias=False)
      self.v = hk.Linear(self.dim_per_shard, with_bias=False)
      self.k = hk.Linear(self.dim_per_shard, with_bias=False)
  
      self.o = hk.Linear(
          self.dim,
          with_bias=False,
          w_init=hk.initializers.TruncatedNormal(
              stddev=init_scale / np.sqrt(self.dim)
          ),
      )
  
      self.dense_proj = hk.Linear(self.dim_per_shard * 4)
      self.dense_proj_o = hk.Linear(
          self.dim,
          w_init=hk.initializers.TruncatedNormal(
              stddev=init_scale / np.sqrt(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 TransformerLayerShard object from inherited base classes: hk with example: obj = TransformerLayerShard(config, name, init_scale)