def METHOD_NAME(x, provenance: frozenset):
    # avoid max-recursion depth error for torch<=2.0
    flat_args, _ = tree_flatten(x)
    if not flat_args or flat_args[0] is x:
        return x
    return tree_map(partial(track_provenance, provenance=provenance), x)