1 | #pragma once |
---|---|
2 | |
3 | #include <torch/csrc/jit/frontend/name_mangler.h> |
4 | #include <torch/csrc/jit/ir/type_hashing.h> |
5 | |
6 | namespace torch { |
7 | namespace jit { |
8 | |
9 | /** |
10 | * class TypeNameUniquer |
11 | * |
12 | * Generates a unique name for every type `t` passed in. Types that compare |
13 | * equal with EqualType will receive the same unique name. |
14 | * |
15 | * This is used during Module::save(), to resolve type name collisions during |
16 | * serialization. |
17 | */ |
18 | class TORCH_API TypeNameUniquer { |
19 | public: |
20 | c10::QualifiedName getUniqueName(c10::ConstNamedTypePtr t); |
21 | |
22 | private: |
23 | NameMangler mangler_; |
24 | std::unordered_set<c10::QualifiedName> used_names_; |
25 | std::unordered_map< |
26 | c10::ConstNamedTypePtr, |
27 | c10::QualifiedName, |
28 | HashType, |
29 | EqualType> |
30 | name_map_; |
31 | }; |
32 | } // namespace jit |
33 | } // namespace torch |
34 |