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