[jit] preserve keys on dictionary input tracing (#40792)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/40792
Fixes https://github.com/pytorch/pytorch/issues/40529.
One followup should be to produce a better error message when a new
dictionary has different keys than the traced input. Right now it
presents as a fairly opaque `KeyError`.
Test Plan: Imported from OSS
Differential Revision: D22311731
Pulled By: suo
fbshipit-source-id: c9fbe0b54cf69daed2f11a191d988568521a3932
