pytorch/torch
Jordan Fix 987f146185 [fx] Improve support for tuple subclasses such as NamedTuple (#73198)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/73198

Previously, if an arg to an FX node is a subclass of tuple then it gets sanitized essentially back to that base class. An example here is when setting an arg to be a TensorMetadata object, which is a NamedTuple, it will be set as a tuple instead.

- Change `map_aggregate` to repack the tuple to `type(a)` when it's not directly a tuple (try/except for best attempt)
- During codegen, call `add_global` for `type(a)` if it's not directly a tuple.
- Add an option for an arg to provide a `_custom_fx_repr_fn` for use inside stringifying via `_format_arg`

Test Plan: Added unit test coverage, where we inline the named tuple into arg/kwarg.

Reviewed By: jamesr66a

Differential Revision: D34381888

fbshipit-source-id: bd672a8542e2bba5aa604b448bec920efc256440
(cherry picked from commit 68f99c12dd)
2022-02-23 11:31:10 +00:00
..
_C Integrate full ONNX check into ONNX export API (#71125) 2022-02-18 18:40:09 +00:00
_masked Generate static docstrings for torch._masked functions. (#72865) 2022-02-17 02:44:16 +00:00
ao [quant] Add ConvTranspose reference module - Reland #73031 (#73094) 2022-02-23 02:31:42 +00:00
autograd Add autograd tests for addmm, addmv, mm, mv and CSR matrix input (#71949) 2022-02-11 23:14:02 +00:00
backends Fix type annotation for torch.backends.cudnn.allow_tf32 (#72757) 2022-02-14 17:26:37 +00:00
contrib
cpu
csrc [AutoAccept][Codemod][FBSourceClangFormatLinter] Daily arc lint --take CLANGFORMAT 2022-02-23 10:29:48 +00:00
cuda Avoid collections deprecation warning (#72239) 2022-02-23 02:31:42 +00:00
distributed [ZeRO] (Reland) Add ctor support for multiple param groups (#72932) 2022-02-22 16:29:55 +00:00
distributions Adding details to kl.py (#72845) 2022-02-19 06:33:08 +00:00
fft
for_onnx
futures #71946 Remove Python 3.6 references (#72211) 2022-02-08 03:46:20 +00:00
fx [fx] Improve support for tuple subclasses such as NamedTuple (#73198) 2022-02-23 11:31:10 +00:00
jit Make fusion strategy api public (#72639) 2022-02-16 03:45:15 +00:00
legacy
lib
linalg Extrapolated on equiv between linalg @ and solve (#71769) 2022-02-22 12:29:32 +00:00
monitor torch/monitor: TensorboardEventHandler (#71658) 2022-01-27 08:33:55 +00:00
multiprocessing Rename Typed/UntypedStorage to _Typed/_UntypedStorage (#72540) 2022-02-15 23:53:01 +00:00
nn [quant] Add ConvTranspose reference module - Reland #73031 (#73094) 2022-02-23 02:31:42 +00:00
onnx [ONNX] Add Squeeze/Unsqueeze dynamic dimensions support when opset >= 13 (#71158) 2022-02-23 06:41:15 +00:00
optim [PyTorch/d2go] fix optim _multi_tensor (#73215) 2022-02-23 10:29:48 +00:00
package Rename Typed/UntypedStorage to _Typed/_UntypedStorage (#72540) 2022-02-15 23:53:01 +00:00
profiler Add low level torch.profiler.kineto_profile base class (#63302) 2021-12-14 14:47:43 -08:00
quantization [quant][fx][graphmode][be] Change the type for output of convert to be torch.nn.Module (#69959) 2021-12-29 20:33:32 -08:00
sparse Add missing entry for sampled_addmm in sparse.rst (#72312) 2022-02-08 00:07:20 +00:00
special
testing Change how cuda available memory is calculated in largeTensorTest decorator (#72207) 2022-02-23 02:31:42 +00:00
utils Documenting cuda 11.5 windows issue (#73013) 2022-02-19 02:34:59 +00:00
__config__.py
__future__.py
__init__.py Rename Typed/UntypedStorage to _Typed/_UntypedStorage (#72540) 2022-02-15 23:53:01 +00:00
_appdirs.py
_classes.py
_deploy.py Rename Typed/UntypedStorage to _Typed/_UntypedStorage (#72540) 2022-02-15 23:53:01 +00:00
_jit_internal.py (2/2) Make TorchScript Preserve Fully Qualified Class Name for Python Exceptions: frontend change (#72899) 2022-02-16 03:45:15 +00:00
_linalg_utils.py
_lobpcg.py Fix trivial typo at the doc of torch.lobpcg (#71464) 2022-01-20 00:07:39 +00:00
_lowrank.py
_namedtensor_internals.py
_ops.py Add torch.ops per overload API (#72206) 2022-02-11 17:19:06 +00:00
_python_dispatcher.py Revert D34034848: free up dispatch key space (in C++) 2022-02-14 23:29:00 +00:00
_six.py Update min python version to 3.7 in setup.py and mypy configs (#71494) 2022-01-20 00:03:57 +00:00
_sources.py
_storage_docs.py
_tensor_docs.py scatter_reduce documentation (#73125) 2022-02-22 19:33:46 +00:00
_tensor_str.py fixed compilations on xla tensor print (#71147) 2022-01-27 02:28:19 +00:00
_tensor.py Fix tensor.__deepcopy__ for lazy device (#73197) 2022-02-23 02:31:42 +00:00
_torch_docs.py scatter_reduce documentation (#73125) 2022-02-22 19:33:46 +00:00
_utils_internal.py
_utils.py Rename Typed/UntypedStorage to _Typed/_UntypedStorage (#72540) 2022-02-15 23:53:01 +00:00
_VF.py
_vmap_internals.py More aggressively market functorch.vmap when torch.vmap gets called (#67347) 2021-11-12 16:10:16 -08:00
abi-check.cpp
autocast_mode.py
CMakeLists.txt Allow caffe2-specific graph transformations for OperatorExportTypes.ONNX_ATEN_FALLBACK when BUILD_CAFFE2 is ON (#67460) (#68490) 2022-02-10 03:26:48 +00:00
custom_class_detail.h
custom_class.h [jit] Decouple ivalue.h from jit_type.h (#70119) 2022-01-07 18:34:17 -08:00
deploy.h
extension.h
functional.py Update stft tests to support latest librosa (#72833) 2022-02-23 02:31:42 +00:00
hub.py making import_module private and deprecating public method (#67990) 2021-11-09 07:27:57 -08:00
library.h Codegen: python_torch_functions only include relevant operators (#68693) 2022-01-21 15:37:06 +00:00
overrides.py scatter_reduce documentation (#73125) 2022-02-22 19:33:46 +00:00
py.typed
quasirandom.py
random.py
README.txt
return_types.py expose return_types in Python (#66614) 2021-12-06 09:05:29 -08:00
script.h
serialization.py Rename Typed/UntypedStorage to _Typed/_UntypedStorage (#72540) 2022-02-15 23:53:01 +00:00
storage.py Rename Typed/UntypedStorage to _Typed/_UntypedStorage (#72540) 2022-02-15 23:53:01 +00:00
torch_version.py Import packaging.version in torch_version, if available (#71902) 2022-02-22 21:30:14 +00:00
types.py

Note [TH abstraction violation]
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

TH/THC provide some hpp headers, which are proper C++ headers rather than
C headers.  These headers serve double duty as *internal implementation
detail* headers, whose contents should largely not be used by external
clients.

Ideally, we would not install these headers at all; instead, you should
use public functions (in headers like `THTensor.h`, NOT `THTensor.hpp`)
to manipulate these structs.  However, there are a few places
in torch/csrc where we violate this abstraction.  They are marked with
a pointer to this note.  Each of those sites will have to be refactored
when we refactor the guts of THTensor and related structures.