pytorch/torch
Jiashen Cao ca7ce2fca1 [ts-migration][1/N]: Add prim::Loop for constant number of iterations and condition (#131418)
#### Description
This PR adds prim::Loop support for the simplest case where the number of iteration is constant and the loop termination condition is also a constant.

[PR by stages](https://docs.google.com/document/d/1q6OprW3HBHbYPwEyE_DikBn-uzmhnN284Cmen_CnlhI/edit?usp=sharing)

#### Test Plan
Add reprod example.
* `pytest test/export/test_converter.py -s -k test_ts2ep_with_loop`
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131418
Approved by: https://github.com/angelayi
2024-08-06 16:51:08 +00:00
..
_awaits
_C [c10d] Add a new API for adding ephemeral timeout for one local rank and the timeout will reset when the first collective finishes (#130905) 2024-08-06 03:47:58 +00:00
_C_flatbuffer
_custom_op Tighten torch.library.infer_schema input types (#130705) 2024-07-29 16:01:19 +00:00
_decomp Add None return type to init (#132335) 2024-08-01 15:26:45 +00:00
_dispatch [BE][Easy][15/19] enforce style for empty lines in import segments in torch/_d*/ (#129767) 2024-07-31 21:18:11 +00:00
_dynamo Add cudagraph static inputs logging (#132726) 2024-08-06 12:01:20 +00:00
_export [ts-migration][1/N]: Add prim::Loop for constant number of iterations and condition (#131418) 2024-08-06 16:51:08 +00:00
_functorch Add cudagraph static inputs logging (#132726) 2024-08-06 12:01:20 +00:00
_higher_order_ops [hop] makes compiled hops not share code objects (#132427) 2024-08-05 22:59:05 +00:00
_inductor Revert "[inductor] export kernel for gemm template. (#132580)" 2024-08-06 13:15:15 +00:00
_lazy Add None return type to init (#132335) 2024-08-01 15:26:45 +00:00
_library [Doc] fix some typos (found by codespell and typos) (#132544) 2024-08-05 17:21:56 +00:00
_logging Add cudagraph static inputs logging (#132726) 2024-08-06 12:01:20 +00:00
_numpy [test/torch_np] Fix usages of deprecated NumPy 2.0 APIs in numpy_tests (#131909) 2024-08-05 16:21:08 +00:00
_prims [BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765) 2024-07-31 10:42:50 +00:00
_prims_common [BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765) 2024-07-31 10:42:50 +00:00
_refs Don't attempt to compute hints for unbacked expressions (#132060) 2024-08-02 16:39:14 +00:00
_strobelight [BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765) 2024-07-31 10:42:50 +00:00
_subclasses Revert "Conversions between strided and jagged layouts for Nested Tensors (#115749)" 2024-08-06 01:55:41 +00:00
_vendor
amp Revert "[MPS] Add support for autocast in MPS (#99272)" 2024-08-05 19:59:04 +00:00
ao [pt2e][quant] Ensure BN node is erased after convert (#131651) 2024-08-06 16:37:39 +00:00
autograd [BE]: Simplify some list comps to generators C419 (#132578) 2024-08-04 17:46:26 +00:00
backends [BE][Easy][17/19] enforce style for empty lines in import segments in torch/[a-c]*/ and torch/[e-n]*/ (#129769) 2024-08-04 10:24:09 +00:00
compiler
contrib
cpu [BE][Easy][17/19] enforce style for empty lines in import segments in torch/[a-c]*/ and torch/[e-n]*/ (#129769) 2024-08-04 10:24:09 +00:00
csrc [c10d] Add a new API for adding ephemeral timeout for one local rank and the timeout will reset when the first collective finishes (#130905) 2024-08-06 03:47:58 +00:00
cuda Allow torch.cuda.memory.mem_get_info to take a device str argument with an unspecified device index. (#132616) 2024-08-06 13:19:46 +00:00
distributed [c10d] Add a new API for adding ephemeral timeout for one local rank and the timeout will reset when the first collective finishes (#130905) 2024-08-06 03:47:58 +00:00
distributions [BE][Easy][18/19] enforce style for empty lines in import segments in torch/d*/ (#129770) 2024-08-01 04:22:50 +00:00
export [export][reland] Convert autocast to HOO (#132677) 2024-08-05 22:34:52 +00:00
fft
func
futures
fx Converted __all__ literal tuple to literal list. (#132404) 2024-08-06 15:12:32 +00:00
jit [BE][Easy][17/19] enforce style for empty lines in import segments in torch/[a-c]*/ and torch/[e-n]*/ (#129769) 2024-08-04 10:24:09 +00:00
legacy
lib
linalg
masked [BE][Easy][17/19] enforce style for empty lines in import segments in torch/[a-c]*/ and torch/[e-n]*/ (#129769) 2024-08-04 10:24:09 +00:00
monitor [pytorch][counters] Pybind for WaitCounter (#132357) 2024-08-02 16:08:10 +00:00
mps [BE][Easy][17/19] enforce style for empty lines in import segments in torch/[a-c]*/ and torch/[e-n]*/ (#129769) 2024-08-04 10:24:09 +00:00
mtia [BE][Easy][17/19] enforce style for empty lines in import segments in torch/[a-c]*/ and torch/[e-n]*/ (#129769) 2024-08-04 10:24:09 +00:00
multiprocessing [BE][Easy][17/19] enforce style for empty lines in import segments in torch/[a-c]*/ and torch/[e-n]*/ (#129769) 2024-08-04 10:24:09 +00:00
nested Revert "Conversions between strided and jagged layouts for Nested Tensors (#115749)" 2024-08-06 01:55:41 +00:00
nn Format tensor (#127992) 2024-08-06 07:10:16 +00:00
onnx [ONNX] Rename _internal/exporter to _exporter_legacy (#132429) 2024-08-03 04:23:05 +00:00
optim Deprecate torch._utils.is_compiling() and torch._dynamo.external_utils.is_compiling() (#127690) 2024-08-03 09:43:38 +00:00
package [BE][Easy][19/19] enforce style for empty lines in import segments in torch/[o-z]*/ (#129771) 2024-08-01 17:07:14 +00:00
profiler [BE][Easy][19/19] enforce style for empty lines in import segments in torch/[o-z]*/ (#129771) 2024-08-01 17:07:14 +00:00
quantization [BE][Easy][19/19] enforce style for empty lines in import segments in torch/[o-z]*/ (#129771) 2024-08-01 17:07:14 +00:00
signal Revert "[BE] typing for decorators - signal/windows/windows (#131582)" 2024-07-28 03:29:31 +00:00
sparse Add tests to bsr_dense_addmm_meta. Tune bsr_dense_addmm kernel for ViT shapes. (#132646) 2024-08-05 20:22:33 +00:00
special
testing [TP] verify numeric parity on Transfromers for multiple iterations (#132543) 2024-08-04 06:43:27 +00:00
utils Refactor thunkify to return proper thunk abstraction (#132407) 2024-08-06 02:35:45 +00:00
xpu [BE][Easy][19/19] enforce style for empty lines in import segments in torch/[o-z]*/ (#129771) 2024-08-01 17:07:14 +00:00
__config__.py
__future__.py
__init__.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_appdirs.py [BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765) 2024-07-31 10:42:50 +00:00
_classes.py Add None return type to init (#132335) 2024-08-01 15:26:45 +00:00
_compile.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_custom_ops.py [BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765) 2024-07-31 10:42:50 +00:00
_deploy.py
_guards.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_jit_internal.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_linalg_utils.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_lobpcg.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_lowrank.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_meta_registrations.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_namedtensor_internals.py [BE][Easy][14/19] enforce style for empty lines in import segments in torch/_[a-c]*/ and torch/_[e-h]*/ and torch/_[j-z]*/ (#129765) 2024-07-31 10:42:50 +00:00
_ops.py [easy] fix f-string messages in torch/_ops.py (#132531) 2024-08-05 18:58:33 +00:00
_python_dispatcher.py Add None return type to init (#132335) 2024-08-01 15:26:45 +00:00
_size_docs.py
_sources.py
_storage_docs.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_streambase.py
_tensor_docs.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_tensor_str.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_tensor.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_torch_docs.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_utils_internal.py
_utils.py Deprecate torch._utils.is_compiling() and torch._dynamo.external_utils.is_compiling() (#127690) 2024-08-03 09:43:38 +00:00
_VF.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
_vmap_internals.py
_weights_only_unpickler.py
abi-check.cpp
CMakeLists.txt Revert "Add wrappers for synchronous GPUDirect Storage APIs (#130633)" 2024-07-26 18:08:20 +00:00
custom_class_detail.h
custom_class.h
extension.h
functional.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
hub.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
library.h
library.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
overrides.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
py.typed
quasirandom.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
random.py
README.txt
return_types.py
script.h
serialization.py [BE] Format uncategorized Python files with ruff format (#132576) 2024-08-04 17:13:31 +00:00
storage.py Fix public API tests (#131386) 2024-07-30 18:42:54 +00:00
torch_version.py
types.py
version.py.tpl

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.