# 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
      decode: Parameter to pass through to all embedders.
      deterministic: Parameter to pass through to all embedders.
      **input_kwargs: The input tensors to be embedded, with a name that matches
        the embedder in self.embedders, and each shaped: <int64>[..., seq_len].

    Returns:
      A tensor with the combined embeddings <float32>[..., embedding_size] in
      case of summing, and
      <float32>[..., embedding_size_1 + embedding_size_2 + ..] in case of
      concatenation.
    """
    return self.combine_embeddings(
        self.get_individual_embeddings(
            segment_ids=segment_ids,
            decode=decode,
            deterministic=deterministic,
            **input_kwargs),
        combine_method=combine_method)


class FixedEmbed(nn.Module, EmbedderWithDecode[Array]):
  """Fixed (not learnable) embeddings specified by the initializer function.

  Attributes:
    init_fn: The initializer function that defines the embeddings.
    max_length: The maximum supported length.
    dtype: The DType to use for the embeddings.
  """
  features: int
  max_length: int = 2048

# 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

      decode: Parameter to pass through to all embedders.
      deterministic: Parameter to pass through to all embedders.
      **input_kwargs: The input tensors to be embedded, with a name that matches
        the embedder in self.embedders, and each shaped: <int64>[..., seq_len].

    Returns:
      A tensor with the combined embeddings <float32>[..., embedding_size] in
      case of summing, and
      <float32>[..., embedding_size_1 + embedding_size_2 + ..] in case of
      concatenation.
    """
    return self.combine_embeddings(
        self.get_individual_embeddings(
            segment_ids=segment_ids,
            decode=decode,
            deterministic=deterministic,
            **input_kwargs),
        combine_method=combine_method)


class FixedEmbed(nn.Module, EmbedderWithDecode[Array]):
  """Fixed (not learnable) embeddings specified by the initializer function.

  Attributes:
    init_fn: The initializer function that defines the embeddings.
    max_length: The maximum supported length.
    dtype: The DType to use for the embeddings.
  """
  features: int
  max_length: int = 2048

  embedding_init: Initializer = initializers.sinusoidal()