import transformers


class MLPConfig(transformers.PretrainedConfig):
  model_type = 'MLP'

  def __init__(
    self,
    input_dim: int = 16,
    hidden_dim: int = 32,
    output_dim: int = 16,
    **kwargs):
    super().__init__(**kwargs)
    self.input_dim = input_dim
    self.hidden_dim = hidden_dim
    self.output_dim = output_dim