


osd = nonwrapped_optim.state_dict()
         >>> rekeyed_osd = FSDP.rekey_optim_state_dict(osd, OptimStateKeyType.PARAM_NAME, nonwrapped_model)
         >>> wrapped_model, wrapped_optim = ...
         >>> sharded_osd = FSDP.shard_full_optim_state_dict(rekeyed_osd, wrapped_model)
         >>> wrapped_optim.load_state_dict(sharded_osd)



  Returns:
     The optimizer state dict re-keyed using the parameter keys
     specified by "optim_state_key_type".

  Return type:
     Dict[str, Any]

static scatter_full_optim_state_dict(full_optim_state_dict, model, optim_input=None, optim=None, group=None)
  Scatters the full optimizer state dict from rank 0 to all other
  ranks, returning the sharded optimizer state dict on each rank.
  The return value is the same as "shard_full_optim_state_dict()",
  and on rank 0, the first argument should be the return value of
  "full_optim_state_dict()".

  Example:
