keras/src/tree/dmtree_impl.py/map_shape_structure
def map_shape_structure(func, structure):
    if is_shape_tuple(structure):
        return func(tuple(structure))
    if isinstance(structure, list):
        return [map_shape_structure(func, e) for e in structure]
    if isinstance(structure, tuple):
        return tuple(map_shape_structure(func, e) for e in structure)
    if isinstance(structure, dict):
        return {k: map_shape_structure(func, v) for k, v in structure.items()}
    else:
        raise ValueError(f"Cannot map function to unknown object {structure}")