tf.keras.tree.map_shape_structure

Variant of keras.tree.map_structure that operates on shape tuples.