pytorch
19ede75e - [JIT] Enable ModuleDict non-literal indexing (#45716)

Commit
4 years ago
[JIT] Enable ModuleDict non-literal indexing (#45716) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/45716 **Summary** This commit enables indexing into `ModuleDict` using a non-literal index if the `ModuleDict` is annotated with `Dict[str, X]`, where `X` is a module interface type. These annotations must be expressed using a class attribute named `__annotations__`, which is a `Dict[str, Type]` where the keys are the names of module attributes and the values are their types. The approach taken by this commit is that these annotations are stored as "hints" along with the corresponding module attributes in the `ConcreteSubmoduleTypeBuilder` instance for each module (which might be a `ModuleDict`). These hints are passed into the `ModuleValue` that is created for desugaring operations on submodules so that indexing into a `ModuleDict` can be emitted as a getitem op into a dict emitted into the graph that represents the `ModuleDict`. **Test Plan** This commit adds unit tests to `TestModuleContainers` to test this feature (`test_typed_module_dict`). Differential Revision: D24070606 Test Plan: Imported from OSS Reviewed By: ansley Pulled By: SplitInfinity fbshipit-source-id: 6019a7242d53d68fbfc1aa5a49df6cfc0507b992
Author
Meghan Lele
Parents
Loading