code: def __init__(self, config, name=None):
      super().__init__(name=name)
      in_dim = config["n_vocab"]
      out_dim = config["d_model"]
      shards = config["cores_per_replica"]
  
      assert in_dim % shards == 0
  
      self.in_dim = in_dim
      self.out_dim = out_dim
  
      self.proj = hk.Linear(
          self.out_dim,
          w_init=hk.initializers.TruncatedNormal(
              stddev=1 / np.sqrt(in_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 EmbeddingShardV2 object from inherited base classes: hk with example: obj = EmbeddingShardV2(config, name)