1#pragma once
2#include <c10/core/DispatchKey.h>
3#include <c10/core/impl/LocalDispatchKeySet.h>
4#include <torch/csrc/python_headers.h>
5
6namespace 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.
11bool torch_function_enabled();
12PyObject* disabled_torch_function_impl();
13PyObject* disabled_torch_dispatch_impl();
14void set_disabled_torch_function_impl(PyObject* value);
15void 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__
19bool check_has_torch_function(PyObject* obj, bool ignore_mode = false);
20
21struct 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
31PyObject* THPModule_isEnabledTorchFunction(PyObject* self, PyObject* unused);
32PyObject* THPModule_DisableTorchFunctionType();
33PyObject* THPModule_DisableTorchFunctionSubclassType();
34PyObject* THPModule_disable_torch_function(PyObject* self, PyObject* args);
35PyObject* THPModule_disable_torch_dispatch(PyObject* self, PyObject* args);
36PyObject* THPModule_has_torch_function(PyObject*, PyObject* arg);
37PyObject* THPModule_has_torch_function_unary(PyObject*, PyObject* obj);
38PyObject* THPModule_has_torch_function_variadic(
39 PyObject*,
40 PyObject* const* args,
41 Py_ssize_t nargs);
42