Re-keys the optimizer state dict "optim_state_dict" to use the
      key type "optim_state_key_type". This can be used to achieve
      compatibility between optimizer state dicts from models with
      FSDP instances and ones without.
  To re-key an FSDP full optimizer state dict (i.e. from
  "full_optim_state_dict()") to use parameter IDs and be loadable
  to a non-wrapped model:

     >>> wrapped_model, wrapped_optim = ...
     >>> full_osd = FSDP.full_optim_state_dict(wrapped_model, wrapped_optim)
     >>> nonwrapped_model, nonwrapped_optim = ...
     >>> rekeyed_osd = FSDP.rekey_optim_state_dict(full_osd, OptimStateKeyType.PARAM_ID, nonwrapped_model)
     >>> nonwrapped_optim.load_state_dict(rekeyed_osd)

  To re-key a normal optimizer state dict from a non-wrapped model
  to be loadable to a wrapped model:

     >>> nonwrapped_model, nonwrapped_optim = ...
     >>> osd = nonwrapped_optim.state_dict()
