ColossalAI/colossalai/kernel/jit/__init__.py

9 lines
308 B
Python
Raw Normal View History

2022-01-18 05:37:20 +00:00
from .option import set_jit_fusion_options
2022-01-13 08:47:17 +00:00
from .bias_dropout_add import bias_dropout_add_fused_train, bias_dropout_add_fused_inference
from .bias_gelu import bias_gelu_impl
2021-12-21 04:19:52 +00:00
2022-01-13 08:47:17 +00:00
__all__ = [
"bias_dropout_add_fused_train", "bias_dropout_add_fused_inference", "bias_gelu_impl",
2022-01-18 05:37:20 +00:00
"set_jit_fusion_options"
2022-01-13 08:47:17 +00:00
]