1 | #pragma once |
2 | #include <c10/core/DispatchKey.h> |
3 | #include <c10/core/impl/LocalDispatchKeySet.h> |
4 | #include <torch/csrc/python_headers.h> |
5 | |
6 | namespace torch { |
7 | // Sometimes we don't want infinite recursion for subclasses, |
8 | // Or a way to achieve the old behaviour. |
9 | |
10 | // This is an internal utility, not exposed to users. |
11 | bool torch_function_enabled(); |
12 | PyObject* disabled_torch_function_impl(); |
13 | PyObject* disabled_torch_dispatch_impl(); |
14 | void set_disabled_torch_function_impl(PyObject* value); |
15 | void set_disabled_torch_dispatch_impl(PyObject* value); |
16 | // Set ignore_mode to true if you're trying to collect overloaded arguments; |
17 | // using mode here will improperly cause you to add ALL objects to the |
18 | // overloaded list even if they don't actually have __torch_function__ |
19 | bool check_has_torch_function(PyObject* obj, bool ignore_mode = false); |
20 | |
21 | struct DisableTorchDispatch { |
22 | DisableTorchDispatch() |
23 | : guard_(c10::DispatchKey::Python), |
24 | guard_tls_snapshot_(c10::DispatchKey::PythonTLSSnapshot) {} |
25 | c10::impl::ExcludeDispatchKeyGuard guard_; |
26 | c10::impl::ExcludeDispatchKeyGuard guard_tls_snapshot_; |
27 | }; |
28 | |
29 | } // namespace torch |
30 | |
31 | PyObject* THPModule_isEnabledTorchFunction(PyObject* self, PyObject* unused); |
32 | PyObject* THPModule_DisableTorchFunctionType(); |
33 | PyObject* THPModule_DisableTorchFunctionSubclassType(); |
34 | PyObject* THPModule_disable_torch_function(PyObject* self, PyObject* args); |
35 | PyObject* THPModule_disable_torch_dispatch(PyObject* self, PyObject* args); |
36 | PyObject* THPModule_has_torch_function(PyObject*, PyObject* arg); |
37 | PyObject* THPModule_has_torch_function_unary(PyObject*, PyObject* obj); |
38 | PyObject* THPModule_has_torch_function_variadic( |
39 | PyObject*, |
40 | PyObject* const* args, |
41 | Py_ssize_t nargs); |
42 | |