# 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/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py



import functools
import operator
import jax.numpy as jnp
import numpy as np
from typing import Any, Callable, Iterable, Mapping, Optional, Sequence, Tuple, Union
from flax import linen as nn
from flax.core import frozen_dict
from flax.linen import partitioning
from flax.linen.linear import default_kernel_init
from jax import lax
from flaxformer import activation_partitioning
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer
from aqt import flax_layers as aqt_flax_layers
from aqt import quant_config as aqt_config
from aqt import quantization as aqt
      features: tuple with numbers of output features.
      use_bias: whether to add a bias to the output (default: False).
      axis: tuple with axes to apply the transformation on.
      dtype: the dtype of the computation (default: float32).
      kernel_init: initializer function for the weight matrix.
      bias_init: initializer function for the bias.
      precision: numerical precision of the computation see `jax.lax.Precision`
        for details.
      kernel_axis_names: logical axis names to use for kernel sharding. Each
        should be one of _VALID_AXIS_NAMES in sharding.py.
      reshaped_kernel_axis_name_map: Rules for renaming fused kernel axes. We
        keep this as a separate parameter than kernel_axis_names so that
        experiments can toggle `reshape_kernel` without having to keep
        `kernel_axis_names` in sync.
      reshape_kernel: whether to reshape the kernel parameter to 2D for
        Adafactor.
  """
  features: Union[Iterable[int], int]
  use_bias: bool
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  kernel_init: Initializer = default_kernel_init
  bias_init: Initializer = nn.initializers.zeros
  precision: Any = None
  kernel_axis_names: Optional[Sequence[str]] = None
  reshaped_kernel_axis_name_map: Mapping[str, str] = (
      _RESHAPED_KERNEL_AXIS_NAME_MAP)
  reshape_kernel: bool = True

  @nn.compact

# 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/types.py

# 
# Path: flaxformer/types.py

# 
# Path: flaxformer/types.py



# Path: flaxformer/components/dense.py
import functools
import operator
import jax.numpy as jnp
import numpy as np
from typing import Any, Callable, Iterable, Mapping, Optional, Sequence, Tuple, Union
from flax import linen as nn
from flax.core import frozen_dict
from flax.linen import partitioning
from flax.linen.linear import default_kernel_init
from jax import lax
from flaxformer import activation_partitioning
from flaxformer.types import Array
from flaxformer.types import DType
from flaxformer.types import Initializer
from aqt import flax_layers as aqt_flax_layers
from aqt import quant_config as aqt_config
from aqt import quantization as aqt

      features: tuple with numbers of output features.
      use_bias: whether to add a bias to the output (default: False).
      axis: tuple with axes to apply the transformation on.
      dtype: the dtype of the computation (default: float32).
      kernel_init: initializer function for the weight matrix.
      bias_init: initializer function for the bias.
      precision: numerical precision of the computation see `jax.lax.Precision`
        for details.
      kernel_axis_names: logical axis names to use for kernel sharding. Each
        should be one of _VALID_AXIS_NAMES in sharding.py.
      reshaped_kernel_axis_name_map: Rules for renaming fused kernel axes. We
        keep this as a separate parameter than kernel_axis_names so that
        experiments can toggle `reshape_kernel` without having to keep
        `kernel_axis_names` in sync.
      reshape_kernel: whether to reshape the kernel parameter to 2D for
        Adafactor.
  """
  features: Union[Iterable[int], int]
  use_bias: bool
  axis: Union[Iterable[int], int] = -1
  dtype: DType = jnp.float32
  kernel_init: Initializer = default_kernel_init
  bias_init: Initializer = nn.initializers.zeros
  precision: Any = None
  kernel_axis_names: Optional[Sequence[str]] = None
  reshaped_kernel_axis_name_map: Mapping[str, str] = (
      _RESHAPED_KERNEL_AXIS_NAME_MAP)
  reshape_kernel: bool = True

  @nn.compact

  def __call__(self, inputs: Array) -> Array: