# Path: vmoe/nn/external.py
# class MlpMixer(models.MlpMixer):
# class VisionTransformer(models.VisionTransformer):
#   def __post_init__(self):
#   def __call__(self, inputs):
#   def __post_init__(self):
#   def __call__(self, inputs):


import abc
import chex
import jax
import jax.numpy as jnp
import ml_collections
from absl.testing import absltest
from vmoe.nn import external
                                           grads_norm, zeros)


class MlpMixerTest(absltest.TestCase, BaseExternalTest):

  def create_model(self, deterministic: bool):
    config = ml_collections.ConfigDict()
    config.patches = ml_collections.ConfigDict({'size': (8, 8)})
    config.hidden_dim = 64
    config.num_blocks = 1
    config.tokens_mlp_dim = 128
    config.channels_mlp_dim = 128
    config.num_classes = 10
    return external.MlpMixer(**config, deterministic=deterministic)


class VisionTransformerTest(absltest.TestCase, BaseExternalTest):

  def create_model(self, deterministic: bool):
    config = ml_collections.ConfigDict()
    config.patches = ml_collections.ConfigDict({'size': (8, 8)})
    config.hidden_size = 64
    config.num_classes = 10
    config.transformer = ml_collections.ConfigDict({
        'num_layers': 1,
        'mlp_dim': 128,
        'num_heads': 2,
        'dropout_rate': 0.0,
        'attention_dropout_rate': 0.0,
    })

# Path: vmoe/nn/external.py
# class MlpMixer(models.MlpMixer):
# class VisionTransformer(models.VisionTransformer):
#   def __post_init__(self):
#   def __call__(self, inputs):
#   def __post_init__(self):
#   def __call__(self, inputs):


# Path: vmoe/nn/external_test.py
import abc
import chex
import jax
import jax.numpy as jnp
import ml_collections
from absl.testing import absltest
from vmoe.nn import external

                                           grads_norm, zeros)


class MlpMixerTest(absltest.TestCase, BaseExternalTest):

  def create_model(self, deterministic: bool):
    config = ml_collections.ConfigDict()
    config.patches = ml_collections.ConfigDict({'size': (8, 8)})
    config.hidden_dim = 64
    config.num_blocks = 1
    config.tokens_mlp_dim = 128
    config.channels_mlp_dim = 128
    config.num_classes = 10
    return external.MlpMixer(**config, deterministic=deterministic)


class VisionTransformerTest(absltest.TestCase, BaseExternalTest):

  def create_model(self, deterministic: bool):
    config = ml_collections.ConfigDict()
    config.patches = ml_collections.ConfigDict({'size': (8, 8)})
    config.hidden_size = 64
    config.num_classes = 10
    config.transformer = ml_collections.ConfigDict({
        'num_layers': 1,
        'mlp_dim': 128,
        'num_heads': 2,
        'dropout_rate': 0.0,
        'attention_dropout_rate': 0.0,
    })

    return external.VisionTransformer(**config, deterministic=deterministic)