code/flash-linear-attention/fla/models/__init__.py

from fla.models.abc import ABCConfig, ABCForCausalLM, ABCModel
from fla.models.bitnet import BitNetConfig, BitNetForCausalLM, BitNetModel
from fla.models.comba import CombaConfig, CombaForCausalLM, CombaModel
from fla.models.delta_net import DeltaNetConfig, DeltaNetForCausalLM, DeltaNetModel
from fla.models.deltaformer import DeltaFormerConfig, DeltaFormerForCausalLM, DeltaFormerModel
from fla.models.forgetting_transformer import (
    ForgettingTransformerConfig,
    ForgettingTransformerForCausalLM,
    ForgettingTransformerModel,
)
from fla.models.gated_deltanet import GatedDeltaNetConfig, GatedDeltaNetForCausalLM, GatedDeltaNetModel
from fla.models.gated_deltaproduct import GatedDeltaProductConfig, GatedDeltaProductForCausalLM, GatedDeltaProductModel
from fla.models.gla import GLAConfig, GLAForCausalLM, GLAModel
from fla.models.gsa import GSAConfig, GSAForCausalLM, GSAModel
from fla.models.hgrn import HGRNConfig, HGRNForCausalLM, HGRNModel
from fla.models.hgrn2 import HGRN2Config, HGRN2ForCausalLM, HGRN2Model
from fla.models.kda import KDAConfig, KDAForCausalLM, KDAModel
from fla.models.lightnet import LightNetConfig, LightNetForCausalLM, LightNetModel
from fla.models.linear_attn import LinearAttentionConfig, LinearAttentionForCausalLM, LinearAttentionModel
from fla.models.log_linear_mamba2 import LogLinearMamba2Config, LogLinearMamba2ForCausalLM, LogLinearMamba2Model
from fla.models.mamba import MambaConfig, MambaForCausalLM, MambaModel
from fla.models.mamba2 import Mamba2Config, Mamba2ForCausalLM, Mamba2Model
from fla.models.mesa_net import MesaNetConfig, MesaNetForCausalLM, MesaNetModel
from fla.models.mla import MLAConfig, MLAForCausalLM, MLAModel
from fla.models.mom import MomConfig, MomForCausalLM, MomModel
from fla.models.nsa import NSAConfig, NSAForCausalLM, NSAModel
from fla.models.path_attn import PaTHAttentionConfig, PaTHAttentionForCausalLM, PaTHAttentionModel
from fla.models.retnet import RetNetConfig, RetNetForCausalLM, RetNetModel
from fla.models.rodimus import RodimusConfig, RodimusForCausalLM, RodimusModel
from fla.models.rwkv6 import RWKV6Config, RWKV6ForCausalLM, RWKV6Model
from fla.models.rwkv7 import RWKV7Config, RWKV7ForCausalLM, RWKV7Model
from fla.models.samba import SambaConfig, SambaForCausalLM, SambaModel
from fla.models.sse import SSEConfig, SSEForCausalLM, SSEModel
from fla.models.transformer import TransformerConfig, TransformerForCausalLM, TransformerModel

__all__ = [
    'ABCConfig', 'ABCForCausalLM', 'ABCModel',
    'BitNetConfig', 'BitNetForCausalLM', 'BitNetModel',
    'CombaConfig', 'CombaForCausalLM', 'CombaModel',
    'DeltaNetConfig', 'DeltaNetForCausalLM', 'DeltaNetModel',
    'DeltaFormerConfig', 'DeltaFormerForCausalLM', 'DeltaFormerModel',
    'ForgettingTransformerConfig', 'ForgettingTransformerForCausalLM', 'ForgettingTransformerModel',
    'GatedDeltaNetConfig', 'GatedDeltaNetForCausalLM', 'GatedDeltaNetModel',
    'GatedDeltaProductConfig', 'GatedDeltaProductForCausalLM', 'GatedDeltaProductModel',
    'GLAConfig', 'GLAForCausalLM', 'GLAModel',
    'GSAConfig', 'GSAForCausalLM', 'GSAModel',
    'HGRNConfig', 'HGRNForCausalLM', 'HGRNModel',
    'HGRN2Config', 'HGRN2ForCausalLM', 'HGRN2Model',
    'KDAConfig', 'KDAForCausalLM', 'KDAModel',
    'LightNetConfig', 'LightNetForCausalLM', 'LightNetModel',
    'LinearAttentionConfig', 'LinearAttentionForCausalLM', 'LinearAttentionModel',
    'LogLinearMamba2Config', 'LogLinearMamba2ForCausalLM', 'LogLinearMamba2Model',
    'MambaConfig', 'MambaForCausalLM', 'MambaModel',
    'Mamba2Config', 'Mamba2ForCausalLM', 'Mamba2Model',
    'MesaNetConfig', 'MesaNetForCausalLM', 'MesaNetModel',
    'MomConfig', 'MomForCausalLM', 'MomModel',
    'MLAConfig', 'MLAForCausalLM', 'MLAModel',
    'NSAConfig', 'NSAForCausalLM', 'NSAModel',
    'PaTHAttentionConfig', 'PaTHAttentionForCausalLM', 'PaTHAttentionModel',
    'RetNetConfig', 'RetNetForCausalLM', 'RetNetModel',
    'RodimusConfig', 'RodimusForCausalLM', 'RodimusModel',
    'RWKV6Config', 'RWKV6ForCausalLM', 'RWKV6Model',
    'RWKV7Config', 'RWKV7ForCausalLM', 'RWKV7Model',
    'SambaConfig', 'SambaForCausalLM', 'SambaModel',
    'SSEConfig', 'SSEForCausalLM', 'SSEModel',
    'TransformerConfig', 'TransformerForCausalLM', 'TransformerModel',
]