1#pragma once
2
3#include <c10/core/SafePyObject.h>
4#include <c10/macros/Macros.h>
5#include <unordered_map>
6
7namespace at {
8namespace impl {
9
10struct TORCH_API ThreadLocalPythonObjects {
11 static void set(const std::string& key, std::shared_ptr<SafePyObject> value);
12 static const std::shared_ptr<SafePyObject>& get(const std::string& key);
13 static bool contains(const std::string& key);
14
15 static const ThreadLocalPythonObjects& get_state();
16 static void set_state(ThreadLocalPythonObjects state);
17
18 private:
19 std::unordered_map<std::string, std::shared_ptr<c10::SafePyObject>> obj_dict_;
20};
21
22} // namespace impl
23} // namespace at
24