(
self,
tx,
args: "List[VariableTracker]",
kwargs: "Dict[str, VariableTracker]",
)
| 236 | return variables.GetAttrVariable(self, name, source=source) |
| 237 | |
| 238 | def call_function( |
| 239 | self, |
| 240 | tx, |
| 241 | args: "List[VariableTracker]", |
| 242 | kwargs: "Dict[str, VariableTracker]", |
| 243 | ) -> "VariableTracker": |
| 244 | mod = tx.output.get_submodule(self.module_key) |
| 245 | |
| 246 | with record_nn_module_stack(self.module_key, self.source, tx, mod): |
| 247 | is_lazy = is_lazy_module(mod) |
| 248 | if ( |
| 249 | isinstance(mod, torch.nn.Sequential) |
| 250 | and mod.__class__.forward is torch.nn.Sequential.forward |
| 251 | ): |
| 252 | if nnmodule_has_hooks(mod): |
| 253 | # We do not want to unroll sequential if it has hooks, since evaporating it |
| 254 | # will cause hooks to not fire! |
| 255 | # This terminates and restart the tracing process |
| 256 | self.convert_to_unspecialized(tx) |
| 257 | |
| 258 | # Unroll sequential |
| 259 | assert ( |
| 260 | not is_lazy |
| 261 | ), "Expected lazy sequential isn't a valid combination?" |
| 262 | assert not kwargs |
| 263 | (arg,) = args |
| 264 | # TODO: Use named_children when it supports remove_duplicate=False. |
| 265 | for child_name, submod in mod._modules.items(): |
| 266 | tx.call_function( |
| 267 | tx.output.register_attr_or_module( |
| 268 | submod, |
| 269 | self.module_key, |
| 270 | child_name, |
| 271 | source=NNModuleSource(AttrSource(self.source, child_name)), |
| 272 | ), |
| 273 | [arg], |
| 274 | {}, |
| 275 | ) |
| 276 | arg = tx.pop() |
| 277 | return arg |
| 278 | |
| 279 | if is_lazy: |
| 280 | # The module type will change after it is called |
| 281 | if mod.cls_to_become is not None: |
| 282 | self.module_type = mod.cls_to_become |
| 283 | |
| 284 | # The pre-hook runs to initialize the module shapes, then deletes itself. After this, |
| 285 | # the module is more or less not lazy and can be treated as a normal module regardless of |
| 286 | # is_allowed or other variations. |
| 287 | initialize_lazy_module(tx, mod, args, kwargs) |
| 288 | |
| 289 | # If we are tracing the higher order op, we want Dynamo to step |
| 290 | # inside the module call so that Dynamo can see the underlying |
| 291 | # parameters and buffers and raise them as inputs to the graph. |
| 292 | if tx.output.is_root_tracer() and is_allowed(mod.__class__): |
| 293 | if nnmodule_has_hooks( |
| 294 | mod, check_forward_hooks=True, check_backward_hooks=True |
| 295 | ): |
no test coverage detected