def from_config(
      cls,
      vocab_size: int,
      embedding_axes: dict[str, int],
      vocabulary_axis: str = "vocabulary",
      initializer: linear_and_affine.LinearOperatorWeightInitializer = (
          functools.partial(
              linear_and_affine.variance_scaling_initializer,
              scale=1.0,
              mode="fan_out",
              distribution="normal",
          )
      ),
      dtype: np.typing.DTypeLike = np.float32,
  ) -> EmbeddingTable:
    """Constructs an ``EmbeddingTable`` with uninitialized parameters.

    Args:
      vocab_size: The size of the vocabulary.
      embedding_axes: A dictionary mapping embedding axis names to their sizes.
        Will usually be a single-element dictionary of the form ``{"embedding":
        embedding_size}``.
      vocabulary_axis: The name of the axis that corresponds to the vocabulary.
        This axis will be indexed into when performing embedding lookups. Must
        not appear in ``embedding_axes``.
      initializer: A weight initializer that will be used to initialize the
        parameters. For the purposes of initialization, the "input axes" are of
        dimension 1, and the "output axes" are the ``embedding_axes``; the
        vocabulary axis is treated as a parallel axis. Defaults to fan-out
        normalization over the embedding axes.
      dtype: The data type of the embedding parameters.

    Returns:
      An EmbeddingTable with uninitialized parameters of the given shape.
    """
    if vocabulary_axis in embedding_axes:
      raise ValueError(
          f"`vocabulary_axis` {vocabulary_axis} should not appear in"
          f"`embedding_axes` {embedding_axes}"
      )

    return cls(  # pytype: disable=wrong-arg-types
        embeddings=parameters.UninitializedParameter(
            initializer=functools.partial(
                initializer,
                input_axes={},
                output_axes=embedding_axes,
                parallel_axes={vocabulary_axis: vocab_size},
                convolution_spatial_axes={},
                dtype=dtype,
            ),
            name="embeddings",
        ),
        vocabulary_axis=vocabulary_axis,
    )
Constructs an ``EmbeddingTable`` with uninitialized parameters.

    Args:
      vocab_size: The size of the vocabulary.
      embedding_axes: A dictionary mapping embedding axis names to their sizes.
        Will usually be a single-element dictionary of the form ``{"embedding":
        embedding_size}``.
      vocabulary_axis: The name of the axis that corresponds to the vocabulary.
        This axis will be indexed into when performing embedding lookups. Must
        not appear in ``embedding_axes``.
      initializer: A weight initializer that will be used to initialize the
        parameters. For the purposes of initialization, the "input axes" are of
        dimension 1, and the "output axes" are the ``embedding_axes``; the
        vocabulary axis is treated as a parallel axis. Defaults to fan-out
        normalization over the embedding axes.
      dtype: The data type of the embedding parameters.

    Returns:
      An EmbeddingTable with uninitialized parameters of the given shape.
    
penzai/deprecated/v1/nn/embeddings.py