pytorch/torch/_dynamo/variables
Shangdi Yu 348f208504 Autocast re-tracibility (#138082)
Summary:
Support autocast re-tracing by giving it the same treatment as set_grad.

In re-tracing, when dynamo encounters an autocast HOP, we want it to trace through `with torch.autocast()` again, and replace the HOP with the traced subgraph.

Test Plan:
```
buck2 run 'fbcode//mode/dev-nosan' fbcode//caffe2/test:test_export  -- -r  test_export_with_autocast
```

Differential Revision: D63856081

Pull Request resolved: https://github.com/pytorch/pytorch/pull/138082
Approved by: https://github.com/ydwu4
2024-10-17 16:09:11 +00:00
..
__init__.py Add host-side Triton TMA support to Dynamo (#137677) 2024-10-16 02:18:48 +00:00
base.py Revert "Fix constant propagation in builtins and UserClasses (#131354)" 2024-09-25 23:01:03 +00:00
builder.py Add host-side Triton TMA support to Dynamo (#137677) 2024-10-16 02:18:48 +00:00
builtin.py [Dynamo] Ensure torch function modes are dispatched on builtin ops (#137117) 2024-10-09 02:29:40 +00:00
constant.py [dynamo] Fix support for classmethod(property(...)) (#134968) 2024-09-18 04:47:51 +00:00
ctx_manager.py [Dynamo] Trace enter/exit of TorchFunctionModes (#135422) (#137114) 2024-10-09 02:29:40 +00:00
dicts.py Remove vt argument in raise_observed_exception (#136037) 2024-09-24 02:36:57 +00:00
distributed.py Revert "[compiled autograd] Compiled autograd configs in TLS (#137821)" 2024-10-16 16:38:29 +00:00
functions.py Add host-side Triton TMA support to Dynamo (#137677) 2024-10-16 02:18:48 +00:00
higher_order_ops.py Autocast re-tracibility (#138082) 2024-10-17 16:09:11 +00:00
iter.py Remove vt argument in raise_observed_exception (#136037) 2024-09-24 02:36:57 +00:00
lazy.py [Dynamo] Handle torch function subclass/mode dispatch on generic tensor methods (#137119) 2024-10-09 02:29:40 +00:00
lists.py [dynamo] fix NamedTupleVariable for PyStructSequence (torch.return_types.*) support (#137776) 2024-10-13 06:46:41 +00:00
misc.py Revert "[compiled autograd] Compiled autograd configs in TLS (#137821)" 2024-10-16 16:38:29 +00:00
nn_module.py [alt] fix unroll in successive unflatten (#137646) 2024-10-12 15:53:52 +00:00
optimizer.py Add warning for non static grads in optimizer variable (#137554) 2024-10-10 01:23:21 +00:00
script_object.py Add basic mypy annotations to dynamo (#132415) 2024-08-04 18:43:36 +00:00
sdpa.py Grouped Query Attention (#132689) 2024-08-07 05:35:36 +00:00
tensor.py Revert "[compiled autograd] Compiled autograd configs in TLS (#137821)" 2024-10-16 16:38:29 +00:00
torch_function.py [Dynamo] Handle torch function subclass/mode dispatch on generic tensor methods (#137119) 2024-10-09 02:29:40 +00:00
torch.py [Dynamo] Handle torch function subclass/mode dispatch on generic tensor methods (#137119) 2024-10-09 02:29:40 +00:00
user_defined.py [dynamo] fix NamedTupleVariable for PyStructSequence (torch.return_types.*) support (#137776) 2024-10-13 06:46:41 +00:00