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