|
|
|
|
|
|
|
|
from fla.modules.convolution import ImplicitLongConvolution, LongConvolution, ShortConvolution |
|
|
from fla.modules.fused_bitlinear import BitLinear, FusedBitLinear |
|
|
from fla.modules.fused_cross_entropy import FusedCrossEntropyLoss |
|
|
from fla.modules.fused_kl_div import FusedKLDivLoss |
|
|
from fla.modules.fused_linear_cross_entropy import FusedLinearCrossEntropyLoss |
|
|
from fla.modules.fused_linear_listnet_loss import FusedLinearListNetLoss |
|
|
from fla.modules.fused_norm_gate import ( |
|
|
FusedLayerNormGated, |
|
|
FusedLayerNormSwishGate, |
|
|
FusedLayerNormSwishGateLinear, |
|
|
FusedRMSNormGated, |
|
|
FusedRMSNormSwishGate, |
|
|
FusedRMSNormSwishGateLinear |
|
|
) |
|
|
from fla.modules.layernorm import GroupNorm, GroupNormLinear, LayerNorm, LayerNormLinear, RMSNorm, RMSNormLinear |
|
|
from fla.modules.mlp import GatedMLP |
|
|
from fla.modules.rotary import RotaryEmbedding |
|
|
|
|
|
__all__ = [ |
|
|
'ImplicitLongConvolution', 'LongConvolution', 'ShortConvolution', |
|
|
'BitLinear', 'FusedBitLinear', |
|
|
'FusedCrossEntropyLoss', 'FusedLinearCrossEntropyLoss', 'FusedKLDivLoss', |
|
|
'GroupNorm', 'GroupNormLinear', 'LayerNorm', 'LayerNormLinear', 'RMSNorm', 'RMSNormLinear', |
|
|
'FusedLayerNormGated', 'FusedLayerNormSwishGate', 'FusedLayerNormSwishGateLinear', |
|
|
'FusedRMSNormGated', 'FusedRMSNormSwishGate', 'FusedRMSNormSwishGateLinear', |
|
|
'GatedMLP', |
|
|
'RotaryEmbedding' |
|
|
] |
|
|
|
