TensorFlow: Number of elements was larger than representable by 32-bit output type

Question

For a large tensor shape such as (72, 7007313, 5) and using a large network such as this:

import keras
import numpy as np
from keras.optimizers import Adam
from keras.models import Sequential
from keras.layers import Dense, Flatten

Y = np.random.randint(low=0, high=1, dtype=np.int16, size=(72, 2))
X = np.random.randint(low=1, high=5, dtype=np.int16, size=(72, 7007313, 5))

model = keras.models.Sequential()
model.add(Dense(1024, input_shape=X.shape[1:], activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Dense(1024, activation='relu'))
model.add(Flatten())
model.add(Dense(2, activation='softmax'))
model.compile(optimizer=Adam(),loss='binary_crossentropy', metrics=['accuracy'])
model.summary()
model.fit(X, Y, epochs=1, batch_size=1)

Training this tensor using this neural network with batch_size=1 gives the following error:

_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_1 (Dense)              (None, 7007313, 1024)     6144      
_________________________________________________________________
dense_2 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
dense_3 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
dense_4 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
dense_5 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
dense_6 (Dense)              (None, 7007313, 1024)     1049600   
_________________________________________________________________
flatten_1 (Flatten)          (None, 7175488512)        0         
_________________________________________________________________
dense_7 (Dense)              (None, 2)                 1435097702
=================================================================
Total params: 14,356,231,170
Trainable params: 14,356,231,170
Non-trainable params: 0
_________________________________________________________________
Epoch 1/1
Traceback (most recent call last):
tensorflow.python.framework.errors_impl.InvalidArgumentError:
Number of elements was larger than representable by 32-bit output type

That is because the network has 14,356,231,170 trainable items, which is above the default TensorFlow setup which indexes tensors using int32 (i.e. a max of 4,294,967,296 trainable items).

So my question is: How can I change TensorFlow to index tensors using int64 instead?

Answer 1

OP_REQUIRES(
          ctx, FastBoundsCheck(size, std::numeric_limits<int32>::max()),
          errors::InvalidArgument("Number of elements was larger than "
                                  "representable by 32-bit output type"));

By design, you cannot have TensorFlow to have tensor sizes above the int32 limit.

Source code reference