(self, *args, **kwargs)
| 350 | return self.trace_without_host(*args, **kwargs) |
| 351 | |
| 352 | def trace_normal(self, *args, **kwargs): |
| 353 | global active_trace |
| 354 | symbolic_shape = None |
| 355 | outputs = None |
| 356 | try: |
| 357 | active_trace = self |
| 358 | self._trace.enter() |
| 359 | if self._capture_as_const: |
| 360 | self._process_inputs(*args, **kwargs) |
| 361 | symbolic_shape = set_symbolic_shape(self._symbolic_shape) |
| 362 | outputs = self.__wrapped__(*args, **kwargs) |
| 363 | finally: |
| 364 | handling_exc = sys.exc_info() != (None,) * 3 |
| 365 | active_trace = None |
| 366 | if symbolic_shape is not None: |
| 367 | symbolic_shape = set_symbolic_shape(symbolic_shape) |
| 368 | assert symbolic_shape == self._symbolic_shape |
| 369 | if self._capture_as_const and (outputs is not None): |
| 370 | self._process_outputs(outputs) |
| 371 | try: |
| 372 | # may raise TraceError |
| 373 | self._trace.exit() |
| 374 | except TraceError: |
| 375 | if not handling_exc: |
| 376 | raise |
| 377 | return outputs |
| 378 | |
| 379 | def trace_without_host(self, *args, **kwargs): |
| 380 | from ..traced_module.pytree import tree_flatten, SUPPORTED_LEAF_CLS |
no test coverage detected