Repository: pytorch
Issue: Applying python 'next' directly on a tensor is painfully slow
Labels: triaged,enhancement,module: python frontend

### 🐛 Describe the bug

Applying the standard python `next` function directly to a pytorch tensor is mysteriously (and painfully) slow relative to iterating over the indices of that tensor and accessing each value independently. This is not due to, at least explicitly, breaking `next`: when printing the values being considered during the loop, only the expected values are considered.

Example: consider this code
```python
import torch
from time import process_time

def timed(f):  # will allow us to time the functions below
    def timed_f(*a, **kw):
        start = process_time()
        res = f(*a, **kw)
        total = process_time() - start
        print("TIME:", (" " * 4) + f.__name__, "took:", total, "s", flush=True)
        return res
    return timed_f


def predicate(loud): 
    def _predicate(b):
        if loud:   # will allow us to check which items are being considered in the `next` run
            print(b)
        return b>0.1 and b<0.9
    return _predicate

@timed
def get_next_direct(a, pred):  # apply next directly to a flat tensor
    return next(v for v in a if pred(v))

@timed
def get_next_indirect(a, pred):  # apply next indirectly, by iterating over the positions in the tensor
    return next(a[i] for i in range(len(a)) if pred(a[i]))

a = torch.rand(int(1e7))
pred = predicate(False)  # quiet run, for clean timing
get_next_direct(a, pred), get_next_indirect(a, pred)

print("now with printing calls to the predicate")
pred = predicate(True)  # loud run, to see if the next is failing to break once the condition is satisfied - it is not, this is not the problem
get_next_direct(a, pred), get_next_indirect(a, pred)
```

The time difference between `get_next_direct` and `get_next_indirect` is enormous (demonstration with my specific run). As we can see when printing from the predicate, this is not due to `next` failing to break - only the expected evaluations are made:
```
TIME:     get_next_direct took: 10.626148 s
TIME:     get_next_indirect took: 0.00022200000000083264 s
now with printing calls to the predicate
testing: tensor(0.7398)
TIME:     get_next_direct took: 10.294364999999992 s
testing: tensor(0.7398)
TIME:     get_next_indirect took: 0.0006169999999912079 s
```

### Versions

PyTorch version: 2.2.0
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 15.0.1 (arm64)
GCC version: Could not collect
Clang version: 16.0.0 (clang-1600.0.26.3)
CMake version: Could not collect
Libc version: N/A

Python version: 3.11.5 (main, Sep 11 2023, 08:31:25) [Clang 14.0.6 ] (64-bit runtime)
Python platform: macOS-15.0.1-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 Pro

Versions of relevant libraries:
[pip3] msgpack-numpy==0.4.8
[pip3] numpy==1.26.3
[pip3] pytorch-lightning==2.0.3
[pip3] torch==2.2.0
[pip3] torchaudio==2.2.0
[pip3] torchmetrics==1.1.2
[pip3] torchvision==0.15.2a0
[conda] msgpack-numpy             0.4.8                    pypi_0    pypi
[conda] numpy                     1.26.3          py311he598dae_0  
[conda] numpy-base                1.26.3          py311hfbfe69c_0  
[conda] pytorch                   2.2.0                  py3.11_0    pytorch
[conda] pytorch-lightning         2.0.3           py311hca03da5_0  
[conda] torchaudio                2.2.0                 py311_cpu    pytorch
[conda] torchmetrics              1.1.2           py311hca03da5_0  
[conda] torchvision               0.15.2          cpu_py311he74fb5d_0  

cc @msaroufim @albanD