# 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
# See the License for the specific language governing permissions and
# limitations under the License.

"""A library with embedding classes and functions."""




# Note: We don't use this in real models, but keep the default initializers the
# same as in Flax.
default_embed_init = nn.initializers.variance_scaling(
    1.0, 'fan_in', 'normal', out_axis=0)


@enum.unique
class EmbedCombineMethod(enum.IntEnum):
  # We use IntEnum here so that this class is serializable. Enum isn't.
  SUM = 1
  CONCAT = 2


_Inputs = TypeVar('_Inputs')


class Embedder(Generic[_Inputs], metaclass=abc.ABCMeta):

  @abc.abstractmethod
  def __call__(self,
               inputs: _Inputs,
               *,

# 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

# See the License for the specific language governing permissions and
# limitations under the License.

"""A library with embedding classes and functions."""




# Note: We don't use this in real models, but keep the default initializers the
# same as in Flax.
default_embed_init = nn.initializers.variance_scaling(
    1.0, 'fan_in', 'normal', out_axis=0)


@enum.unique
class EmbedCombineMethod(enum.IntEnum):
  # We use IntEnum here so that this class is serializable. Enum isn't.
  SUM = 1
  CONCAT = 2


_Inputs = TypeVar('_Inputs')


class Embedder(Generic[_Inputs], metaclass=abc.ABCMeta):

  @abc.abstractmethod
  def __call__(self,
               inputs: _Inputs,
               *,

               segment_ids: Optional[Array] = None,