File: disable_torch_function.h

package info (click to toggle)
pytorch-cuda 2.6.0%2Bdfsg-7
  • links: PTS, VCS
  • area: contrib
  • in suites: forky, sid, trixie
  • size: 161,620 kB
  • sloc: python: 1,278,832; cpp: 900,322; ansic: 82,710; asm: 7,754; java: 3,363; sh: 2,811; javascript: 2,443; makefile: 597; ruby: 195; xml: 84; objc: 68
file content (45 lines) | stat: -rw-r--r-- 1,868 bytes parent folder | download | duplicates (3)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
#pragma once
#include <c10/core/DispatchKey.h>
#include <c10/core/impl/LocalDispatchKeySet.h>
#include <torch/csrc/python_headers.h>

namespace torch {
// Sometimes we don't want infinite recursion for subclasses,
// Or a way to achieve the old behaviour.

// This is an internal utility, not exposed to users.
bool torch_function_enabled();
PyObject* disabled_torch_function_impl();
PyObject* disabled_torch_dispatch_impl();
void set_disabled_torch_function_impl(PyObject* value);
void set_disabled_torch_dispatch_impl(PyObject* value);
// Set ignore_mode to true if you're trying to collect overloaded arguments;
// using mode here will improperly cause you to add ALL objects to the
// overloaded list even if they don't actually have __torch_function__
bool check_has_torch_function(PyObject* obj, bool ignore_mode = false);

struct DisableTorchDispatch {
  DisableTorchDispatch()
      : guard_(c10::DispatchKeySet(
            {c10::DispatchKey::Python, c10::DispatchKey::PreDispatch})),
        guard_tls_snapshot_(c10::DispatchKey::PythonTLSSnapshot) {}
  c10::impl::ExcludeDispatchKeyGuard guard_;
  c10::impl::ExcludeDispatchKeyGuard guard_tls_snapshot_;
};

} // namespace torch

PyObject* THPModule_isEnabledTorchFunction(PyObject* self, PyObject* unused);
PyObject* THPModule_isAllDisabledTorchFunction(
    PyObject* self,
    PyObject* unused);
PyObject* THPModule_DisableTorchFunctionType();
PyObject* THPModule_DisableTorchFunctionSubclassType();
PyObject* THPModule_disable_torch_function(PyObject* self, PyObject* args);
PyObject* THPModule_disable_torch_dispatch(PyObject* self, PyObject* args);
PyObject* THPModule_has_torch_function(PyObject*, PyObject* arg);
PyObject* THPModule_has_torch_function_unary(PyObject*, PyObject* obj);
PyObject* THPModule_has_torch_function_variadic(
    PyObject*,
    PyObject* const* args,
    Py_ssize_t nargs);