operators.h 9.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. /*
  3. pybind11/operator.h: Metatemplates for operator overloading
  4. Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
  5. All rights reserved. Use of this source code is governed by a
  6. BSD-style license that can be found in the LICENSE file.
  7. */
  8. #pragma once
  9. #include "pybind11.h"
  10. PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
  11. PYBIND11_NAMESPACE_BEGIN(detail)
  12. /// Enumeration with all supported operator types
  13. enum op_id : int {
  14. op_add,
  15. op_sub,
  16. op_mul,
  17. op_div,
  18. op_mod,
  19. op_divmod,
  20. op_pow,
  21. op_lshift,
  22. op_rshift,
  23. op_and,
  24. op_xor,
  25. op_or,
  26. op_neg,
  27. op_pos,
  28. op_abs,
  29. op_invert,
  30. op_int,
  31. op_long,
  32. op_float,
  33. op_str,
  34. op_cmp,
  35. op_gt,
  36. op_ge,
  37. op_lt,
  38. op_le,
  39. op_eq,
  40. op_ne,
  41. op_iadd,
  42. op_isub,
  43. op_imul,
  44. op_idiv,
  45. op_imod,
  46. op_ilshift,
  47. op_irshift,
  48. op_iand,
  49. op_ixor,
  50. op_ior,
  51. op_complex,
  52. op_bool,
  53. op_nonzero,
  54. op_repr,
  55. op_truediv,
  56. op_itruediv,
  57. op_hash
  58. };
  59. enum op_type : int {
  60. op_l, /* base type on left */
  61. op_r, /* base type on right */
  62. op_u /* unary operator */
  63. };
  64. struct self_t {};
  65. static const self_t self = self_t();
  66. /// Type for an unused type slot
  67. struct undefined_t {};
  68. /// Don't warn about an unused variable
  69. inline self_t __self() { return self; }
  70. /// base template of operator implementations
  71. template <op_id, op_type, typename B, typename L, typename R>
  72. struct op_impl {};
  73. /// Operator implementation generator
  74. template <op_id id, op_type ot, typename L, typename R>
  75. struct op_ {
  76. static constexpr bool op_enable_if_hook = true;
  77. template <typename Class, typename... Extra>
  78. void execute(Class &cl, const Extra &...extra) const {
  79. using Base = typename Class::type;
  80. using L_type = conditional_t<std::is_same<L, self_t>::value, Base, L>;
  81. using R_type = conditional_t<std::is_same<R, self_t>::value, Base, R>;
  82. using op = op_impl<id, ot, Base, L_type, R_type>;
  83. cl.def(op::name(), &op::execute, is_operator(), extra...);
  84. }
  85. template <typename Class, typename... Extra>
  86. void execute_cast(Class &cl, const Extra &...extra) const {
  87. using Base = typename Class::type;
  88. using L_type = conditional_t<std::is_same<L, self_t>::value, Base, L>;
  89. using R_type = conditional_t<std::is_same<R, self_t>::value, Base, R>;
  90. using op = op_impl<id, ot, Base, L_type, R_type>;
  91. cl.def(op::name(), &op::execute_cast, is_operator(), extra...);
  92. }
  93. };
  94. #define PYBIND11_BINARY_OPERATOR(id, rid, op, expr) \
  95. template <typename B, typename L, typename R> \
  96. struct op_impl<op_##id, op_l, B, L, R> { \
  97. static char const *name() { return "__" #id "__"; } \
  98. static auto execute(const L &l, const R &r) -> decltype(expr) { return (expr); } \
  99. static B execute_cast(const L &l, const R &r) { return B(expr); } \
  100. }; \
  101. template <typename B, typename L, typename R> \
  102. struct op_impl<op_##id, op_r, B, L, R> { \
  103. static char const *name() { return "__" #rid "__"; } \
  104. static auto execute(const R &r, const L &l) -> decltype(expr) { return (expr); } \
  105. static B execute_cast(const R &r, const L &l) { return B(expr); } \
  106. }; \
  107. inline op_<op_##id, op_l, self_t, self_t> op(const self_t &, const self_t &) { \
  108. return op_<op_##id, op_l, self_t, self_t>(); \
  109. } \
  110. template <typename T> \
  111. op_<op_##id, op_l, self_t, T> op(const self_t &, const T &) { \
  112. return op_<op_##id, op_l, self_t, T>(); \
  113. } \
  114. template <typename T> \
  115. op_<op_##id, op_r, T, self_t> op(const T &, const self_t &) { \
  116. return op_<op_##id, op_r, T, self_t>(); \
  117. }
  118. #define PYBIND11_INPLACE_OPERATOR(id, op, expr) \
  119. template <typename B, typename L, typename R> \
  120. struct op_impl<op_##id, op_l, B, L, R> { \
  121. static char const *name() { return "__" #id "__"; } \
  122. static auto execute(L &l, const R &r) -> decltype(expr) { return expr; } \
  123. static B execute_cast(L &l, const R &r) { return B(expr); } \
  124. }; \
  125. template <typename T> \
  126. op_<op_##id, op_l, self_t, T> op(const self_t &, const T &) { \
  127. return op_<op_##id, op_l, self_t, T>(); \
  128. }
  129. #define PYBIND11_UNARY_OPERATOR(id, op, expr) \
  130. template <typename B, typename L> \
  131. struct op_impl<op_##id, op_u, B, L, undefined_t> { \
  132. static char const *name() { return "__" #id "__"; } \
  133. static auto execute(const L &l) -> decltype(expr) { return expr; } \
  134. static B execute_cast(const L &l) { return B(expr); } \
  135. }; \
  136. inline op_<op_##id, op_u, self_t, undefined_t> op(const self_t &) { \
  137. return op_<op_##id, op_u, self_t, undefined_t>(); \
  138. }
  139. PYBIND11_BINARY_OPERATOR(sub, rsub, operator-, l - r)
  140. PYBIND11_BINARY_OPERATOR(add, radd, operator+, l + r)
  141. PYBIND11_BINARY_OPERATOR(mul, rmul, operator*, l *r)
  142. PYBIND11_BINARY_OPERATOR(truediv, rtruediv, operator/, l / r)
  143. PYBIND11_BINARY_OPERATOR(mod, rmod, operator%, l % r)
  144. PYBIND11_BINARY_OPERATOR(lshift, rlshift, operator<<, l << r)
  145. PYBIND11_BINARY_OPERATOR(rshift, rrshift, operator>>, l >> r)
  146. PYBIND11_BINARY_OPERATOR(and, rand, operator&, l &r)
  147. PYBIND11_BINARY_OPERATOR(xor, rxor, operator^, l ^ r)
  148. PYBIND11_BINARY_OPERATOR(eq, eq, operator==, l == r)
  149. PYBIND11_BINARY_OPERATOR(ne, ne, operator!=, l != r)
  150. PYBIND11_BINARY_OPERATOR(or, ror, operator|, l | r)
  151. PYBIND11_BINARY_OPERATOR(gt, lt, operator>, l > r)
  152. PYBIND11_BINARY_OPERATOR(ge, le, operator>=, l >= r)
  153. PYBIND11_BINARY_OPERATOR(lt, gt, operator<, l < r)
  154. PYBIND11_BINARY_OPERATOR(le, ge, operator<=, l <= r)
  155. // PYBIND11_BINARY_OPERATOR(pow, rpow, pow, std::pow(l, r))
  156. PYBIND11_INPLACE_OPERATOR(iadd, operator+=, l += r)
  157. PYBIND11_INPLACE_OPERATOR(isub, operator-=, l -= r)
  158. PYBIND11_INPLACE_OPERATOR(imul, operator*=, l *= r)
  159. PYBIND11_INPLACE_OPERATOR(itruediv, operator/=, l /= r)
  160. PYBIND11_INPLACE_OPERATOR(imod, operator%=, l %= r)
  161. PYBIND11_INPLACE_OPERATOR(ilshift, operator<<=, l <<= r)
  162. PYBIND11_INPLACE_OPERATOR(irshift, operator>>=, l >>= r)
  163. PYBIND11_INPLACE_OPERATOR(iand, operator&=, l &= r)
  164. PYBIND11_INPLACE_OPERATOR(ixor, operator^=, l ^= r)
  165. PYBIND11_INPLACE_OPERATOR(ior, operator|=, l |= r)
  166. PYBIND11_UNARY_OPERATOR(neg, operator-, -l)
  167. PYBIND11_UNARY_OPERATOR(pos, operator+, +l)
  168. // WARNING: This usage of `abs` should only be done for existing STL overloads.
  169. // Adding overloads directly in to the `std::` namespace is advised against:
  170. // https://en.cppreference.com/w/cpp/language/extending_std
  171. PYBIND11_UNARY_OPERATOR(abs, abs, std::abs(l))
  172. PYBIND11_UNARY_OPERATOR(hash, hash, std::hash<L>()(l))
  173. PYBIND11_UNARY_OPERATOR(invert, operator~, (~l))
  174. PYBIND11_UNARY_OPERATOR(bool, operator!, !!l)
  175. PYBIND11_UNARY_OPERATOR(int, int_, (int) l)
  176. PYBIND11_UNARY_OPERATOR(float, float_, (double) l)
  177. #undef PYBIND11_BINARY_OPERATOR
  178. #undef PYBIND11_INPLACE_OPERATOR
  179. #undef PYBIND11_UNARY_OPERATOR
  180. PYBIND11_NAMESPACE_END(detail)
  181. using detail::self;
  182. // Add named operators so that they are accessible via `py::`.
  183. using detail::hash;
  184. PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)
  185. #else
  186. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  187. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)