code: def __call__(self, x, dtype=jnp.bfloat16):
      shard_start_index = (
          jax.lax.axis_index("shard") * self.in_dim_per_shard
      )
  
      input_onehot = jax.nn.one_hot(
          x - shard_start_index, self.in_dim_per_shard
      )
      proj_out = self.proj(input_onehot)
  
      proj_out = g_psum(proj_out)
  
      if self.positional_embeddings is not None:
          all_pos_embed = jax.lax.all_gather(
              self.positional_embeddings, "shard"
          )
  
          all_pos_embed = hk.Flatten()(
              jnp.transpose(all_pos_embed, (1, 0, 2))
          )
  
          proj_out += all_pos_embed
  
      return proj_out
  
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 2 for calling the EmbeddingShard obj's underlying member method code with example usage: obj.__call__(x, dtype) and returns: proj_out