xref: /aosp_15_r20/external/pytorch/torch/csrc/Generator.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <ATen/core/Generator.h>
4 #include <torch/csrc/Export.h>
5 #include <torch/csrc/python_headers.h>
6 
7 // NOLINTNEXTLINE(cppcoreguidelines-pro-type-member-init)
8 struct THPGenerator {
9   PyObject_HEAD at::Generator cdata;
10 };
11 
12 // Creates a new Python object wrapping the default at::Generator. The reference
13 // is borrowed. The caller should ensure that the at::Generator object lifetime
14 // last at least as long as the Python wrapper.
15 TORCH_PYTHON_API PyObject* THPGenerator_initDefaultGenerator(
16     at::Generator cdata);
17 
18 #define THPGenerator_Check(obj) PyObject_IsInstance(obj, THPGeneratorClass)
19 
20 TORCH_PYTHON_API extern PyObject* THPGeneratorClass;
21 
22 bool THPGenerator_init(PyObject* module);
23 
24 TORCH_PYTHON_API PyObject* THPGenerator_Wrap(at::Generator gen);
25 
26 TORCH_PYTHON_API at::Generator THPGenerator_Unwrap(PyObject* state);
27 
28 // Creates a new Python object for a Generator. The Generator must not already
29 // have a PyObject* associated with it.
30 PyObject* THPGenerator_NewWithVar(PyTypeObject* type, at::Generator gen);
31