# Path: flaxformer/activation_partitioning.py
# def global_mesh_defined():
# def with_sharding(x, partitioning_dims: int):
# def with_sharding_migration(
#     x: T,
#     activation_partitioning_dims: Optional[int],
#     logical_axis_names: Tuple[str, ...],
# ) -> T:
# T = TypeVar('T')
# 
# Path: flaxformer/components/initializers.py
# def sinusoidal(min_scale: float = 1.0,
#                max_scale: float = 10000.0,
#                dtype: DType = jnp.float32) -> Initializer:
#   def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
# def truncated_normal(mean: Union[float, Array] = 0.0,
#                      stddev: Union[float, Array] = 0.05,
#                      dtype: DType = jnp.float32) -> Initializer:
#   def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py



import abc
import collections
import dataclasses
import enum
import functools
import math
import chex
import jax
import numpy as np
from typing import (Any, Callable, DefaultDict, Dict, Generic, List, Mapping,
                    Optional, Sequence, TypeVar, Union)
from flax import linen as nn
from flax.linen import partitioning
from flax.training import common_utils
from jax import lax
from jax import numpy as jnp
from flaxformer import activation_partitioning
from flaxformer.components import initializers
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer
    return self.embeddings_combiner(embeddings)


class EmbedderWithDecode(Generic[_Inputs], Embedder[_Inputs]):
  """Denotes embedder classes that support the `decode` parameter."""
  pass


class EmbedderWithDeterministic(Generic[_Inputs], Embedder[_Inputs]):
  """Denotes embedder classes that support the `deterministic` parameter."""
  pass


class Embed(nn.Module, Embedder[Array]):
  """An embedder for `Array`s.

  A parameterized function from integers [0, n) to d-dimensional vectors.

  Attributes:
    num_embeddings: number of embeddings.
    features: number of feature dimensions for each embedding.
    dtype: the dtype of the embedding vectors (default: float32).
    embedding_init: embedding initializer.
    one_hot: performs the gather with a one-hot contraction rather than a true
      gather. This is currently needed for SPMD partitioning.
    axes: default axis metadata names for the embedding table.
  """
  num_embeddings: int
  features: int
  cast_input_dtype: Optional[DType] = None

# Path: flaxformer/activation_partitioning.py
# def global_mesh_defined():
# def with_sharding(x, partitioning_dims: int):
# def with_sharding_migration(
#     x: T,
#     activation_partitioning_dims: Optional[int],
#     logical_axis_names: Tuple[str, ...],
# ) -> T:
# T = TypeVar('T')
# 
# Path: flaxformer/components/initializers.py
# def sinusoidal(min_scale: float = 1.0,
#                max_scale: float = 10000.0,
#                dtype: DType = jnp.float32) -> Initializer:
#   def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
# def truncated_normal(mean: Union[float, Array] = 0.0,
#                      stddev: Union[float, Array] = 0.05,
#                      dtype: DType = jnp.float32) -> Initializer:
#   def init(key: PRNGKey, shape: Shape, dtype: DType = dtype) -> Array:
# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py



# Path: flaxformer/components/embedding.py
import abc
import collections
import dataclasses
import enum
import functools
import math
import chex
import jax
import numpy as np
from typing import (Any, Callable, DefaultDict, Dict, Generic, List, Mapping,
                    Optional, Sequence, TypeVar, Union)
from flax import linen as nn
from flax.linen import partitioning
from flax.training import common_utils
from jax import lax
from jax import numpy as jnp
from flaxformer import activation_partitioning
from flaxformer.components import initializers
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer

    return self.embeddings_combiner(embeddings)


class EmbedderWithDecode(Generic[_Inputs], Embedder[_Inputs]):
  """Denotes embedder classes that support the `decode` parameter."""
  pass


class EmbedderWithDeterministic(Generic[_Inputs], Embedder[_Inputs]):
  """Denotes embedder classes that support the `deterministic` parameter."""
  pass


class Embed(nn.Module, Embedder[Array]):
  """An embedder for `Array`s.

  A parameterized function from integers [0, n) to d-dimensional vectors.

  Attributes:
    num_embeddings: number of embeddings.
    features: number of feature dimensions for each embedding.
    dtype: the dtype of the embedding vectors (default: float32).
    embedding_init: embedding initializer.
    one_hot: performs the gather with a one-hot contraction rather than a true
      gather. This is currently needed for SPMD partitioning.
    axes: default axis metadata names for the embedding table.
  """
  num_embeddings: int
  features: int
  cast_input_dtype: Optional[DType] = None

  dtype: DType = jnp.float32