Save and restore TLS has to also reset torch_function disabled state

[ezyang]

https://github.com/pytorch/pytorch/pull/73942 is blocked because torch function gets disabled in the default implementation and then if you fall through to Python key it is still disabled, which breaks FX tracing. So we need to somehow reenable torch function by the time we get to python key (it can't be saved via the snapshot TLS key because by the time we get to dispatcher it's already disabled). Alban suggests just unconditionally turning it back on when we go Python key.