aten::bucketize.Scalar doesn't work in compile mode

[Silv3S]

🐛 Describe the bug

aten::bucketize.Scalar works in eager mode, but it can't run compile mode. Scalar is passed to https://github.com/pytorch/pytorch/blob/main/torch/_refs/__init__.py#L5901 and it tries to call tensor.shape() on it.

aten::bucketize.Tensor works as expected in both modes.

Error logs

TorchRuntimeError: Failed running call_function <built-in method bucketize of type object at 0x7c6c17281760>(*(6.3, FakeTensor(..., size=(s0,))), **{}):
'float' object has no attribute 'shape'

from user code:
File "/usr/local/lib/python3.10/dist-packages/torch/_dynamo/external_utils.py", line 36, in inner
return fn(*args, **kwargs)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

You can suppress this exception and fall back to eager by setting:
import torch._dynamo
torch._dynamo.config.suppress_errors = True

Minified repro

import torch
fn = torch.bucketize
fn(6.3, torch.tensor([1.4, 3.2, 5.1, 8.9]))

fn = torch.compile(fn)
fn(6.3, torch.tensor([1.4, 3.2, 5.1, 8.9]))

Versions

PyTorch version: 2.3.1+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.3 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.30.1
Libc version: glibc-2.35

Python version: 3.10.12 (main, Mar 22 2024, 16:50:05) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-6.1.85+-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: Could not collect
Nvidia driver version: Could not collect
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture: x86_64
CPU op-mode(s): 32-bit, 64-bit
Address sizes: 46 bits physical, 48 bits virtual
Byte Order: Little Endian
CPU(s): 2
On-line CPU(s) list: 0,1
Vendor ID: GenuineIntel
Model name: Intel(R) Xeon(R) CPU @ 2.20GHz
CPU family: 6
Model: 79
Thread(s) per core: 2
Core(s) per socket: 1
Socket(s): 1
Stepping: 0
BogoMIPS: 4399.99
Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities
Hypervisor vendor: KVM
Virtualization type: full
L1d cache: 32 KiB (1 instance)
L1i cache: 32 KiB (1 instance)
L2 cache: 256 KiB (1 instance)
L3 cache: 55 MiB (1 instance)
NUMA node(s): 1
NUMA node0 CPU(s): 0,1
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit: Not affected
Vulnerability L1tf: Mitigation; PTE Inversion
Vulnerability Mds: Vulnerable; SMT Host state unknown
Vulnerability Meltdown: Vulnerable
Vulnerability Mmio stale data: Vulnerable
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed: Vulnerable
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass: Vulnerable
Vulnerability Spectre v1: Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers
Vulnerability Spectre v2: Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Vulnerable (Syscall hardening enabled)
Vulnerability Srbds: Not affected
Vulnerability Tsx async abort: Vulnerable

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.3.1+cu121
[pip3] torchaudio==2.3.1+cu121
[pip3] torchsummary==1.5.1
[pip3] torchtext==0.18.0
[pip3] torchvision==0.18.1+cu121
[pip3] triton==2.3.1
[conda] Could not collect

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames

[ezyang]

It's likely a mistake in the meta function forgetting scalars are allowed

[Harry262000]

Hey, @ezyang The aten::bucketize.Scalar function works in eager mode but fails in compile mode when a Python float is passed as the first argument. It tries to call tensor.shape() on the scalar, resulting in an error because scalars do not have a shape attribute. However, when the scalar is wrapped in a torch.tensor, the function works correctly in both modes.

1. compiled_fn = torch.compile(fn)

2. compiled_fn(6.3, torch.tensor([1.4, 3.2, 5.1, 8.9])) # This raises an error

3. compiled_fn(torch.tensor(6.3), torch.tensor([1.4, 3.2, 5.1, 8.9])) # This works correctly

I'm new to open source. If I described anything incorrectly, please correct me. I appreciate any guidance on how to better articulate this issue.

Thank you!!

[Harry262000]

Hi @ezyang , @Silv3S

I wanted to follow up on the issue I raised about the aten::bucketize.Scalar function in TorchScript when passing a Python float as the first argument. It works fine in eager mode but fails in compiled mode, as described in my previous comment.

I understand that you may be busy, but I would appreciate any updates or guidance on this issue.

[Silv3S]

However, when the scalar is wrapped in a torch.tensor, the function works correctly in both modes.

If you wrap scalar as torch.tensor then bucketize operator will be executed as bucketize.Tensor. Goal of this issue is to make changes in internal PyTorch code, so user can pass scalar and get correct results, without converting it to tensor beforehand