# 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
      embedding_shards.append(shard_embeddings)
    # RESULT: <float>[batch, seq, features]
    return jnp.concatenate(embedding_shards, axis=-1)


class NgramHashEmbed(nn.Module, Embedder[Array]):
  """Produces embeddings for ngrams of identifiers.

  This is similar to `HashEmbed`, but instead of embedding just the input IDs,
  it embeds ngrams of those IDs.

  Attributes:
    ngram_orders: The sizes of ngrams to embed.
    padding_id: The ID to be used for padding the ends of the inputs.
    features: Dimensionality of final embedding.
    num_embeddings_per_table: Size ("vocabulary") of each embedding table.
    num_tables: Number of embedding tables (a.k.a. hash functions / shards).
    cast_input_dtype: DType to cast input to.
    dtype: DType of resulting embeddings.
    embedding_init: Initializer for embeddings.
    one_hot: Performs the gather with a one-hot contraction rather than a true
      gather. This is currently needed for SPMD partitioning.
  """
  ngram_orders: Sequence[int]
  padding_id: int
  features: int
  num_embeddings_per_table: int
  num_tables: int = 8
  cast_input_dtype: Optional[DType] = None
  dtype: DType = jnp.float32

# 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

      embedding_shards.append(shard_embeddings)
    # RESULT: <float>[batch, seq, features]
    return jnp.concatenate(embedding_shards, axis=-1)


class NgramHashEmbed(nn.Module, Embedder[Array]):
  """Produces embeddings for ngrams of identifiers.

  This is similar to `HashEmbed`, but instead of embedding just the input IDs,
  it embeds ngrams of those IDs.

  Attributes:
    ngram_orders: The sizes of ngrams to embed.
    padding_id: The ID to be used for padding the ends of the inputs.
    features: Dimensionality of final embedding.
    num_embeddings_per_table: Size ("vocabulary") of each embedding table.
    num_tables: Number of embedding tables (a.k.a. hash functions / shards).
    cast_input_dtype: DType to cast input to.
    dtype: DType of resulting embeddings.
    embedding_init: Initializer for embeddings.
    one_hot: Performs the gather with a one-hot contraction rather than a true
      gather. This is currently needed for SPMD partitioning.
  """
  ngram_orders: Sequence[int]
  padding_id: int
  features: int
  num_embeddings_per_table: int
  num_tables: int = 8
  cast_input_dtype: Optional[DType] = None
  dtype: DType = jnp.float32

  embedding_init: Initializer = default_embed_init