Is there a way for tensorflow's scatter_nd_update to work with indices of type tf.int64 instead of tf.int32?

[BoubakerA]

I'm working with a tf.Variable for which I need to update the values using the scatter_nd_update function. The problem is that my variable has a number of elements equals to 512x512x2048x7 = 3 758 096 384 which is larger than the range converd by int32 (2147483647). And apparently the scatter_nd_update function works only with tf.int32 number and does not work with tf.int64.

Here is an example of the code that produces an overflow error:

import tensorflow as tf

image = tf.Variable(tf.ones((512,512,2048,7), dtype=tf.float32))

update = tf.constant([0.0])
index = tf.constant([[0,0,0,0]], dtype=tf.int64)
image.scatter_nd_update(index, update)

print(f"image shape {tf.shape(image)}")

The error that is raised is the following:
F ./tensorflow/core/util/gpu_launch_config.h:129] Check failed: work_element_count > 0 (-536870912 vs. 0) Aborted (core dumped)

Do you have any idea about how this problem can be resolved?

[gototophao]

i also met this problem, have you solved it?

[BoubakerA]

i also met this problem, have you solved it?

No, I didn't.