TPU v3-8 CrossReplicaSum_33 Error

[vivekjoshy]

Issue type

Support

Have you reproduced the bug with TensorFlow Nightly?

No

Source

source

TensorFlow version

2.15.0

Custom code

Yes

OS platform and distribution

No response

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'm getting this error when I run model.fit

tensorflow/core/tpu/kernels/tpu_compilation_cache_external.cc:112] Asked to propagate a dynamic dimension from hlo transpose.3750@{}@0 to hlo %all-reduce.3755 = f32[<=70,256]{1,0} all-reduce(f32[<=70,256]{1,0} %transpose.3750), replica_groups={{0,1,2,3,4,5,6,7}}, to_apply=%sum.3751, metadata={op_type="CrossReplicaSum" op_name="CrossReplicaSum_33" source_file="dummy_file_name" source_line=10}, which is not implemented.

1013 tpu_program_group.cc:90] Check failed: xla_tpu_programs.size() > 0 (0 vs. 0) 

I have a very large code base, and am unable to reproduce. I was using AUC metric, but removed it as seen here #33890. The issue still persists.

Another relevant issue is #41590.

I have attached an MVCE example to reproduce, but it's still not enough to know that explicit shapes are needed as explained in #41590 since I don't know where I'm using transpose. I need a stack trace to point to where in the code it's causing the issue.

Standalone code to reproduce the issue

https://colab.research.google.com/drive/1bYuuwG0pFnIQe1X6jA7FJMLjZtvAmqlj?usp=sharing

Relevant log output

No response

[vivekjoshy]

What's even more weirder is that I'm defining the input sizes explicitly like so:

Config.COMPUTED_BATCH_SIZE = 128

with strategy.scope():
    model = my_model()
    input_shapes = [
        [Config.COMPUTED_BATCH_SIZE, 192], 
        [Config.COMPUTED_BATCH_SIZE, 192], 
        [COMPUTED_CHANNELS, 105, 129, 100], 
        [COMPUTED_CHANNELS, 105, 129, 100], 
        [COMPUTED_CHANNELS, 105, 129, 100], 
        [Config.COMPUTED_BATCH_SIZE, 70], 
        [Config.COMPUTED_BATCH_SIZE, 320]
    ]
    model.build(input_shape=input_shapes)

But the batch size shown in the error log is 256 in [70, 256] which I assume is the transposed tensor.

Ignore this since it appears that Hidden Size * 2 is also 256. So it could be a transpose inside a Dense layer.

Edit:

I've tracked it down to this code:

hidden_size = 128
self.descriptor_embedding = layers.Dense(
    hidden_size * 2, # 256
    activation='relu',
    input_shape=(Config.COMPUTED_BATCH_SIZE, 70)
)


learned_descriptors = tf.expand_dims(
    self.descriptor_embedding(descriptors),
    1
) # [BS, 1, HS * 2]

[vivekjoshy]

It seems to be an issue with tf.expand_dims(x, axis=1). Any axis other than 0 will cause this error. tf.transpose also causes this error.

[tilakrayal]

@vivekjoshy,
When I tried to execute the code by explicitly setting the size after you decode, using tf.reshape, the code was executed successfully. While with the other approach it was executed with the error.

tf.reshape(tf.image.decode_jpeg(image, channels = 3),[256,256, 3]), class_idx

Kindly find the gist of it here.

[vivekjoshy]

Fixed by passing in explicit shapes everywhere.

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No