tf.keras.tree.map_structure_up_to

Maps func through given structures up to shallow_structure.

tf.keras.tree.map_structure_up_to(
    shallow_structure, func, *structures
)

This is a variant of map_structure which only maps the given structures up to shallow_structure. All further nested components are retained as-is.

Examples:

shallow_structure = [None, None]
structure = [[1, 1], [2, 2]]
keras.tree.map_structure_up_to(shallow_structure, len, structure)
[2, 2]

shallow_structure = [None, [None, None]]
keras.tree.map_structure_up_to(shallow_structure, str, structure)
['[1, 1]', ['2', '2']]

shallow_structure A structure with layout common to all structures.
func A callable that accepts as many arguments as there are structures.
*structures Arbitrarily nested structures of the same layout.

A new structure with the same layout as shallow_structure.