fix torch.mean

[qsh-zh]

description

torch.mean(a_tensor) gives torch.sum results

reproduce

env

Collecting environment information...
PyTorch version: 2.7.0a0+79aa17489c.nv25.04
Is debug build: False
CUDA used to build PyTorch: 12.9
ROCM used to build PyTorch: N/A

OS: Ubuntu 24.04.2 LTS (x86_64)
GCC version: (Ubuntu 13.3.0-6ubuntu2~24.04) 13.3.0
Clang version: Could not collect
CMake version: version 3.31.6
Libc version: glibc-2.39

Python version: 3.12.3 (main, Feb  4 2025, 14:48:35) [GCC 13.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-1055-aws-x86_64-with-glibc2.39
Is CUDA available: True
CUDA runtime version: 12.9.41
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA H100 80GB HBM3
GPU 1: NVIDIA H100 80GB HBM3
GPU 2: NVIDIA H100 80GB HBM3
GPU 3: NVIDIA H100 80GB HBM3
GPU 4: NVIDIA H100 80GB HBM3
GPU 5: NVIDIA H100 80GB HBM3
GPU 6: NVIDIA H100 80GB HBM3
GPU 7: NVIDIA H100 80GB HBM3

Nvidia driver version: 550.54.15
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.9.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.9.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.9.0
Is XPU available: False
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:                      48 bits physical, 48 bits virtual
Byte Order:                         Little Endian
CPU(s):                             192
On-line CPU(s) list:                0-191
Vendor ID:                          AuthenticAMD
Model name:                         AMD EPYC 7R13 Processor
CPU family:                         25
Model:                              1
Thread(s) per core:                 2
Core(s) per socket:                 48
Socket(s):                          2
Stepping:                           1
BogoMIPS:                           5299.99
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm cmp_legacy cr8_legacy abm sse4a misalignsse 3dnowprefetch topoext perfctr_core invpcid_single ssbd ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 invpcid rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 clzero xsaveerptr rdpru wbnoinvd arat npt nrip_save vaes vpclmulqdq rdpid
Hypervisor vendor:                  KVM
Virtualization type:                full
L1d cache:                          3 MiB (96 instances)
L1i cache:                          3 MiB (96 instances)
L2 cache:                           48 MiB (96 instances)
L3 cache:                           384 MiB (12 instances)
NUMA node(s):                       2
NUMA node0 CPU(s):                  0-47,96-143
NUMA node1 CPU(s):                  48-95,144-191
Vulnerability Gather data sampling: Not affected
Vulnerability Itlb multihit:        Not affected
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Not affected
Vulnerability Retbleed:             Not affected
Vulnerability Spec rstack overflow: Vulnerable
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
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Not affected

Versions of relevant libraries:
[pip3] DISTS-pytorch==0.1
[pip3] flake8==7.1.0
[pip3] intel-openmp==2021.4.0
[pip3] mkl==2021.1.1
[pip3] mkl-devel==2021.1.1
[pip3] mkl-include==2021.1.1
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] nvidia-cudnn-frontend==1.11.0
[pip3] nvidia-pytriton==0.5.5
[pip3] nvtx==0.2.11
[pip3] onnx==1.17.0
[pip3] onnx_graphsurgeon==0.5.8
[pip3] onnxruntime==1.22.0
[pip3] open_clip_torch==2.32.0
[pip3] optree==0.14.1
[pip3] pynvjitlink==0.3.0
[pip3] pytorch-lightning==2.5.1.post0
[pip3] pytorch-ranger==0.1.1
[pip3] pytorch-triton==3.2.0+git4b3bb1f8b.nvinternal
[pip3] slangtorch==1.3.8
[pip3] tbb==2021.13.1
[pip3] torch==2.7.0a0+79aa17489c.nv25.4
[pip3] torch_automated_profiler==1.10.0
[pip3] torch-fidelity==0.3.0
[pip3] torch-geometric==2.6.1
[pip3] torch-optimizer==0.3.0
[pip3] torch_tensorrt==2.7.0a0
[pip3] torchdata==0.11.0
[pip3] torchmetrics==1.7.2
[pip3] torchprofile==0.0.4
[pip3] torchvision==0.22.0a0
[pip3] tritonclient==2.51.0
[conda] Could not collect

code

from batch_invariant_ops import set_batch_invariant_mode
import torch

torch.set_default_device("cuda")
torch.set_default_dtype(torch.bfloat16)

rand_tensor = torch.randn(1, 1, 128).cuda()

with set_batch_invariant_mode(True):
    print(
        torch.mean(rand_tensor)
    )

print(torch.mean(rand_tensor))

two torch.mean call gives different results

issue

dim is empty list when using torch.mean(a_tensor)