# Path: flaxformer/transformer_common.py
# class LayerSequence(nn.Module, param_remapping.ParameterRemappable):
# class TransparentLayerSequence:
#   def setup(self):
#   def __call__(self, inputs: Array, *args, **kwargs) -> Array:
#   def apply_range_of_layers(self, start_idx: int, end_idx: Optional[int],
#                             inputs: Array, *args, **kwargs) -> Array:
#   def __call__(self, inputs: Array, *args, **kwargs) -> Array:
#   def apply_range_of_layers(self, start_idx: int, end_idx: Optional[int],
#                             inputs: Array, *args, **kwargs) -> Array:


from absl.testing import absltest
from flax import linen as nn
from jax import random
from flaxformer import transformer_common as common
import jax.numpy as jnp
import numpy as np
    """Tests that a decoder layer stack outputs are of the correct shape."""
    batch_size, seq_len, embedding_size, hidden_size = 2, 3, 4, 5
    input_shape = (batch_size, seq_len, embedding_size)
    encoder_outputs_shape = (batch_size, seq_len, hidden_size)
    inputs = random.uniform(self.rng, input_shape, dtype=jnp.float32)
    encoder_outputs = random.uniform(
        self.rng, encoder_outputs_shape, dtype=jnp.float32)
    num_layers = 3

    def make_layer():
      return MockDecoderLayer(hidden_size)

    model = common.LayerSequence(num_layers=num_layers, make_layer=make_layer)
    result, _ = model.init_with_output(self.rng, inputs, encoder_outputs)
    expected = (batch_size, seq_len, hidden_size)
    self.assertEqual(expected, result.shape)

  def test_decoder_layer_stack_applies_layer_range_correctly(self):
    """Tests that a decoder layer stack correctly applies a range of layers."""
    batch_size, seq_len, embedding_size, hidden_size = 2, 3, 4, 5
    input_shape = (batch_size, seq_len, embedding_size)
    encoder_outputs_shape = (batch_size, seq_len, hidden_size)
    inputs = random.uniform(self.rng, input_shape, dtype=jnp.float32)
    encoder_outputs = random.uniform(
        self.rng, encoder_outputs_shape, dtype=jnp.float32)
    num_layers = 3

    def make_layer():
      return MockDecoderLayer(hidden_size)


# Path: flaxformer/transformer_common.py
# class LayerSequence(nn.Module, param_remapping.ParameterRemappable):
# class TransparentLayerSequence:
#   def setup(self):
#   def __call__(self, inputs: Array, *args, **kwargs) -> Array:
#   def apply_range_of_layers(self, start_idx: int, end_idx: Optional[int],
#                             inputs: Array, *args, **kwargs) -> Array:
#   def __call__(self, inputs: Array, *args, **kwargs) -> Array:
#   def apply_range_of_layers(self, start_idx: int, end_idx: Optional[int],
#                             inputs: Array, *args, **kwargs) -> Array:


# Path: flaxformer/transformer_common_test.py
from absl.testing import absltest
from flax import linen as nn
from jax import random
from flaxformer import transformer_common as common
import jax.numpy as jnp
import numpy as np

    """Tests that a decoder layer stack outputs are of the correct shape."""
    batch_size, seq_len, embedding_size, hidden_size = 2, 3, 4, 5
    input_shape = (batch_size, seq_len, embedding_size)
    encoder_outputs_shape = (batch_size, seq_len, hidden_size)
    inputs = random.uniform(self.rng, input_shape, dtype=jnp.float32)
    encoder_outputs = random.uniform(
        self.rng, encoder_outputs_shape, dtype=jnp.float32)
    num_layers = 3

    def make_layer():
      return MockDecoderLayer(hidden_size)

    model = common.LayerSequence(num_layers=num_layers, make_layer=make_layer)
    result, _ = model.init_with_output(self.rng, inputs, encoder_outputs)
    expected = (batch_size, seq_len, hidden_size)
    self.assertEqual(expected, result.shape)

  def test_decoder_layer_stack_applies_layer_range_correctly(self):
    """Tests that a decoder layer stack correctly applies a range of layers."""
    batch_size, seq_len, embedding_size, hidden_size = 2, 3, 4, 5
    input_shape = (batch_size, seq_len, embedding_size)
    encoder_outputs_shape = (batch_size, seq_len, hidden_size)
    inputs = random.uniform(self.rng, input_shape, dtype=jnp.float32)
    encoder_outputs = random.uniform(
        self.rng, encoder_outputs_shape, dtype=jnp.float32)
    num_layers = 3

    def make_layer():
      return MockDecoderLayer(hidden_size)


    model = common.LayerSequence(num_layers=num_layers, make_layer=make_layer)