(self, a: Argument)
| 207 | raise ExportPassBaseError("ExportTracer doesn't support trace().") |
| 208 | |
| 209 | def create_arg(self, a: Argument) -> torch.fx.Node: |
| 210 | if isinstance(a, torch.nn.Module): |
| 211 | if a not in self.submodules: |
| 212 | name_submodule = f"submodule_{len(self.submodules)}" |
| 213 | self.root.add_module(name_submodule, a) |
| 214 | self.submodules[a] = name_submodule |
| 215 | elif isinstance(a, FakeTensor): |
| 216 | if not hasattr(a, "constant") or a.constant is None: |
| 217 | raise ExportPassBaseError(f"Cannot add {a} to graph.") |
| 218 | a = a.constant |
| 219 | elif isinstance(a, torch.SymInt): |
| 220 | if a.node.constant is not None: |
| 221 | return a.node.constant |
| 222 | else: |
| 223 | return a |
| 224 | node = super().create_arg(a) |
| 225 | if ( |
| 226 | isinstance(a, torch.Tensor) |
| 227 | and isinstance(node, torch.fx.Node) |
| 228 | and node.op == "get_attr" |
| 229 | ): |
| 230 | self.set_metadata(node, a) |
| 231 | self.callback.on_attr(ProxyValue(a, node)) |
| 232 | return node |
| 233 | |
| 234 | def set_metadata( # noqa: C901 |
| 235 | self, |
nothing calls this directly
no test coverage detected