pytorch/torch/fx/experimental
Edward Z. Yang f7365eca90 Add unbacked symints support; item works now (#90624)
The big idea is to add `create_unbacked_symfloat` and `create_unbacked_symint` to ShapeEnv, allowing you to allocate symbolic floats/ints corresponding to data you don't know about at compile time. Then, instead of immediately erroring out when you try to call local_scalar_dense on a FakeTensor, we instead create a fresh symint/symfloat and return that.

There a bunch of odds and ends that need to be handled:

* A number of `numel` calls converted to `sym_numel`
* When we finally return from item(), we need to ensure we actually produce a SymInt/SymFloat when appropriate. The previous binding code assumed that you would have to get a normal Python item. I add a pybind11 binding for Scalar (to PyObject only) and refactor the code to use that. There is some trickiness where you are NOT allowed to go through c10::SymInt if there isn't actually any SymInt involved. See comment.
* One of our unit tests tripped an implicit data dependent access which occurs when you pass a Tensor as an argument to a sizes parameter. This is also converted to support symbolic shapes
* We now support tracking bare SymInt/SymFloat returns in proxy tensor mode (this was already in symbolic-shapes branch)
* Whenever we allocate an unbacked symint, we record the stack trace it was allocated at. These get printed when you attempt data dependent access on the symint (e.g., you try to guard on it)
* Subtlety: unbacked symints are not necessarily > 1. I added a test for this.

These unbacked symints are not very useful right now as you will almost always immediately raise an error later when you try to guard on them. The next logical step is adding an assertion refinement system that lets ShapeEnv learn facts about unbacked symints so it can do a better job eliding guards that are unnecessary.

Signed-off-by: Edward Z. Yang <ezyang@fb.com>
Pull Request resolved: https://github.com/pytorch/pytorch/pull/90624
Approved by: https://github.com/Skylion007, https://github.com/voznesenskym
2022-12-12 13:33:07 +00:00
..
migrate_gradual_types fix torch._C._nn.linear bug (#83682) 2022-08-22 02:16:11 +00:00
unification Fix exception causes in fx, nn and onnx packages (#90134) 2022-12-06 04:34:58 +00:00
__init__.py
accelerator_partitioner.py Fix typos in messages under torch (#89049) 2022-11-17 04:18:14 +00:00
const_fold.py Allow fx.Graph.owning_module to be used as attribute. (#86822) 2022-10-19 00:12:59 +00:00
debug.py
graph_gradual_typechecker.py Fix typos in messages under torch (#89049) 2022-11-17 04:18:14 +00:00
merge_matmul.py
meta_tracer.py Fix dumb make_fx issue (#84011) 2022-08-25 06:52:01 +00:00
normalize.py [fx] Make NormalizeArgs preserve node type (#85637) 2022-09-26 21:30:16 +00:00
optimization.py
partitioner_utils.py
proxy_tensor.py Add unbacked symints support; item works now (#90624) 2022-12-12 13:33:07 +00:00
refinement_types.py
rewriter.py
schema_type_annotation.py Revert "[quant][fx] Support keyword arguments for functional linear (#79095)" 2022-07-09 21:45:01 +00:00
symbolic_shapes.py Add unbacked symints support; item works now (#90624) 2022-12-12 13:33:07 +00:00
unify_refinements.py