keras/src/tree/optree_impl.py/map_structure
def map_structure(func, *structures):
    if not callable(func):
        raise TypeError(f"`func` must be callable. Received: func={func}")
    if not structures:
        raise ValueError("Must provide at least one structure")
    for other in structures[1:]:
        assert_same_structure(structures[0], other, check_types=False)
    return optree.tree_map(
        func, *structures, none_is_leaf=True, namespace="keras"
    )