utils.py 387 B

123456789101112131415161718
  1. """
  2. Copyright 2020 The Microsoft DeepSpeed Team
  3. """
  4. from .builder import OpBuilder
  5. class UtilsBuilder(OpBuilder):
  6. BUILD_VAR = "DS_BUILD_UTILS"
  7. NAME = "utils"
  8. def __init__(self):
  9. super().__init__(name=self.NAME)
  10. def absolute_name(self):
  11. return f'deepspeed.ops.{self.NAME}_op'
  12. def sources(self):
  13. return ['csrc/utils/flatten_unflatten.cpp']