TensorFlow in version 2.3 can produce misleading output summary

TensorFlow in version 2.3 can produce misleading output summary when initialized in a specific way.

Issue description

When producing summary using model.summary(), the summary can have misleading output. An example can be found in the original issue posted in the TensorFlow issue tracker:

import tensorflow as tf
from tensorflow import keras

layers = tf.keras.layers


class LayerTest(layers.Layer):
    def __init__(self):
        super(LayerTest, self).__init__()

    def call(self, inputs) -> tf.Tensor:
        predictions = inputs
        
        for k in predictions.keys():
            predictions[k] = tf.math.l2_normalize(predictions[k], axis=-1)

        for step_name in predictions.keys():
            loss = tf.reduce_mean(predictions[k])
            
        return loss
    
    
def Model(target_dim: int = 64):
    input_tensor = layers.Input(
        shape=[target_dim], name="input_tensor"
    )
    predictions = {'step_0': layers.Lambda(lambda _x: _x)(input_tensor), 'step_1': layers.Lambda(lambda _x: _x)(input_tensor)}
    logits = LayerTest()(predictions)
    return keras.Model(inputs=input_tensor, outputs=logits)

model = Model()
model.compile()
model.summary()

The output produced using TensorFlow 2.3:

__________________________________________________________________________________________________
Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_tensor (InputLayer)       [(None, 64)]         0                                            
__________________________________________________________________________________________________
lambda_1 (Lambda)               (None, 64)           0           input_tensor[0][0]               
__________________________________________________________________________________________________
tf_op_layer_layer_test/l2_norma [(None, 64)]         0           lambda_1[0][0]                   
__________________________________________________________________________________________________
tf_op_layer_layer_test/l2_norma [(None, 1)]          0           tf_op_layer_layer_test/l2_normali
__________________________________________________________________________________________________
tf_op_layer_layer_test/l2_norma [(None, 1)]          0           tf_op_layer_layer_test/l2_normali
__________________________________________________________________________________________________
tf_op_layer_layer_test/l2_norma [(None, 1)]          0           tf_op_layer_layer_test/l2_normali
__________________________________________________________________________________________________
tf_op_layer_layer_test/l2_norma [(None, 64)]         0           lambda_1[0][0]                   
                                                                 tf_op_layer_layer_test/l2_normali
__________________________________________________________________________________________________

Expected output (as produced by TensorFlow 2.2):

Layer (type)                    Output Shape         Param #     Connected to                     
==================================================================================================
input_tensor (InputLayer)       [(None, 64)]         0                                            
__________________________________________________________________________________________________
lambda (Lambda)                 (None, 64)           0           input_tensor[0][0]               
__________________________________________________________________________________________________
lambda_1 (Lambda)               (None, 64)           0           input_tensor[0][0]               
__________________________________________________________________________________________________
layer_test (LayerTest)          ()                   0           lambda[0][0]                     
                                                                 lambda_1[0][0]                   
=================================================================================================

The fix is to create a new dictionary:

import tensorflow as tf

layers = tf.keras.layers

class LayerTest(layers.Layer):
    def __init__(self):
        super(LayerTest, self).__init__()

    def call(self, inputs) -> tf.Tensor:
        predictions = inputs

        predictions2 = {}
        
        for k in predictions.keys():
            predictions2[k] = tf.math.l2_normalize(predictions[k], axis=-1)

        for step_name in predictions2.keys():
            loss = tf.reduce_mean(predictions2[k])
            
        return loss

See the original issue posted in the TensorFlow issue tracker for more info.

Affected packages:

Severity

Issue fix

Upgrade or downgrade TensorFlow used.

Pipeline units

Recommendation types

This message can be produced for all the recommendation types.

See this document that describes recommendation types.

Issues associated