keras/src/tree/optree_impl.py/assert_same_structure
def assert_same_structure(a, b, check_types=True):
    a_structure = optree.tree_structure(a, none_is_leaf=True, namespace="keras")
    b_structure = optree.tree_structure(b, none_is_leaf=True, namespace="keras")
    if a_structure != b_structure:
        raise ValueError(
            "`a` and `b` don't have the same structure. "
            f"Received: structure of a={a_structure}, "
            f"structure of b={b_structure}"
        )
    if check_types:
        type_structure = optree.tree_map(
            lambda x, y: type(x) is type(y),
            a,
            b,
            none_is_leaf=True,
            namespace="keras",
        )
        if not optree.tree_all(
            type_structure, none_is_leaf=True, namespace="keras"
        ):
            raise TypeError(
                "The type of the leaves of `a` and `b` doesn't match."
            )