# 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
      for idx, act_fn in enumerate(self.activations):
        x = jnp.squeeze(lax.dynamic_slice_in_dim(xs, idx, 1, -2), -2)
        x = _convert_to_activation_function(act_fn)(x)
        if idx == 0 and self.intermediate_conv is not None:
          x = self.intermediate_conv(  # pylint: disable=not-callable
              x,
              decode=decode,
              prefill=prefill,
              prefill_lengths=prefill_lengths)
        activations.append(x)

    # Take elementwise product of above intermediate activations.
    x = functools.reduce(operator.mul, activations)
    # Apply dropout and final dense output projection.
    # TODO: Change the `None` branch to not applying dropout
    # instead of fallback to default dropout.
    if self.intermediate_dropout:
      x = self.intermediate_dropout(x, deterministic=not enable_dropout)  # pylint: disable=not-callable
    else:
      x = nn.Dropout(
          rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
              x, deterministic=not enable_dropout)  # Broadcast along length.

    # Note: We don't use `activation_partitioning.with_sharding_migration` here
    # because we do often want this 2D sharded. However, if rules are valid,
    # they should result in 2D sharding. We don't need to raise errors if both
    # result in 2D sharding (which with_sharding_migration does).
    if partitioning.get_axis_rules():
      x = partitioning.with_sharding_constraint(x, ('batch', 'length', 'mlp'))
    else:

# 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

      for idx, act_fn in enumerate(self.activations):
        x = jnp.squeeze(lax.dynamic_slice_in_dim(xs, idx, 1, -2), -2)
        x = _convert_to_activation_function(act_fn)(x)
        if idx == 0 and self.intermediate_conv is not None:
          x = self.intermediate_conv(  # pylint: disable=not-callable
              x,
              decode=decode,
              prefill=prefill,
              prefill_lengths=prefill_lengths)
        activations.append(x)

    # Take elementwise product of above intermediate activations.
    x = functools.reduce(operator.mul, activations)
    # Apply dropout and final dense output projection.
    # TODO: Change the `None` branch to not applying dropout
    # instead of fallback to default dropout.
    if self.intermediate_dropout:
      x = self.intermediate_dropout(x, deterministic=not enable_dropout)  # pylint: disable=not-callable
    else:
      x = nn.Dropout(
          rate=self.intermediate_dropout_rate, broadcast_dims=(-2,))(
              x, deterministic=not enable_dropout)  # Broadcast along length.

    # Note: We don't use `activation_partitioning.with_sharding_migration` here
    # because we do often want this 2D sharded. However, if rules are valid,
    # they should result in 2D sharding. We don't need to raise errors if both
    # result in 2D sharding (which with_sharding_migration does).
    if partitioning.get_axis_rules():
      x = partitioning.with_sharding_constraint(x, ('batch', 'length', 'mlp'))
    else:

      x = activation_partitioning.with_sharding(