2021-12-21 04:19:52 +00:00
|
|
|
from .layer_norm import MixedFusedLayerNorm as LayerNorm
|
2023-08-04 05:46:22 +00:00
|
|
|
from .mha.mha import ColoAttention
|
2022-01-13 08:47:17 +00:00
|
|
|
from .multihead_attention import MultiHeadAttention
|
2023-02-03 01:47:13 +00:00
|
|
|
from .scaled_softmax import FusedScaleMaskSoftmax, ScaledUpperTriangMaskedSoftmax
|
|
|
|
|
|
2023-08-04 05:46:22 +00:00
|
|
|
__all__ = [
|
|
|
|
|
'LayerNorm', 'MultiHeadAttention', 'FusedScaleMaskSoftmax', 'ScaledUpperTriangMaskedSoftmax', 'ColoAttention'
|
|
|
|
|
]
|
