DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        num_layers=121,  # Must be one of [121, 169, 201, 161]
        growth_rate=32,
        reduction=0.5,
        layer_rescale_factors={},
        lr_hparams={
            'schedule': 'constant',
            'base_lr': 0.2,
        },
        normalizer='batch_norm',
        optimizer='momentum',
        opt_hparams={
            'momentum': 0.9,
        },
        batch_size=128,
        l2_decay_factor=0.0001,
        l2_decay_rank_threshold=2,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
        grad_clip=None,
        normalize_classifier_input='none',
        classification_scale_factor=1.0,
    ))
class BottleneckBlock(nn.Module):
class TransitionBlock(nn.Module):
class DenseNet(nn.Module):
class AdaBeliefDensenetModel(base_model.BaseModel):
  def __call__(self, x, train):
  def __call__(self, x, train):
  def __call__(self, x, train):
    def dense_layers(y, block, num_blocks, growth_rate):
    def update_num_features(num_features, num_blocks, growth_rate, reduction):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        num_filters=16,
        num_layers=18,  # Must be one of [18, 34, 50, 101, 152, 200]
        layer_rescale_factors={},
        lr_hparams={
            'schedule': 'constant',
            'base_lr': 0.2,
        },
        optimizer='momentum',
        opt_hparams={
            'momentum': 0.9,
        },
        batch_size=128,
        l2_decay_factor=0.0001,
        l2_decay_rank_threshold=2,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        batch_norm_momentum=0.9,
        batch_norm_epsilon=1e-5,
        # Make this a string to avoid having to import jnp into the configs.
        model_dtype='float32',
        virtual_batch_size=None,
        total_accumulated_batch_size=None,
        data_format='NHWC',
        grad_clip=None,
    ))
class BasicResidualBlock(nn.Module):
class BottleneckResidualBlock(nn.Module):
class ResNet(nn.Module):
class AdaBeliefResnetModel(base_model.BaseModel):
  def __call__(self, x, train):
  def __call__(self, x, train):
  def __call__(self, x, train):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        num_layers=11,  # Must be one of [11, 13, 16, 19]
        layer_rescale_factors={},
        lr_hparams={
            'schedule': 'constant',
            'base_lr': 0.2,
        },
        normalizer='none',
        optimizer='momentum',
        opt_hparams={
            'momentum': 0.9,
        },
        batch_size=128,
        l2_decay_factor=0.0001,
        l2_decay_rank_threshold=2,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
        grad_clip=None,
    ))
def classifier(x, num_outputs, dropout_rate, deterministic):
def features(x, num_layers, normalizer, dtype, train):
  def __call__(self, x, train):
  def build_flax_module(self):
class VGG(nn.Module):
class AdaBeliefVGGModel(base_model.BaseModel): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        hid_sizes=[128, 64, 32, 64, 128],
        activation_function=['relu', 'relu', 'relu', 'relu', 'relu'],
        kernel_scales=[0.5, 0.5, 0.5, 0.5, 0.5, 0.5],
        lr_hparams={
            'base_lr': 0.1,
            'schedule': 'constant'
        },
        layer_rescale_factors={},
        optimizer='hessian_free',
        opt_hparams={
            'cg_max_iters': 250,
        },
        batch_size=128,
        l2_decay_factor=None,
        l2_decay_rank_threshold=0,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
    ))
class AutoEncoderModel(base_model.BaseModel):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        encoder={
            'filter_sizes': [16, 8, 8],
            'kernel_sizes': [(3, 3), (3, 3), (3, 3)],
            'kernel_paddings': ['SAME', 'SAME', 'SAME'],
            'window_sizes': [(2, 2), (2, 2), (2, 2)],
            'window_paddings': ['SAME', 'SAME', 'SAME'],
            'strides': [(2, 2), (2, 2), (2, 2)],
            'activations': ['relu', 'relu', 'relu'],
        },
        decoder={
            'filter_sizes': [8, 8, 16, 1],
            'kernel_sizes': [(3, 3), (3, 3), (3, 3), (3, 3)],
            'window_sizes': [(2, 2), (2, 2), (2, 2), None],
            'paddings': ['SAME', ((1, 0), (1, 0)), 'SAME', 'SAME'],
            'activations': ['relu', 'relu', 'relu', 'id'],
        },

        activation_function='relu',
        lr_hparams={
            'base_lr': 0.02,
            'schedule': 'constant'
        },
        layer_rescale_factors={},
        optimizer='momentum',
        opt_hparams={
            'momentum': 0,
        },
        batch_size=128,
        l2_decay_factor=None,
        l2_decay_rank_threshold=0,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
        grad_clip=None,
    ))
class ConvAutoEncoder(nn.Module):
class ConvAutoEncoderModel(base_model.BaseModel):
  def __call__(self, x, train):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        rng_seed=-1,
        model_dtype='float32',
        vocab_sizes=[1024] * 26,
        mlp_bottom_dims=[128, 128],
        mlp_top_dims=[256, 128, 1],
        output_shape=(1,),
        embed_dim=64,
        keep_diags=True,
        optimizer='adam',
        batch_size=128,
        num_dense_features=13,
        lr_hparams={
            'base_lr': 0.01,
            'schedule': 'constant'
        },
        opt_hparams={
            'beta1': 0.9,
            'beta2': 0.999,
            'epsilon': 1e-5,
        },
        l2_decay_factor=1e-5,
        l2_decay_rank_threshold=2,
    ))
def dot_interact(concat_features, keep_diags=True):
  def __call__(self, x, train):
    def scaled_init(key, shape, scale, init, dtype=jnp.float_):
  def build_flax_module(self):
class DLRM(nn.Module):
class DLRMModel(base_model.BaseModel): DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    hid_sizes=[20, 10],
    kernel_scales=[1.0, 1.0, 1.0],
    lr_hparams={
        'base_lr': 0.1,
        'schedule': 'constant'
    },
    layer_rescale_factors={},
    optimizer='momentum',
    opt_hparams={
        'momentum': 0.9,
    },
    batch_size=128,
    activation_function='relu',
    l2_decay_factor=.0005,
    l2_decay_rank_threshold=2,
    label_smoothing=None,
    rng_seed=-1,
    use_shallue_label_smoothing=False,
    model_dtype='float32',
    grad_clip=None,
))
class FullyConnected(nn.Module):
class FullyConnectedModel(base_model.BaseModel):
  def __call__(self, x, train):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        rng_seed=-1,
        model_dtype='float32',
        latent_dim=256,
        optimizer='adam',
        hidden_dims=(256,),
        batch_size=256,
        lr_hparams={
            'base_lr': 0.01,
            'schedule': 'constant'
        },
        opt_hparams={
            'beta1': 0.9,
            'beta2': 0.999,
            'epsilon': 1e-8,
        },
        activation_function='relu',
        l2_decay_factor=.0005,
        l2_decay_rank_threshold=2,
        num_message_passing_steps=5,
        normalizer='layer_norm',
        dropout_rate=0.1,
    ))
def _make_embed(latent_dim):
  def make_fn(inputs):
def _make_mlp(hidden_dims, maybe_normalize_fn, dropout):
  def make_fn(inputs):
  def __call__(self, graph, train):
  def get_fake_batch(self, hps):
  def build_flax_module(self):
class GNN(nn.Module):
class GNNModel(base_model.BaseModel): DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    num_filters=[64, 96, 128],
    kernel_sizes=[5, 3, 3],
    kernel_paddings=['VALID', 'VALID', 'SAME'],
    window_sizes=[3, 3, 3],
    window_paddings=['SAME', 'SAME', 'SAME'],
    strides=[2, 2, 2],
    num_dense_units=[512, 256],
    lr_hparams={
        'base_lr': 0.001,
        'schedule': 'constant'
    },
    layer_rescale_factors={},
    optimizer='momentum',
    opt_hparams={
        'momentum': 0.9,
    },
    batch_size=128,
    activation_fn='relu',
    normalizer='none',
    l2_decay_factor=.0005,
    l2_decay_rank_threshold=2,
    label_smoothing=None,
    rng_seed=-1,
    use_shallue_label_smoothing=False,
    model_dtype='float32',
    grad_clip=None,
))
class MaxPoolingCNN(nn.Module):
class MaxPoolingCNNModel(base_model.BaseModel):
  def __call__(self, x, train):
  def build_flax_module(self): FAKE_MODEL_DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    num_filters=16,
    num_layers=18,  # Must be one of [18, 34, 50, 101, 152, 200]
    layer_rescale_factors={},
    lr_hparams={
        'batch_size': 128,
        'base_lr': 10.0,
        'decay_end': -1,
        'end_lr': 1e-4,
        'power': 2.0,
        'schedule': 'mlperf_polynomial',
        'start_lr': 0.0,
        'steps_per_epoch': 10009.250000000002,
        'warmup_steps': 18,
    },
    optimizer='mlperf_lars_resnet',
    opt_hparams={
        'weight_decay': 2e-4,
        'beta': 0.9
    },
    batch_size=128,
    l2_decay_factor=None,
    l2_decay_rank_threshold=2,
    label_smoothing=.1,
    use_shallue_label_smoothing=False,
    model_dtype='float32',
    virtual_batch_size=64,
    data_format='NHWC',
))
MLPERF_DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    num_filters=16,
    # We set default to 18 for faster unit tests.
    num_layers=18,  # Must be one of [18, 34, 50, 101, 152, 200]
    layer_rescale_factors={},
    lr_hparams={
        'schedule': 'constant',
        'base_lr': 0.2,
    },
    optimizer='momentum',
    opt_hparams={
        'momentum': 0.9,
    },
    batch_size=128,
    bn_output_scale=0.0,
    l2_decay_factor=None,
    l2_decay_rank_threshold=2,
    label_smoothing=None,
    rng_seed=-1,
    use_shallue_label_smoothing=False,
    batch_norm_momentum=0.9,
    batch_norm_epsilon=1e-5,
    model_dtype='float32',
    virtual_batch_size=64,
    total_accumulated_batch_size=None,
    data_format='NHWC',
))
def _constant_init(factor):
  def init_fn(key, shape, dtype=jnp.float32):
  def __call__(self, x, train):
  def __call__(self, x, train):
  def __call__(self, x, train):
  def build_flax_module(self):
  def build_flax_module(self):
class ResidualBlock(nn.Module):
class ResNet(nn.Module):
class FakeResNet(nn.Module):
class ResnetModelMLPerf(base_model.BaseModel):
class FakeModel(base_model.BaseModel): DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    optimizer='momentum',
    opt_hparams={
        'momentum': 0.0,
    },
    lr_hparams={
        'base_lr': 0.1,
        'schedule': 'constant'
    },
    batch_size=128,
    rng_seed=-1,
    # Note the dimension is set by input_shape.
    hessian_decay_power=1,
    noise_decay_power=1,
    nqm_mode='diagH_diagC',
    model_dtype='float32',
    l2_decay_factor=None,
))
class NQMLoss(nn.Module):
class NQM(base_model.BaseModel):
  def __call__(self, noise_input, train):
def quadratic_form(u, sigma):
def _get_nqm_matrices(dim,
                      hessian_decay_power=1.0,
                      noise_decay_power=1.0,
                      mode='diagH_noC'):
  def __init__(self, hps, dataset_meta_data, loss_name, metrics_name):
  def evaluate_batch(self, params, batch_stats, batch):
  def training_cost(self, params, batch, batch_stats=None, dropout_rng=None):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    num_filters=16,
    num_layers=18,  # Must be one of [18, 34, 50, 101, 152, 200]
    layer_rescale_factors={},
    lr_hparams={
        'schedule': 'constant',
        'base_lr': 0.2,
    },
    optimizer='momentum',
    opt_hparams={
        'momentum': 0.9,
    },
    batch_size=128,
    l2_decay_factor=0.0001,
    l2_decay_rank_threshold=2,
    label_smoothing=None,
    rng_seed=-1,
    use_shallue_label_smoothing=False,
    batch_norm_momentum=0.9,
    batch_norm_epsilon=1e-5,
    # Make this a string to avoid having to import jnp into the configs.
    model_dtype='float32',
    virtual_batch_size=64,
    total_accumulated_batch_size=None,
    data_format='NHWC',
    block_type='post_activation',  # either pre_activation or post_activation
    bn_relu_conv=True,  # only used for block_type='pre_activation'
    use_bn=True,
    dropout_rate=0.0,
    grad_clip=None,
    activation_function='relu',
))
class PreActResidualBlock(nn.Module):
class ResidualBlock(nn.Module):
class ResNet(nn.Module):
class ResnetModel(base_model.BaseModel):
  def __call__(self, x, train):
    def maybe_normalize(name):
    def _bn_nonlin(y, name):
  def __call__(self, x, train):
    def maybe_normalize(name, scale_init=nn.initializers.ones):
  def __call__(self, x, train):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(dict(
    num_filters=[20, 10],
    kernel_sizes=[3, 3],
    lr_hparams={
        'base_lr': 0.001,
        'schedule': 'constant'
    },
    layer_rescale_factors={},
    optimizer='momentum',
    opt_hparams={
        'momentum': 0.9,
    },
    batch_size=128,
    activation_function='relu',
    l2_decay_factor=.0005,
    l2_decay_rank_threshold=2,
    label_smoothing=None,
    rng_seed=-1,
    use_shallue_label_smoothing=False,
    model_dtype='float32',
))
class SimpleCNN(nn.Module):
class SimpleCNNModel(base_model.BaseModel):
  def __call__(self, x, train):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        batch_size=512,
        emb_dim=128,
        num_heads=8,
        num_layers=6,
        qkv_dim=128,
        mlp_dim=512,
        dropout_rate=0.1,
        attention_dropout_rate=0.1,
        optimizer='adam',
        opt_hparams={
            'beta1': .9,
            'beta2': .98,
            'epsilon': 1e-9,
            'weight_decay': 0.1
        },
        layer_rescale_factors={},
        normalizer='layer_norm',
        lr_hparams={
            'base_lr': 0.0016,
            'warmup_steps': 1000,
            'squash_steps': 1000,
            'schedule': 'rsqrt_normalized_decay_warmup'
        },
        label_smoothing=None,
        l2_decay_factor=None,
        l2_decay_rank_threshold=0,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
        grad_clip=None,
        decode=False,
    ))
def shift_right(x, axis=1):
def shift_inputs(x, segment_ids=None, axis=1):
def sinusoidal_init(max_len=2048):
  def init(key, shape, dtype=np.float32):
  def __call__(self, inputs, inputs_positions=None):
  def __call__(self, inputs, train):
  def __call__(self,
               inputs,
               train,
               decoder_mask=None,
               encoder_decoder_mask=None,
               inputs_positions=None,
               inputs_segmentation=None):
  def __call__(self,
               inputs,
               train,
               inputs_positions=None,
               inputs_segmentation=None):
  def build_flax_module(self):
class AddPositionEmbs(nn.Module):
class MlpBlock(nn.Module):
class Transformer1DBlock(nn.Module):
class TransformerLM(nn.Module):
class TransformerLM1B(base_model.BaseModel): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        num_classes=1000,
        variant='B/32',
        rep_size=True,
        pool_type='gap',
        lr_hparams={
            'base_lr': 1e-3,
            'schedule': 'cosine_warmup',
            'warmup_steps': 10_000
        },
        optimizer='adam',
        opt_hparams={
            'beta1': 0.9,
            'beta2': 0.999,
            'epsilon': 1e-8,
            'weight_decay': 1e-1,
        },
        l2_decay_factor=None,
        batch_size=1024,
        rng_seed=-1,
        model_dtype='float32',
        grad_clip=None
    ))
def posemb_sincos_2d(h, w, width, temperature=10_000., dtype=jnp.float32):
def get_posemb(self, emb_type, seqshape, width, name, dtype=jnp.float32):
  def __call__(self, x, train=True):
  def __call__(self, x, train=True):
  def __call__(self, x, train=True):
  def __call__(self, x):
  def __call__(self, x, *, train=False):
  def build_flax_module(self):
def decode_variant(variant):
class MlpBlock(nn.Module):
class Encoder1DBlock(nn.Module):
class Encoder(nn.Module):
class MAPHead(nn.Module):
class ViT(nn.Module):
class ViTModel(base_model.BaseModel): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        blocks_per_group=3,
        channel_multiplier=2,
        lr_hparams={
            'base_lr': 0.001,
            'schedule': 'cosine'
        },
        normalizer='batch_norm',
        layer_rescale_factors={},
        conv_kernel_scale=1.0,
        dense_kernel_scale=1.0,
        dropout_rate=0.0,
        conv_kernel_init='lecun_normal',
        dense_kernel_init='lecun_normal',
        optimizer='momentum',
        opt_hparams={
            'momentum': 0.9,
        },
        batch_size=128,
        virtual_batch_size=None,
        total_accumulated_batch_size=None,
        l2_decay_factor=0.0001,
        l2_decay_rank_threshold=2,
        label_smoothing=None,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
        grad_clip=None,
        activation_function='relu',
        group_strides=[(1, 1), (2, 2), (2, 2)])
)
class WideResnetBlock(nn.Module):
class WideResnetGroup(nn.Module):
class WideResnet(nn.Module):
class WideResnetModel(base_model.BaseModel):
  def __call__(self, x, train):
  def __call__(self, x, train):
  def __call__(self, x, train):
  def build_flax_module(self): DEFAULT_HPARAMS = config_dict.ConfigDict(
    dict(
        batch_size=64,
        share_embeddings=False,
        logits_via_embedding=False,
        emb_dim=512,
        num_heads=8,
        enc_num_layers=4,
        dec_num_layers=4,
        qkv_dim=512,
        mlp_dim=512,
        dropout_rate=0.1,
        attention_dropout_rate=0.1,
        optimizer='adam',
        opt_hparams={
            'beta1': .9,
            'beta2': .98,
            'epsilon': 1e-9,
            'weight_decay': 0.0
        },
        layer_rescale_factors={},
        normalizer='layer_norm',
        lr_hparams={
            'base_lr': 0.05,
            'warmup_steps': 8000,
            'factors': 'constant * linear_warmup * rsqrt_decay',
            'schedule': 'compound'
        },
        label_smoothing=0.1,
        l2_decay_factor=None,
        l2_decay_rank_threshold=0,
        rng_seed=-1,
        use_shallue_label_smoothing=False,
        model_dtype='float32',
        grad_clip=None,
        enc_self_attn_kernel_init='xavier_uniform',
        dec_self_attn_kernel_init='xavier_uniform',
        dec_cross_attn_kernel_init='xavier_uniform',
        decode=False,
    ))
def _get_dtype(use_bfloat16):
def shift_right(x, axis=1):
def sinusoidal_init(max_len=2048, min_scale=1.0, max_scale=10000.0):
  def init(key, shape, dtype=np.float32):
  def __call__(self,
               inputs,
               inputs_positions=None,
               train=True):
  def __call__(self, inputs, train):
  def __call__(self,
               inputs,
               encoder_mask=None,
               train=True):
  def __call__(self,
               targets,
               encoded,
               decoder_mask=None,
               encoder_decoder_mask=None,
               train=True):
  def __call__(self,
               inputs,
               inputs_positions=None,
               encoder_mask=None,
               train=True):
  def __call__(self,
               encoded,
               targets,
               targets_positions=None,
               decoder_mask=None,
               encoder_decoder_mask=None,
               train=True):
  def setup(self):
  def __call__(self,
               inputs,
               targets,
               inputs_positions=None,
               targets_positions=None,
               inputs_segmentation=None,
               targets_segmentation=None,
               train=False):
  def encode(self,
             inputs,
             inputs_positions=None,
             inputs_segmentation=None,
             train=False):
  def decode(self,
             encoded,
             inputs,
             targets,
             targets_positions=None,
             inputs_segmentation=None,
             targets_segmentation=None,
             train=False):
  def __init__(self, hps, dataset_meta_data, loss_name, metrics_name):
  def evaluate_batch(self, params, batch_stats, batch):
  def training_cost(self, params, batch, batch_stats=None, dropout_rng=None):
  def build_flax_module(self):
class AddPositionEmbs(nn.Module):
class MlpBlock(nn.Module):
class Encoder1DBlock(nn.Module):
class EncoderDecoder1DBlock(nn.Module):
class Encoder(nn.Module):
class Decoder(nn.Module):
class Transformer(nn.Module):
class TransformerTranslate(base_model.BaseModel):
from init2winit.model_lib import adabelief_densenet
from init2winit.model_lib import adabelief_resnet
from init2winit.model_lib import adabelief_vgg
from init2winit.model_lib import autoencoder
from init2winit.model_lib import convolutional_autoencoder
from init2winit.model_lib import dlrm
from init2winit.model_lib import fully_connected
from init2winit.model_lib import gnn
from init2winit.model_lib import max_pooling_cnn
from init2winit.model_lib import mlperf_resnet
from init2winit.model_lib import nqm
from init2winit.model_lib import resnet
from init2winit.model_lib import simple_cnn
from init2winit.model_lib import transformer_lm
from init2winit.model_lib import vit
from init2winit.model_lib import wide_resnet
from init2winit.model_lib import xformer_translate
# coding=utf-8
# Copyright 2022 The init2winit Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Registry for the available models we can train."""


_ALL_MODELS = {
    'fully_connected':
        (fully_connected.FullyConnectedModel, fully_connected.DEFAULT_HPARAMS),
    'simple_cnn': (simple_cnn.SimpleCNNModel, simple_cnn.DEFAULT_HPARAMS),
    'max_pooling_cnn':
        (max_pooling_cnn.MaxPoolingCNNModel, max_pooling_cnn.DEFAULT_HPARAMS),
    'wide_resnet': (wide_resnet.WideResnetModel, wide_resnet.DEFAULT_HPARAMS),
    'resnet': (resnet.ResnetModel, resnet.DEFAULT_HPARAMS),
    'adabelief_densenet': (adabelief_densenet.AdaBeliefDensenetModel,
                           adabelief_densenet.DEFAULT_HPARAMS),
    'adabelief_resnet': (adabelief_resnet.AdaBeliefResnetModel,
                         adabelief_resnet.DEFAULT_HPARAMS),
    'adabelief_vgg':
        (adabelief_vgg.AdaBeliefVGGModel, adabelief_vgg.DEFAULT_HPARAMS),
    'autoencoder': (autoencoder.AutoEncoderModel, autoencoder.DEFAULT_HPARAMS),
    'convolutional_autoencoder':
        (convolutional_autoencoder.ConvAutoEncoderModel,
         convolutional_autoencoder.DEFAULT_HPARAMS),
    'fake_resnet':
        (mlperf_resnet.FakeModel, mlperf_resnet.FAKE_MODEL_DEFAULT_HPARAMS),
    'mlperf_resnet':
        (mlperf_resnet.ResnetModelMLPerf, mlperf_resnet.MLPERF_DEFAULT_HPARAMS),
    'transformer':
        (transformer_lm.TransformerLM1B, transformer_lm.DEFAULT_HPARAMS),
    'nqm': (nqm.NQM, nqm.DEFAULT_HPARAMS),
    'xformer_translate': (xformer_translate.TransformerTranslate,
                          xformer_translate.DEFAULT_HPARAMS),

    'gnn': (gnn.GNNModel, gnn.DEFAULT_HPARAMS),