code: def embedding(indices, vocab_dim, output_dim, variable_dtype, name="embedding"):
    
    weights = embedding_weights(
        indices.mesh, vocab_dim, output_dim, variable_dtype, name)
    return mtf.gather(weights, indices, vocab_dim)
docs: Embedding layer.