Activation.h 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. #pragma once
  3. #include <ATen/native/DispatchStub.h>
  4. #include <ATen/native/Gelu.h>
  5. #include <c10/util/Exception.h>
  6. namespace c10 {
  7. class Scalar;
  8. }
  9. namespace at {
  10. struct TensorIterator;
  11. struct TensorIteratorBase;
  12. class TensorBase;
  13. }
  14. namespace at::native {
  15. using structured_activation_fn = void (*)(TensorIteratorBase&);
  16. using structured_activation_backward_fn = void (*)(TensorIteratorBase&);
  17. using activation_fn = void (*)(TensorIterator&);
  18. using activation_backward_fn = void (*)(TensorIterator&);
  19. using softplus_fn = void (*)(TensorIteratorBase&, const c10::Scalar&, const c10::Scalar&);
  20. using softplus_backward_fn = void (*)(TensorIteratorBase&, const c10::Scalar&, const c10::Scalar&);
  21. using threshold_fn = void (*)(TensorIteratorBase&, const c10::Scalar&, const c10::Scalar&);
  22. using hardtanh_backward_fn = void (*)(TensorIterator&, const c10::Scalar&, const c10::Scalar&);
  23. using hardsigmoid_fn = void(*)(TensorIteratorBase&);
  24. using hardsigmoid_backward_fn = void(*)(TensorIteratorBase&);
  25. using hardswish_fn = void(*)(TensorIterator&);
  26. using hardswish_backward_fn = void(*)(TensorIterator&);
  27. using shrink_fn = void (*)(TensorIteratorBase&, const c10::Scalar&);
  28. using softshrink_fn = void (*)(TensorIteratorBase&, const c10::Scalar&);
  29. using shrink_backward_fn = void (*)(TensorIteratorBase&, const c10::Scalar&);
  30. using elu_fn = void (*)(TensorIteratorBase&, const c10::Scalar&, const c10::Scalar&, const c10::Scalar&);
  31. using elu_backward_fn = void (*)(TensorIteratorBase&, const c10::Scalar&, const c10::Scalar&, const c10::Scalar&, bool);
  32. using leaky_relu_fn = void (*)(TensorIteratorBase&, const c10::Scalar&);
  33. using leaky_relu_backward_fn = void (*)(TensorIteratorBase&, const c10::Scalar&);
  34. using log_sigmoid_cpu_fn = void (*)(TensorBase&, TensorBase&, const TensorBase&);
  35. using gelu_fn = void (*)(TensorIteratorBase&, GeluType);
  36. using gelu_backward_fn = void (*)(TensorIteratorBase&, GeluType);
  37. using glu_jvp_fn = void (*)(TensorIteratorBase&);
  38. DECLARE_DISPATCH(elu_fn, elu_stub)
  39. DECLARE_DISPATCH(elu_backward_fn, elu_backward_stub)
  40. DECLARE_DISPATCH(softplus_fn, softplus_stub)
  41. DECLARE_DISPATCH(softplus_backward_fn, softplus_backward_stub)
  42. DECLARE_DISPATCH(log_sigmoid_cpu_fn, log_sigmoid_cpu_stub)
  43. DECLARE_DISPATCH(activation_backward_fn, log_sigmoid_backward_stub)
  44. DECLARE_DISPATCH(threshold_fn, threshold_stub)
  45. DECLARE_DISPATCH(gelu_fn, GeluKernel)
  46. DECLARE_DISPATCH(gelu_backward_fn, GeluBackwardKernel)
  47. DECLARE_DISPATCH(hardtanh_backward_fn, hardtanh_backward_stub)
  48. DECLARE_DISPATCH(hardsigmoid_fn, hardsigmoid_stub)
  49. DECLARE_DISPATCH(hardsigmoid_backward_fn, hardsigmoid_backward_stub)
  50. DECLARE_DISPATCH(hardswish_fn, hardswish_stub)
  51. DECLARE_DISPATCH(hardswish_backward_fn, hardswish_backward_stub)
  52. DECLARE_DISPATCH(shrink_fn, hardshrink_stub)
  53. DECLARE_DISPATCH(softshrink_fn, softshrink_stub)
  54. DECLARE_DISPATCH(shrink_backward_fn, shrink_backward_stub)
  55. DECLARE_DISPATCH(leaky_relu_fn, leaky_relu_stub)
  56. DECLARE_DISPATCH(leaky_relu_backward_fn, leaky_relu_backward_stub)
  57. DECLARE_DISPATCH(structured_activation_fn, glu_stub)
  58. DECLARE_DISPATCH(activation_backward_fn, glu_backward_stub)
  59. DECLARE_DISPATCH(glu_jvp_fn, glu_jvp_stub)
  60. DECLARE_DISPATCH(structured_activation_fn, silu_stub)
  61. DECLARE_DISPATCH(structured_activation_backward_fn, silu_backward_stub)
  62. DECLARE_DISPATCH(structured_activation_fn, mish_stub)
  63. DECLARE_DISPATCH(activation_backward_fn, mish_backward_stub)
  64. DECLARE_DISPATCH(activation_fn, prelu_stub)
  65. DECLARE_DISPATCH(activation_backward_fn, prelu_backward_stub)
  66. } // namespace at::native
  67. #else
  68. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  69. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)