class_docstr: Embedding Module.
  
  A parameterized function from integers [0, ``num_embeddings``) to
  ``features``-dimensional vectors. This ``Module`` will create an ``embedding``
  matrix with shape ``(num_embeddings, features)``. When calling this layer,
  the input values will be used to 0-index into the ``embedding`` matrix.
  Indexing on a value greater than or equal to ``num_embeddings`` will result
  in ``nan`` values. When ``num_embeddings`` equals to 1, it will
  broadcast the ``embedding`` matrix to input shape with ``features``
  dimension appended.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax, jax.numpy as jnp
  
    >>> layer = nn.Embed(num_embeddings=5, features=3)
    >>> indices_input = jnp.array([[0, 1, 2], [-1, -2, -3]])
    >>> variables = layer.init(jax.random.key(0), indices_input)
    >>> variables
    {'params': {'embedding': Array([[-0.28884724,  0.19018005, -0.414205  ],
           [-0.11768015, -0.54618824, -0.3789283 ],
           [ 0.30428642,  0.49511626,  0.01706631],
           [-0.0982546 , -0.43055868,  0.20654906],
           [-0.688412  , -0.46882293,  0.26723292]], dtype=float32)}}
    >>> # get the first three and last three embeddings
    >>> layer.apply(variables, indices_input)
    Array([[[-0.28884724,  0.19018005, -0.414205  ],
            [-0.11768015, -0.54618824, -0.3789283 ],
            [ 0.30428642,  0.49511626,  0.01706631]],
    <BLANKLINE>
           [[-0.688412  , -0.46882293,  0.26723292],
            [-0.0982546 , -0.43055868,  0.20654906],
            [ 0.30428642,  0.49511626,  0.01706631]]], dtype=float32)
  
  Attributes:
    num_embeddings: number of embeddings / vocab size.
    features: number of feature dimensions for each embedding.
    dtype: the dtype of the embedding vectors (default: same as embedding).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    embedding_init: embedding initializer.
code: def setup(self):
      self.embedding = self.param(
          "embedding",
          self.embedding_init,
          (self.num_embeddings, self.features),
          self.param_dtype,
      )
  
imports: import dataclasses
  from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Union
  import jax
  import jax.numpy
  import numpy
  from jax import eval_shape, lax
  from jax.core import ShapedArray
  from flax.core import meta
  from flax.linen import initializers
  from flax.linen.dtypes import promote_dtype
  from flax.linen.module import Module, compact
label_desc: Use image node_id 1 for calling the Embed obj's underlying member method code with example usage: obj.setup() without return types