# Path: vmoe/nn/vit_moe.py
# def constant_initializer(constant: float) -> InitializerFn:
#   def f(key: PRNGKey, shape: Shape, dtype: DType = jnp.float_) -> Array:
#   def __call__(self, inputs):
#   def __call__(self, inputs):
#   def __call__(self, inputs):
#   def __call__(self, inputs):
#   def __call__(self, inputs):
# class MlpBlock(vit_models.MlpBlock):
# class MlpMoeWithNoisyTopExpertsPerItemRouter(nn.Module):
# class EncoderBlock(nn.Module):
# class EncoderMoe(nn.Module):
# class VisionTransformerMoe(nn.Module):


from absl.testing import absltest
from vmoe.nn import vit_moe
import jax
    shapes = shapes.unfreeze()
    expected_shapes = {
        'params': {
            'Encoder': {
                'encoder_norm': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                'posembed_input': {'pos_embedding': (1, 4, 8)},
                'encoderblock_0': {
                    'LayerNorm_0': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                    'LayerNorm_1': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                    'SelfAttention': EXPECTED_DEFAULT_ATTENTION_SHAPES,
                    'Mlp': EXPECTED_DEFAULT_MLP_SHAPES,
                },
                'encoderblock_1': {
                    'LayerNorm_0': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                    'LayerNorm_1': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                    'SelfAttention': EXPECTED_DEFAULT_ATTENTION_SHAPES,
                    'Moe': {
                        'Mlp': EXPECTED_DEFAULT_MOE_SHAPES,
                        'Router': {'dense': {'kernel': (8, 4)}},
                    },
                },
            },
            'embedding': {'bias': (8,), 'kernel': (2, 2, 3, 8)},
            'head': {'bias': (4,), 'kernel': (8, 4)},
        }
    }
    self.assertDictEqual(shapes, expected_shapes)

  def test_forward(self):
    """Tests that the model runs in forward mode. Correctness is not tested."""

# Path: vmoe/nn/vit_moe.py
# def constant_initializer(constant: float) -> InitializerFn:
#   def f(key: PRNGKey, shape: Shape, dtype: DType = jnp.float_) -> Array:
#   def __call__(self, inputs):
#   def __call__(self, inputs):
#   def __call__(self, inputs):
#   def __call__(self, inputs):
#   def __call__(self, inputs):
# class MlpBlock(vit_models.MlpBlock):
# class MlpMoeWithNoisyTopExpertsPerItemRouter(nn.Module):
# class EncoderBlock(nn.Module):
# class EncoderMoe(nn.Module):
# class VisionTransformerMoe(nn.Module):


# Path: vmoe/nn/vit_moe_test.py
from absl.testing import absltest
from vmoe.nn import vit_moe
import jax

    shapes = shapes.unfreeze()
    expected_shapes = {
        'params': {
            'Encoder': {
                'encoder_norm': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                'posembed_input': {'pos_embedding': (1, 4, 8)},
                'encoderblock_0': {
                    'LayerNorm_0': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                    'LayerNorm_1': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                    'SelfAttention': EXPECTED_DEFAULT_ATTENTION_SHAPES,
                    'Mlp': EXPECTED_DEFAULT_MLP_SHAPES,
                },
                'encoderblock_1': {
                    'LayerNorm_0': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                    'LayerNorm_1': EXPECTED_DEFAULT_LAYER_NORM_SHAPES,
                    'SelfAttention': EXPECTED_DEFAULT_ATTENTION_SHAPES,
                    'Moe': {
                        'Mlp': EXPECTED_DEFAULT_MOE_SHAPES,
                        'Router': {'dense': {'kernel': (8, 4)}},
                    },
                },
            },
            'embedding': {'bias': (8,), 'kernel': (2, 2, 3, 8)},
            'head': {'bias': (4,), 'kernel': (8, 4)},
        }
    }
    self.assertDictEqual(shapes, expected_shapes)

  def test_forward(self):
    """Tests that the model runs in forward mode. Correctness is not tested."""

    model = vit_moe.VisionTransformerMoe(**DEFAULT_TEST_CONFIG)