TPU v3-8 CrossReplicaSum_33 Error about tensorflow

[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