cpu_adagrad.py 622 B

123456789101112131415161718192021222324252627
  1. # Copyright (c) Microsoft Corporation.
  2. # SPDX-License-Identifier: Apache-2.0
  3. # DeepSpeed Team
  4. from .builder import TorchCPUOpBuilder
  5. class CPUAdagradBuilder(TorchCPUOpBuilder):
  6. BUILD_VAR = "DS_BUILD_CPU_ADAGRAD"
  7. NAME = "cpu_adagrad"
  8. def __init__(self):
  9. super().__init__(name=self.NAME)
  10. def absolute_name(self):
  11. return f'deepspeed.ops.adagrad.{self.NAME}_op'
  12. def sources(self):
  13. return ['csrc/adagrad/cpu_adagrad.cpp']
  14. def libraries_args(self):
  15. args = super().libraries_args()
  16. return args
  17. def include_paths(self):
  18. return ['csrc/includes']