[HigherOrderOp] Flatten outputs of `wrap`. (#109433)
Fix: #109247
This PR flattens `wrap` outputs by inlining `pytree.tree_flatten` function after calling
the inner function.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/109433
Approved by: https://github.com/zou3519
ghstack dependencies: #110290
