__init__.py 1.3 KB

12345678910111213141516171819202122232425262728293031323334353637
  1. """
  2. Copyright 2020 The Microsoft DeepSpeed Team
  3. """
  4. from .cpu_adam import CPUAdamBuilder
  5. from .cpu_adagrad import CPUAdagradBuilder
  6. from .fused_adam import FusedAdamBuilder
  7. from .fused_lamb import FusedLambBuilder
  8. from .sparse_attn import SparseAttnBuilder
  9. from .transformer import TransformerBuilder
  10. from .random_ltd import RandomLTDBuilder
  11. from .stochastic_transformer import StochasticTransformerBuilder
  12. from .utils import UtilsBuilder
  13. from .async_io import AsyncIOBuilder
  14. from .transformer_inference import InferenceBuilder
  15. from .quantizer import QuantizerBuilder
  16. from .spatial_inference import SpatialInferenceBuilder
  17. from .builder import get_default_compute_capabilities, OpBuilder
  18. # TODO: This will be removed eventurally when all files containing reference to ALL_OPS redirected to op_builder.all_ops
  19. # TODO: infer this list instead of hard coded
  20. # List of all available ops
  21. __op_builders__ = [
  22. CPUAdamBuilder(),
  23. CPUAdagradBuilder(),
  24. FusedAdamBuilder(),
  25. FusedLambBuilder(),
  26. SparseAttnBuilder(),
  27. TransformerBuilder(),
  28. StochasticTransformerBuilder(),
  29. AsyncIOBuilder(),
  30. UtilsBuilder(),
  31. QuantizerBuilder(),
  32. InferenceBuilder(),
  33. SpatialInferenceBuilder(),
  34. RandomLTDBuilder()
  35. ]
  36. ALL_OPS = {op.name: op for op in __op_builders__}