venv/lib/python3.12/site-packages/functorch/experimental/control_flow.py

from torch import cond  # noqa: F401
from torch._higher_order_ops.map import (  # noqa: F401
    _stack_pytree,
    _unstack_pytree,
    map,
)