blob: dc70e7bc5bcef0159c8f6d5291c3a335e9b72f35 [file] [log] [blame]
#pragma once
#include <ATen/core/ivalue.h>
namespace torch {
namespace jit {
namespace script {
// a stable location that can hold an IValue.
// Currently this is internally implemented as a pointer, but when
// modules become first-class this will be a pair of <module_ivalue, slot_number>
struct Slot {
friend struct NamedIValue;
Slot()
: slot_(nullptr) {}
Slot(at::IValue* slot)
: slot_(slot) {}
at::IValue& operator*() const {
return *slot_;
}
at::IValue* operator->() const {
return slot_;
}
bool operator==(const Slot& rhs) const {
return slot_ == rhs.slot_;
}
private:
at::IValue* slot_;
friend struct std::hash<Slot>;
};
}}}
// slots are hashable, because they are often used as keys in maps
// for remapping uses of a slot from one model to another
namespace std {
template <>
struct hash<torch::jit::script::Slot> {
size_t operator()(const torch::jit::script::Slot& s) const noexcept {
return std::hash<at::IValue*>{}(s.slot_);
}
};
} // namespace std