[reland] Update tls logic to work better with guarded call (#73925)
This PR relands https://github.com/pytorch/pytorch/pull/73925 which we
reverted due to a large breakage in functorch.
As a part of the reland, this PR adds a change we agreed upon in
https://docs.google.com/document/d/1i7Y9VZp9PxtgVcrQh6nGQXkXkPc1uMep0dM-OMOGJ9o/edit
The change is moving the PythonTLSSnapshot key after
DynamicLayerFrontMode.
Test Plan:
- I tested this with an updated version of functorch and all the tests
pass so I think we are out of the woods.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/74577
Approved by: https://github.com/albanD
