`tf.raw_ops.ImmutableConst` causes Segfault (`SIGSEGV`) with large shape dimensions

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

No

Source

source

TensorFlow version

2.20.0

Custom code

Yes

OS platform and distribution

Linux Ubuntu 24.04

Mobile device

No response

Python version

3.12

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

A Segmentation Fault occurs in tf.raw_ops.ImmutableConst when provided with large shape dimensions (approx ~2.1 billion). The operation appears to suffer from an integer overflow or missing bounds check when calculating the required memory size against the provided file region, causing the Python process to crash immediately instead of raising an InvalidArgumentError or ResourceExhaustedError.

Here is a gist

Standalone code to reproduce the issue

import tensorflow as tf
import os

print(f"Tensorflow version {tf.__version__}", flush=True)

# Create a dummy file to pass the file existence check
file_name = "replicate"
with open(file_name, "wb") as f:
    f.write(b"0")

try:
    # Trigger SIGSEGV with large dimensions
    tf.raw_ops.ImmutableConst(
        dtype=tf.float64,
        shape=[2147482841, 2147485163],
        memory_region_name=file_name,
        name="NCHW"
    )
except Exception as e:
    print(f"Caught exception: {e}")
finally:
    if os.path.exists(file_name):
        os.remove(file_name)

print("Print this if the code does not crash", flush=True)

Relevant log output

Tensorflow version 2.20.0
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1769092444.939800 3401849 gpu_device.cc:2431] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 12.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
W0000 00:00:1769092444.947635 3401849 gpu_device.cc:2431] TensorFlow was not built with CUDA kernel binaries compatible with compute capability 12.0. CUDA kernels will be jit-compiled from PTX, which could take 30 minutes or longer.
I0000 00:00:1769092445.031182 3401849 gpu_device.cc:2020] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 29724 MB memory:  -> device: 0, name: NVIDIA GeForce RTX 5090, pci bus id: 0000:02:00.0, compute capability: 12.0
Segmentation fault (core dumped)