model.fit fails when the number of rows exceeds Int32.MaxValue

Issue type

Bug

Have you reproduced the bug with TensorFlow Nightly?

Yes

Source

source

TensorFlow version

2.19.0-dev20241117

Custom code

Yes

OS platform and distribution

MacOS 15.1.0

Mobile device

No response

Python version

3.10

Bazel version

No response

GCC/compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current behavior?

I would expect model.fit to handle training on extremely large NumPy arrays without limitations.

Standalone code to reproduce the issue

import numpy as np
from keras import Sequential
from keras.layers import Dense

n = 2_147_483_648
x = np.zeros(n).astype(np.float32)
y = x

model = Sequential([
    Dense(64, activation="relu", input_shape=(1,)),
    Dense(1, activation="sigmoid")
])
model.compile(optimizer="adam", loss="binary_crossentropy")
model.fit(x=x,y=y, epochs=1, batch_size=1024, verbose=1)

Relevant log output

ValueError: Invalid value in tensor used for shape: -2147483648