gopikrsmscs/torch-issues huggingface dataset:
Using `torch.onnx.export` from file named `onnx.py` results in cryptic error message
Labels: module: onnx, triaged

### 🐛 Describe the bug

Right now, there's a massive footgun when using `torch.onnx.export()` from a file named `onnx.py`.

```python3
# File contents of onnx.py
import torch
import torch.nn as nn
torch.onnx.export(nn.Linear(10, 10), args=torch.ones(10), f='output.onnx')
```
This results in a few dozen lines of rather cryptic error messages:
```
Traceback (most recent call last):
  File "/path/to/project/onnx.py", line 3, in <module>
    torch.onnx.export(nn.Linear(10, 10), args=torch.ones(10), f='output.onnx')
  File "/path/to/python3.9/site-packages/torch/__init__.py", line 1868, in __getattr__
    return importlib.import_module(f".{name}", __name__)
  File "/path/to/python3.9/importlib/__init__.py", line 127, in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
  File "<frozen importlib._bootstrap>", line 1030, in _gcd_import
  File "<frozen importlib._bootstrap>", line 1007, in _find_and_load
  File "<frozen importlib._bootstrap>", line 986, in _find_and_load_unlocked
  File "<frozen importlib._bootstrap>", line 680, in _load_unlocked
  File "<frozen importlib._bootstrap_external>", line 850, in exec_module
  File "<frozen importlib._bootstrap>", line 228, in _call_with_frames_removed
  File "/path/to/python3.9/site-packages/torch/onnx/__init__.py", line 57, in <module>
    from ._internal.onnxruntime import (
  File "/path/to/python3.9/site-packages/torch/onnx/_internal/onnxruntime.py", line 38, in <module>
    import onnx
  File "/path/to/project/onnx.py", line 3, in <module>
    torch.onnx.export(nn.Linear(10, 10), args=torch.ones(10), f='output.onnx')
  File "/path/to/python3.9/site-packages/torch/onnx/utils.py", line 516, in export
    _export(
  File "/path/to/python3.9/site-packages/torch/onnx/utils.py", line 1590, in _export
    with exporter_context(model, training, verbose):
  File "/path/to/python3.9/contextlib.py", line 119, in __enter__
    return next(self.gen)
  File "/path/to/python3.9/site-packages/torch/onnx/utils.py", line 179, in exporter_context
    with select_model_mode_for_export(
  File "/path/to/python3.9/contextlib.py", line 119, in __enter__
    return next(self.gen)
  File "/path/to/python3.9/site-packages/torch/onnx/utils.py", line 166, in setup_onnx_logging
    is_originally_enabled = torch.onnx.is_onnx_log_enabled()
AttributeError: partially initialized module 'torch.onnx' has no attribute 'is_onnx_log_enabled' (most likely due to a circular import)
```

When `torch` imports `onnx` [here](https://github.com/pytorch/pytorch/blob/main/torch/onnx/_internal/onnxruntime.py#L34), the local file seems to be preferred over the pip installed `onnx`. According to #98271, I'm not the first one to encounter this issue, so we should probably throw a more descriptive error.

### Versions

torch==2.2.0.dev20231002+cu118
[fx] Add cache_result option in ShapeProp which can help cache interm tensor in meta
Labels: fb-exported, release notes: fx

Summary: While authoring transformations, we want to compare the result before and after to make sure it works as expected. Extending functionality in shape prop to catch interm result by adding "cache_result=True"

Test Plan:
buck run caffe2/test:fx

```
test_shape_prop_cache_result (test_fx.TestFX) ... ok
...
t_type_check_reshape_true (fx.test_gradual_type.TypeCheckerTest) ... ok
test_type_check_symbolic_inferenceconv2D_maxpool2d_flatten (fx.test_gradual_type.TypeCheckerTest) ... ok
test_type_check_transpose_False (fx.test_gradual_type.TypeCheckerTest) ... ok
test_type_check_transpose_true (fx.test_gradual_type.TypeCheckerTest) ... ok
test_type_maxpool2d_fully_static (fx.test_gradual_type.TypeCheckerTest) ... ok
test_type_typechecl_maxpool2d_3dinput (fx.test_gradual_type.TypeCheckerTest) ... ok
test_typecheck_basicblock (fx.test_gradual_type.TypeCheckerTest) ... ok

----------------------------------------------------------------------
Ran 1234 tests in 48.683s
```

Differential Revision: D49857532
Pytorch for Python 3.12 not available
Labels: module: build, triaged, module: python frontend

### 🚀 The feature, motivation and pitch

On Windows, Pytorch installation is not available:

```
C:\> python312.exe -m pip install torch
ERROR: Could not find a version that satisfies the requirement torch (from versions: none)
ERROR: No matching distribution found for torch
```

### Alternatives

_No response_

### Additional context

_No response_

cc @malfet @seemethere @albanD
bypass nvml for torch.cuda.device_count() if rocm
Labels: module: rocm, triaged, open source, ciflow/trunk, topic: not user facing, ciflow/periodic, rocm

This is a quick-fix to suppress printing "UserWarning: Can't initialize NVML" when calling torch.cuda.device_count() if [NVIDIA Management Library] (https://developer.nvidia.com/nvidia-management-library-nvml) (nvml module) is installed with ROCm. 
Fixes https://ontrack-internal.amd.com/browse/SWDEV-414997


cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang
dynamo support for functorch `grad` with pytree inputs.
Labels: open source, module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110417
* #107618
* #109433
* #110290

This PR removes a graph break when `torch.func.grad` is used with pytree inputs. In
summary: we flatten the inputs for adding graph inputs, followed by an unflatten for
calling the original function.

The following example is extracted from `test_grad_pytree` at
/test/dynamo/test_higher_order_ops.py/. There are no graph breaks after this PR.

```python
def fn(x):
    x1, x2 = x
    return x1.sin().sum() + x2

@torch.compile
def wrapper_fn(x):
    return torch.func.grad(fn)(x)
```

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
DynamicShapesExportTests::test_retracibility_dynamic_shapes times out with ASAN
Labels: oncall: pt2

The test has been timing out for some time, succeeding on retries.

Now it reliably times out: 
see linux-jammy-py3.9-clang12-asan / test (default, 2, 6, linux.4xlarge) at
https://hud.pytorch.org/pytorch/pytorch/commit/16e3f158b947c2c14a98178670f22c047a40807c

> dynamo/test_dynamic_shapes.py::DynamicShapesExportTests::test_retracibility_dynamic_shapes <- test/dynamo/test_export.py Command took >30min, returning 124


cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
Add support for `torch.Generator` type in TorchScript
Labels: triaged, open source, release notes: jit, ciflow/inductor, suppress-bc-linter

- Add support for `torch.Generator` type in TorchScript
- Add `generator` args to all `torch.nn.init` functions that call `uniform_` or `normal_`
- Add support for `torch.Generator` in LTC's TorchScript backend (CC: @wconstab)

CC: @eellison @davidberard98 @GlebKazantaev @behzad-a
WIP / TST: allow testing torch._numpy under Dynamo
Labels: module: numpy, open source, topic: not user facing, module: dynamo, ciflow/inductor, release notes: dynamo, ciflow/slow

Use conditional imports: when running under dynamo, import the original NumPy not torch._numpy. This is what we want to trace, not our implementation.

With this, the test suite passes with and without `PYTORCH_TEST_WITH_DYNAMO=1` (modulo a couple of test modules which are not meant to be compiled, e.g. `test_nep50_examples`). There are two new decorators, `x{fail,pass}ifTorchDynamo`, the `xpass` in most cases indicates a graph break and a fallback to eager for things we do not implement. 

cc @mruberry @rgommers @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
[experiment] do not use slow test json
Labels: ciflow/trunk, module: dynamo, ciflow/inductor, ciflow/slow

Fixes #ISSUE_NUMBER


cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
Add scatter_mm and bsr_scatter_mm operations.
Labels: module: sparse, open source, release notes: sparse, topic: new features, topic: performance

This PR introduces `scatter_mm` operation (compute `mm` of arbitrary pairs of tensors given in batches of tensors) that is used to implement `bsr_scatter_mm` that is equivalent to `bsr_dense_mm` (the `mm` operation on bsr and strided tensors). The implementation is provided both in Triton (when tensor dimensions are multiples of 16) and in PyTorch (otherwise).

The figures below illustrate the performance differences of `bsr_scatter_mm` and `bsr_dense_mm` (GPU: `NVIDIA GeForce RTX 2060 SUPER`). The first figure represents the performance equilibrium point in BSR tensor sparsity at which value `bsr_scatter_mm` or `bsr_dense_mm` have the same performance characteristics as `torch.matmul`. The second figure represents speedups from using `bsr_scatter_mm` at its performance equilibrium points with respect to `bsr_dense_mm`.

<img src="https://github.com/pytorch/pytorch/assets/402156/526d182e-937f-4812-a6c4-904f52d6d5ab" width="48%"> <img src="https://github.com/pytorch/pytorch/assets/402156/ccb606ab-1f3f-4133-887c-b56285f4f168" width="48%">

The same figures for GPU card `NVIDIA A100-SXM4-80GB`:

<img src="https://github.com/pytorch/pytorch/assets/402156/25466f1d-df34-4d1c-a975-afb478e4d9f0" width="48%"> <img src="https://github.com/pytorch/pytorch/assets/402156/6ada91f0-a20f-4f0d-8a48-1f4ccc60d08e" width="48%">

In sum:
- `bsr_scatter_mm` is about 2x faster than `bsr_dense_mm` for small block sizes of 16 and 32 and large tensors [GPU: `NVIDIA GeForce RTX 2060 SUPER`].
- `bsr_scatter_mm` is up to 2x faster than `bsr_dense_mm` for small block sizes of 16 and large tensors [GPU: `NVIDIA A100-SXM4-80GB`].
- `bsr_dense_mm` is up to 20 % faster than `bsr_scatter_mm` for block sizes of 64 or larger [GPU: `NVIDIA GeForce RTX 2060 SUPER`].
- However, `bsr_dense_mm` fails with `OutOfResources` exception for block sizes of 256 or larger whereas `bsr_scatter_mm` succeeds.

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #111489
* #111470
* __->__ #110396



cc @alexsamardzic @nikitaved @cpuhrsch @amjames @bhosmer
Pytorch LoadNativeLibrary issue
Labels: oncall: mobile

### 🐛 LoadNativeLibrary Crash in .ptl model

Hello,

I was trying to run object detection app but it crashes when model in Pytorch is loaded, in start of the model: 


```

  mModule = LiteModuleLoader.load(
                assetFilePath(
                   context,
                   ptlModel
                ))

```
 4080  4080 F /system/bin/app_process64: runtime.cc:1709] LoadNativeLibrary failed for "libjavacore.so": 

I'm using a trained model (.ptl) 

Similar load of model as in demo app : 
https://github.com/pytorch/android-demo-app/blob/master/ObjectDetection/README.md

gradle:
    implementation 'org.pytorch:pytorch_android_lite:1.13.0'
    implementation 'org.pytorch:pytorch_android_torchvision_lite:1.13.0'

How can I fix this error?

### Versions
pytorch_android_lite:1.13.0
pytorch_android_torchvision_lite:1.13.0
[dynamo][nn_module] Save the nn module object as value
Labels: module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110375
* #110535



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
BatchNorm layer 'num_batches_tracked' overwritten with default value when loading empty state_dict
Labels: module: nn, triaged

### 🐛 Describe the bug

BatchNorm layer 'num_batches_tracked' overwritten with default value when loading empty state_dict.

Probably should be fixed in: /torch/nn/modules/batchnorm.py lines 103-108 by checking if a value already exists.

```python
from collections import OrderedDict
from torchvision.models import mobilenet_v2, MobileNet_V2_Weights

sd_key = 'features.0.1.num_batches_tracked'
model = mobilenet_v2(weights=MobileNet_V2_Weights.IMAGENET1K_V1.DEFAULT)
print(f"before loading: {model.state_dict()['features.0.1.num_batches_tracked']}")

empty_dict = OrderedDict()
model.load_state_dict(empty_dict, strict=False)
print(f"after loading: {model.state_dict()['features.0.1.num_batches_tracked']}")
```

expected output: 
```
before loading: 736839
after loading: 736839
```

actual output: 
```
before loading: 736839
after loading: 0
```


### Versions

Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.4 (x86_64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.10 (main, Jan 15 2022, 11:48:04)  [Clang 13.0.0 (clang-1300.0.29.3)] (64-bit runtime)
Python platform: macOS-13.4-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Pro

Versions of relevant libraries:
[pip3] numpy==1.25.1
[pip3] torch==2.0.1
[pip3] torchinfo==1.8.0
[pip3] torchvision==0.15.2
[pip3] torchviz==0.0.2

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki
Dropout signature inconsistent between `torch.dropout`, `torch.nn.Dropout` and `torch.nn.functional.dropout`
Labels: module: nn, triaged, module: python frontend

### 🐛 Describe the bug

Both `torch.nn.Dropout` and `torch.nn.functional.dropout` default to `p=0.5` and `train=True`. However, `torch.dropout` requires them to be set manually.

```python
>>> import torch
>>> x = torch.randn(10, 20)
>>> torch.nn.Dropout()(x)
...
>>> torch.nn.functional.dropout(x)
...
>>> torch.dropout(x)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: dropout() missing 2 required positional argument: "p", "train"
```

Other functions are affected similarly:
- `torch.nn.functional.alpha_dropout`
- `torch.nn.functional.feature_alpha_dropout`

For consistency, we should also set the defaults for these functions in the `torch.*` module.

### Versions

Tested on `torch==2.0.1` and `torch==2.2.0.dev20230912+cu121`

cc @albanD @mruberry @jbschlosser @walterddr @mikaylagawarecki
[dynamo] Added support for type for generated custom object 
Labels: triaged, open source, module: dynamo, ciflow/inductor

Fixes #109056

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @ezyang for the issue.


cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
feat(inductor): Improve compilation speed and add `SGD` Optimizer back to Inductor
Labels: triaged, open source, topic: not user facing, module: inductor, module: dynamo, ciflow/inductor

Originally disabled in https://github.com/pytorch/pytorch/pull/105438


However, compilation times are plenty speedy to me:
![image](https://github.com/pytorch/pytorch/assets/9093549/023a7001-6a25-426d-9535-537d400e8bd3)

Could have benefited from perf changes to scheduler like cycle-detection optimizations.


CC @mlazos as original code author

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
feat(inductor): Add `RAdam` to Inductor by converting data-dependent control-flow to `torch.where`
Labels: triaged, open source, module: inductor, module: dynamo, ciflow/inductor, release notes: optim

For small epochs (adaptive learning rate = 1), this will be more costly than not computing `rect` and `adaptive_lr`, but unlikely by much - they are all fused into a single kernel anyway.


cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
feat(inductor): Improve `Adamax` to be better fused by Inductor and enable it
Labels: triaged, open source, module: inductor, module: dynamo, ciflow/inductor, release notes: optim

In order to avoid running into https://github.com/pytorch/pytorch/issues/110342, we replace complicated `cat + max` logic into a simple `torch.maximum`. (see also [here](https://github.com/facebookresearch/fairseq/blob/7409af7f9a7b6ddac4cbfe7cafccc715b3c1b21e/fairseq/optim/adamax.py#L151) in fairseq repo - the same operation with a different, less-explicit syntax).

Like https://github.com/pytorch/pytorch/pull/110339, we also ensure that `step` is on device.

As noted in https://github.com/pytorch/pytorch/issues/107006#issuecomment-1741841283, it is likely that `num_kernels=2` is optimal, unless one can move the `numel=1, shape=(0,)` scalar computations entirely into CPU.

That being said, in the foreach case, a parallel foreach scalar computation on GPU doesn't seem too bad.




cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[inductor]: Not handling `ConcatKernel/NopKernel` fusions leads to suboptimal fusions
Labels: oncall: pt2

### 🚀 The feature, motivation and pitch

Instead, perhaps try to refactor it as `ComputedBuffer`

Examples: Adamax

Already bad with `config.aggressive_fusion = True`:
![image](https://github.com/pytorch/pytorch/assets/9093549/b9679690-e984-48ae-83e6-ea77db55243f)

Worse when `config.aggressive_fusion = False`:
![image](https://github.com/pytorch/pytorch/assets/9093549/96d30ee2-0dcc-4b56-942b-fc70513dd3d9)


With `config.aggressive_fusion = False`, these kernels are treated as not sharing reads/writes. 

Further, the presence of these `ConcatKernel` nodes seems to obstruct fusion of other kernels, despite later being partly pruned away during kernel codegen. 
### Alternatives

_No response_

### Additional context

_No response_

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[dynamo] Add support for itertools.repeat
Labels: triaged, open source, topic: not user facing, module: dynamo, ciflow/inductor

Implemented call_repeat function in builtin.py to handle itertools.repeat within the Dynamo context. This addresses issue #110286 where support for itertools.repeat was requested.

Fixes #110286


cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
feat(optimizer): `Adagrad` will use `device` when `capturable` - True always when compiling with dynamo
Labels: triaged, open source, module: inductor, release notes: optim

Partial fix: https://github.com/pytorch/pytorch/issues/107006

CC: @mlazos as issue creator

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
fix(inductor): `ForeachKernelSchedulerNode` group shape should be opaque for graph debug
Labels: triaged, open source, release notes: fx, module: inductor, ciflow/inductor

~~Shape is assumed by `TensorMetadata` to be torch.Shape/tuple, however, some of the scheduler node groups utilize `int`, so convert to tuple.~~

Root cause is actually `foreach` scheduler node having silent-error group of int, when in fact it ought to be opaque `foreach`.

**Previously:** silent error / confusing shape of (0,)
![image](https://github.com/pytorch/pytorch/assets/9093549/5bc2a3c7-151f-4433-bbf8-044c7b03e989)

**Now:** clear that it is foreach which does not have well-defined shape:
![image](https://github.com/pytorch/pytorch/assets/9093549/8373080d-4519-4e74-8a3b-da463e9968da)

~~Alternate might be to create list of shapes for each of its subnodes. Actually, for debuggability sake, I may prefer this. We can ensure that the recursive generation of this string is only done dynamically in a debug code path. Else, incrementally computing it on initialization of ForeachKernel may also be feasible.~~ This is quite infeasible for 100s of params.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[pytorch] add should_deepcopy flag to AveragedModel
Labels: fb-exported, release notes: optim

Summary:
# Context
AveragedModel always deepcopies the passed in model. This can pose an issue if the model cannot be deepcopied, like FSDP models. Furthermore, users may want to do this logic themselves

# This diff
Adds `should_deepcopy` flag (default=True for backwards compatibility) so that if users want to pass an FSDP model, they can by turning the flag off

Test Plan: existing unit tests pass

Differential Revision: D49622071
[wip, dynamo] run all guard fail hooks after cache miss
Labels: module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110325

Attempt number 2 at https://github.com/pytorch/pytorch/issues/108950.

Will do testing before this PR is ready for review.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
torch._dynamo.exc.Unsupported: unexpected sourceless type bases: (<class 'torchrec.streamable.Pipelineable'>,)
Labels: good first issue, ezyang's list, oncall: pt2, module: dynamo, module: export

This is happening when trying to use `torch.export` on some code involving KJTs.

Internal stack trace: P840857345
Repro: P840859224

Notes:
- This is happening when we trace the construction of KJT, then pass that KJT as input to a module that we have configured to be `preserve_module_call_signature`. 
- If I remove the KJT construction, and pass a KJT we got directly from the input, things work (strange!)
- Ed suggested commenting out [this check](https://github.com/pytorch/pytorch/blob/359c2a53f59aab460fb15ba60f40fda537c33aeb/torch/_dynamo/variables/builtin.py#L1113). When we do this, things will pass, although I gather that this is unsafe because we would not be properly installing guards.

@ezyang says he knows what's going on and can post an OSS repro when he has some time

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng @avikchaudhuri @gmagogsfm
[ONNX][DO NOT REVIEW] Experiment passing 'dynamic' to dort backend
Labels: open source, release notes: onnx, module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110313
* #110477
* #110178
* #108376



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
[inductor] Decompose boolean min/max into all/any
Labels: open source, topic: not user facing, module: inductor, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110311
* #110310



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[ATen] Support multi dim any and all reductions
Labels: open source, release notes: python_frontend, topic: new features, ciflow/mps

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110310

This adds a new overload to `all` and `any` with support for multiple reduction dims.
```
all.dims(Tensor self, int[1]? dim=None, bool keepdim=False) -> Tensor
any.dims(Tensor self, int[1]? dim=None, bool keepdim=False) -> Tensor
```
Fix for out of bounds read in mobile interpreter FORMAT opcode handler
Labels: fb-exported, release notes: jit

Summary:
The FORMAT opcode for the mobile TorchScript interpreter contained an out of bounds read issue leading to memory corruption.

This change adds an explicit check that the number of inputs passed to the format method called when handling the FORMAT opcode is a valid and within bounds of the stack.

Test Plan: contbuild + OSS signals

Differential Revision: D49739095
Fix for out of bounds read in mobile interpreter INTERFACE_CALL opcode handler
Labels: fb-exported, release notes: mobile

Summary:
The INTERFACE_CALL opcode for the mobile TorchScript interpreter contained an out of bounds read issue leading to memory corruption.

This change adds an explicit check that the number of inputs passed to the format method called when handling the INTERFACE_CALL opcode is a valid and within bounds of the stack.

Test Plan: contbuild + OSS signals

Differential Revision: D49739450
Tests modify global state cause later tests to fail
Labels: module: ci, triaged, module: devx

This is more of a tracker for tests that modify global state somehow, and then tests that run after them change outcome depending on whether the first test was run or not.  I want to keep track of this to see how often it happens and why it happens.  

The tests afterwards generally get marked as flaky because they usually succeed on file level retry, which starts at the failing test and doesn't run the first test.

Examples of this include:
* Test 1 adds to dictionary that is imported or declared at the module level.  Test 2 reads that dictionary and changes outcome after seeing the new addition
* Test 1 imports extra modules.  Test 2 crawls through imports to check functions in each module.  The new module causes the outcome to change.
* Test 1 sets an attribute.  Test 2 sees the new attribute. (https://github.com/pytorch/pytorch/pull/110254)

Fixes can include cleaning up after the function or moving tests to run in different files (then CI will run in different processes).

cc @seemethere @malfet @pytorch/pytorch-dev-infra @ZainRizvi @kit1980 @huydhn
AOT Autograd Neg View Support
Labels: triaged, module: functionalization, oncall: pt2, module: aotdispatch

### 🐛 Describe the bug

The following fails with aot_eager.  It would be great if aot autograd could both a) handle neg_view correctly and b) remove ._neg_view() from the graph that is passed to inductor. 

```
import torch
device = torch.device

def forward(inp):
    a = inp._neg_view()
    b = a.add_(3)

    return a, b

inp = torch.rand([5, 5]).cuda()

out_eager = forward(inp.clone())
out_comp = torch._dynamo.optimize("aot_eager")(forward)(inp)

torch.testing.assert_close(out_comp, out_eager)
``` 

> AssertionError: Tensor-likes are not close!

### Versions

master

cc @bdhirsh @ezyang @msaroufim @wconstab @anijain2305
PYTORCH_TEST_WITH_DYNAMO=1 pytest -n 4 test/test_nestedtensor.py fails
Labels: high priority, triaged, module: nestedtensor, oncall: pt2

Running `PYTORCH_TEST_WITH_DYNAMO=1 pytest -n 4 test/test_nestedtensor.py` causes a non-deterministic number of failures (ranging between 14-70 for me). Perturbing the Dynamo config in CI (e.g. https://github.com/pytorch/pytorch/pull/110107) seems to expose the failures in CI, but CI is otherwise green.  I might disable them for now so that we can land stuff without risking flaky CI. Opening this issue to track the problem.

![image](https://github.com/pytorch/pytorch/assets/5652049/40ae91fa-b8ff-482e-8895-c97e867e912e)

cc @ezyang @gchanan @kadeng @cpuhrsch @jbschlosser @bhosmer @drisspg @msaroufim @wconstab @bdhirsh @anijain2305
TypeError: Got unsupported ScalarType BFloat16
Labels: triaged, module: numpy

### 🐛 Describe the bug

Converting bfloat16 to numpy fails:

```python
tensor = torch.ones((10, ), dtype=torch.bfloat16)
item.numpy()
```

### Versions

```python
(.venv) ➜  lightning git:(introduce_cache) ✗ python collect_env.py
Collecting environment information...
PyTorch version: 2.0.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 13.5.2 (arm64)
GCC version: Could not collect
Clang version: 14.0.3 (clang-1403.0.22.14.1)
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.15 (main, Dec  5 2022, 15:51:18)  [Clang 14.0.0 (clang-1400.0.29.202)] (64-bit runtime)
Python platform: macOS-13.5.2-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Apple M1 Max

Versions of relevant libraries:
[pip3] mypy==1.0.0
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.1
[pip3] pytorch-lightning==1.9.0
[pip3] torch==2.0.1
[pip3] torchdata==0.6.1
[pip3] torchmetrics==0.11.1
[pip3] torchvision==0.15.1
[conda] Could not collect
```

cc @mruberry @rgommers
feat(Pipeline Parallelism): use mincut optimization for local communication optimization
Labels: oncall: distributed

### 🚀 The feature, motivation and pitch

Currently, the pipeline parallelism API requires the user to make a decision about the partitioning:
https://github.com/pytorch/pytorch/blob/758735b739e119dfa66c447b81e8bd391ab765f7/torch/distributed/pipeline/sync/pipe.py#L224

However, I believe an _automated partitioner_ that reduces the cross-device communication could be employed. 


### Proposal
Of course, the device placement issue is hard enough (partitioning the bwd + fwd workload / memory to be evenly distributed across devices), but even once the graphs have been approximately placed, one may still want to search over a small radius to repartition operations so as to reduce cross-device communication. 

To do so, one may employ the mincut optimization to find the minimum cut of communication in a local region, similar to the min_cut_partitioner in inductor:

https://github.com/pytorch/pytorch/blob/758735b739e119dfa66c447b81e8bd391ab765f7/torch/_functorch/partitioners.py#L253

### Considerations
1. In e.g. [torchgpipe](https://arxiv.org/pdf/2004.09910.pdf) only a coarse-grained global partitioning strategy (`torchgpipe.balance`) is utilized, but not a local strategy.

### Alternatives
1. Unclear if this should be the role of PyTorch, which should provide the _capability_ for pipeline parallelism but not necessarily the fine-grained decisions.
    - Nonetheless, I believe that providing a default Torch partitioner / local repartitioner may still be useful to many users. 

### References
1. [Alpa - Inter+Intra Operator Parallelism](https://www.usenix.org/system/files/osdi22-zheng-lianmin.pdf) - unclear if the "local" min cut approach would be complementary with e2e principled, topology-aware approach with cost model like Alpa.

CC @wconstab 

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu @penguinwu
Allow nn.Transformer to be exported as ONNX.
Labels: module: onnx, triaged, open source, release notes: onnx

Fixes #110255.

During ONNX export tracing, the `is_causal` parameter is not actually a `bool` but instead `tensor(bool)`. The native function `scaled_dot_product_attention` does not seem to handle this case, so we can just cast it to a Python `bool` for now.
[S366352] Avoid creating ncclStartEvent for deadlocking
Labels: fb-exported, release notes: distributed (c10d)

Summary: enableTiming will make us create one more cudaEvent per collective: https://fburl.com/code/6hqzik63. From S366352, it looks we run into deadlock in cudaEventDestroy() so reducing the number of events may help. Also I don't really see ncclStartEvent used anywhere for desyncDebug.

Test Plan: sandcastle

Differential Revision: D49760869
[FX/const_fold][bugfix] Set codegen on wrapping GraphModule
Labels: fb-exported, release notes: fx

Summary: The split module pass does not respect retaining custom codegen. This meant const fold was broken if we use concrete args or torch.export etc. that set codegen. Fix this by copying over codegen from the original graph we're doing folding on.

Test Plan: Added test to test_fx_const_fold.

Differential Revision: D49760518
[dynamo][guard-refactor] Global Guard Accessor
Labels: module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #110590
* #110735
* #110589
* __->__ #110265
* #108839



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
Cannot avoid kineto_LIBRARY-NOTFOUND error when using pre-built pytorch
Labels: oncall: binaries, oncall: profiler

### 🐛 Describe the bug

When using the pre-built pytorch like this example in the docs:
https://pytorch.org/cppdocs/installing.html

I am getting the following error:
```
CMake Warning at cmake-build-debug/_deps/pytorch-src/torch/share/cmake/Torch/TorchConfig.cmake:22 (message):
  static library kineto_LIBRARY-NOTFOUND not found.
```

I inspected the zip and saw that Kineto is expected to be `ON`, and will warn with that message. However, Kineto is not included in the releases.

### Versions

N/A

cc @seemethere @malfet @robieta @chaekit @aaronenyeshi @nbcsm @guotuofeng @guyang3532 @gaoteng-git @tiffzhaofb @dzhulgakov @davidberard98
[fbcode] s,\bnp\.bool\b,bool,
Labels: caffe2, module: mkldnn, fb-exported

Test Plan: sandcastle and visual inspection

Differential Revision: D49196618




cc @gujinghui @PenghuiCheng @XiaobingSuper @jianyuh @jgong5 @mingfeima @sanchitintel @ashokei @jingxu10 @min-jean-cho @yanbing-j @Guobing-Chen @Xia-Weiwen
cuda/tf32 docs are outdated
Labels: module: docs, module: cuda, triaged, module: tf32

### 📚 The doc issue

Would it be possible to update:

https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices

1. there is the new way to define precision which is not just `torch.backends.cuda.matmul.allow_tf32` but 3 levels as defined by:
https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
it'd be awesome to update the cuda doc to reflect that and to update the benchmarks and the error with medium precision.

2. I'd imagine that `torch.set_float32_matmul_precision("medium")` is really relevant for H100 as it says:
> if a fast matrix multiplication algorithm using that datatype internally is available. 

or is it available in A100 as well?

Thank you!


cc @svekars @carljparker @ptrblck @zasdfgbnm
Accessing Particular Nightly Builds Don't Work
Labels: oncall: binaries, triaged

### 🐛 Describe the bug

I would like to use a particular nightly release version of libtorch. I need to do this because:

1. Official version releases are sparse
2. Relying on "latest" is dangerous -- consistent builds/deployments is much safer

On the official page:
https://pytorch.org/get-started/locally/

I can only see the official releases or the nightly releases:
![image](https://github.com/pytorch/pytorch/assets/17555572/d3f36dcb-5767-4dee-b737-7834e13f9798)

However, in the github workflows, I can see that date stamped releases are being uploaded as well. For example, with this run flavor:
https://github.com/pytorch/pytorch/actions/runs/6335907738/job/17207923080

I can see:
```
record_file=libtorch-cxx11-abi-shared-with-deps-2.2.0.dev20230928+cpu.zip
..
record_file=libtorch-cxx11-abi-shared-with-deps-latest.zip
```

However, attempting to access these urls results in an 404. If I attempt a similar URL with the wheel storage:
https://download.pytorch.org/whl/nightly/cpu/torch-2.1.0.dev20230901-cp310-none-macosx_11_0_arm64.whl

It works as expected. Can this be fixed?

Thanks,
Anthony


### Versions

N/A

cc @seemethere @malfet
DISABLED test_detach_cpu_float16 (__main__.TestNestedTensorDeviceTypeCPU)
Labels: triaged, module: flaky-tests, module: nestedtensor, skipped, oncall: pt2

Platforms: dynamo

This test was disabled because it is failing in CI. See [recent examples](https://hud.pytorch.org/flakytest?name=test_detach_cpu_float16&suite=TestNestedTensorDeviceTypeCPU) and the most recent trunk [workflow logs](https://github.com/pytorch/pytorch/runs/17230279068).

Over the past 3 hours, it has been determined flaky in 11 workflow(s) with 33 failures and 11 successes.

**Debugging instructions (after clicking on the recent samples link):**
DO NOT ASSUME THINGS ARE OKAY IF THE CI IS GREEN. We now shield flaky tests from developers so CI will thus be green but it will be harder to parse the logs.
To find relevant log snippets:
1. Click on the workflow logs linked above
2. Click on the Test step of the job so that it is expanded. Otherwise, the grepping will not work.
3. Grep for `test_detach_cpu_float16`
4. There should be several instances run (as flaky tests are rerun in CI) from which you can study the logs.


Test file path: `test_nestedtensor.py`

cc @cpuhrsch @jbschlosser @bhosmer @drisspg @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
[inductor] Triton mm codegen uses `make_block_ptr`
Labels: triaged, open source, topic: not user facing, module: inductor, ciflow/inductor

Partially fixes: https://github.com/pytorch/pytorch/issues/109420

No perf degr (on RTX 4070 Laptop Ada Lovelace):

Before:
![image](https://github.com/pytorch/pytorch/assets/9093549/a87fd2c7-2ee4-415c-9006-93c7337bc956)

After:
![image](https://github.com/pytorch/pytorch/assets/9093549/ab173ed8-b277-46a7-b3b7-41f5992fca12)

The reason that the curve is not equal to cublas/triton for smaller M/N is because the triton autotune configs are not tuned for commercial GPUs. See: https://github.com/pytorch/pytorch/issues/109489

`uint4x2_mixed_mm` is skipped due to complicated use of the ptrs to do bitwise hacks.

CC: @jansel @drisspg 

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[ONNX] Replace torchscript with new graph building
Labels: module: onnx, open source, release notes: onnx

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110241

This PR works with https://github.com/microsoft/onnxscript/pull/1071, providing a general idea/starting point of new graph building without TorchScriptGraph. Tested with only `test_add` unittest case in `test_fx_to_onnx_onnxruntime.py`.

NOTE: still need to test with:
1. `add_module_call`
2. multiple outputs
3. initializers
4. sequence type and split node

More detail should be found in https://github.com/microsoft/onnxscript/pull/1071
add full permuted
Labels: module: inductor, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110234
* #110242



cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler
[fbcode] small fix on dataloader and content_understanding/main
Labels: fb-exported, release notes: dataloader

Summary:
* when dl shuts down, sometimes the order of object `__del__` and object clean up may mess up, which may raise exception complaining certain properties of the obj is not defined. adding a catch to avoid it from happening.

* remove the extra `register_new_resolver` as pointed in https://www.internalfb.com/diff/D49448666?dst_version_fbid=132094109961892&transaction_fbid=1285292328796395

Test Plan: cu-cli launch -p xray_video -c xrv_mae_vit_test --no-maybe-pull -l

Reviewed By: ejguan

Differential Revision: D49710645
[RELAND] Disallow skipping dynamo
Labels: ciflow/trunk, topic: not user facing, module: dynamo, ciflow/inductor, module: export

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110222

Previous discussion: https://github.com/pytorch/pytorch/pull/109476

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng
DISABLED test_noncontiguous_samples__native_batch_norm_legit_cuda_float32 (__main__.TestCommonCUDA)
Labels: triaged, module: flaky-tests, skipped, oncall: pt2

Platforms: inductor

This test was disabled because it is failing in CI. See [recent examples](https://hud.pytorch.org/flakytest?name=test_noncontiguous_samples__native_batch_norm_legit_cuda_float32&suite=TestCommonCUDA) and the most recent trunk [workflow logs](https://github.com/pytorch/pytorch/runs/17203542961).

Over the past 3 hours, it has been determined flaky in 2 workflow(s) with 2 failures and 2 successes.

**Debugging instructions (after clicking on the recent samples link):**
DO NOT ASSUME THINGS ARE OKAY IF THE CI IS GREEN. We now shield flaky tests from developers so CI will thus be green but it will be harder to parse the logs.
To find relevant log snippets:
1. Click on the workflow logs linked above
2. Click on the Test step of the job so that it is expanded. Otherwise, the grepping will not work.
3. Grep for `test_noncontiguous_samples__native_batch_norm_legit_cuda_float32`
4. There should be several instances run (as flaky tests are rerun in CI) from which you can study the logs.


Test file path: `test_ops.py`

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
CI: rocm `(default, 1, 3, linux.rocm.gpu)` is very slow
Labels: module: rocm, module: ci, triaged, module: devx


## Current Status
ongoing

## Issue looks like
Taking ~2.5 hours on https://github.com/pytorch/pytorch/pull/110167
![image](https://github.com/pytorch/pytorch/assets/9093549/f36d4d3b-93eb-48ee-a0f9-e945643376c6)

Taking 3.5+ hours on https://github.com/pytorch/pytorch/pull/109976
![image](https://github.com/pytorch/pytorch/assets/9093549/68af3813-84be-4034-ac24-2e61ec4e1856)


## User impact
Slower merging

## Root cause
Seems like it may have been introduced in https://github.com/pytorch/pytorch/pull/109817 @malfet 
![image](https://github.com/pytorch/pytorch/assets/9093549/18b78912-8304-4e95-bc80-a65b4b88e41d)

## Mitigation
Not sure

## Prevention/followups
Investigate cause of slow running time or split up tests into smaller test jobs. Try to make the tests run in similar time to CUDA tests (~1.5 hours)


cc @jeffdaily @sunway513 @jithunnair-amd @pruthvistony @ROCmSupport @dllehr-amd @jataylo @hongxiayang @seemethere @malfet @pytorch/pytorch-dev-infra @ZainRizvi @kit1980 @huydhn @clee2000
```[tasklist]
### Tasks
```
[ONNX] Add sanity check in CI for onnxbench
Labels: open source, topic: not user facing

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* #110313
* #110477
* __->__ #110178
* #108376

ONNX CI to run benchmark with `--quick` to validate the onnxbench infra.
Module-level bufferization for torch dynamo module spanning multiple `fx.GraphModule`
Labels: feature, triaged, oncall: pt2, module: aotdispatch

### 🚀 The feature, motivation and pitch

Some details across:
https://github.com/pytorch/pytorch/issues/109240#issuecomment-1737764251
https://github.com/pytorch/pytorch/issues/106596#issuecomment-1737835937

In summary:
1. It is likely that since we do not analyze buffer aliases (read/write) which may occur in graph breaks, we require copying inputs whenever they are mutated by the fx graphs. (TODO: add source reference?)
2. This can be complicated for scenarios like https://github.com/pytorch/pytorch/issues/109240#issuecomment-1737764251 where we are unsure when we want to commit the mutations back to the original buffer - which would be unnecessarily constrained by the time-sensitive nature of potential eager-mode mutations.

Solution:
1. If we can analyze the code in the eager mode computations (r/w accesses to given buffer), then we can better plan for bufferization across the entire module, reducing copies in between graph breaks and fx graphs. However, pointer alias analysis is generally impossible to do in opaque code (due to pointer encapsulation), so this may be completely infeasible.

However, it _may_ be the case that we know that _certain_ fx graph ouputs can _never_ be mutated by eager code - because they are **owned** (i.e. created) by the graph, hence there can be no pointer alias to it without explicit reference to the pointer in the eager code.

### Alternatives

Unclear

### Additional context
Unclear if any of this is already tracked in torch dynamo [`SideEffects`](https://github.com/pytorch/pytorch/blob/main/torch/_dynamo/side_effects.py)

cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305
Add cuSPARSELt to the nightlies
Labels: oncall: binaries, module: build, triaged

## Issue description

cuSPASRELt is a library that offers accelerated semi-structured sparse matmul. We are looking to integrate v0.5.0
The download link can be found here: https://developer.nvidia.com/cusparselt-downloads 
The library is ~ 20 Mb. 


We merged build support for cuSPASRELt in [this PR](https://github.com/pytorch/pytorch/pull/103700)
To build pytorch with cuSPASRELt, we need to set two env variables, 

```
USE_CUSPARSELT=1
CUSPARSELT_ROOT=/path/to/cusparselt/download
```

cc @seemethere @malfet
Fix for PyTorch mobile flatbuffer loader out of bounds reads
Labels: fb-exported, release notes: mobile

Summary:
The mobile_ivalue_size field in the mobile_bytecode flatbuffer schema can be larger than the ivalues vector. This introduces potential for memory corruption when parsing the mobile_bytecode Module.

This diff fixes the issue by ensuring that  mobile_ivalue_size is less than the size of the ivalues vector.

Test Plan: contbuild & OSS CI

Differential Revision: D49687548
Add _worker_end_fn_t to the DataLoader
Labels: module: dataloader, triaged, module: data

### 🚀 The feature, motivation and pitch

For my particular use case, I need the dataset to perform some operations when the data loaders workers are terminated. When the `iteration_end` value in the `worker.py` file is True.

Similar to `_worker_init_fn_t` for setting the seeds, I would like to have a `_worker_end_fn_t`  that takes the dataset as an argument.

### Alternatives

I haven't considered any other solutions but right now, I need to patch PyTorch.

### Additional context

This is required for the dataset to persist some data on termination.

cc @SsnL @VitalyFedyunin @ejguan @dzhulgakov
[TESTING] Capture scalar/dynamic by default
Labels: module: dynamo, ciflow/inductor

Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom):
* __->__ #110155
* #109893

Signed-off-by: Edward Z. Yang <ezyang@meta.com>

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @Xia-Weiwen @wenzhe-nrv @jiayisunx @chenyang78 @aakhundov @kadeng