sub-units). It then performs all necessary processing and
           returns "None"; 2) "hook: Callable[torch.Tensor,
           torch.Tensor] -> None": This function takes in two Python
           tensors, the first one represents the full, flattened,
           unsharded gradient with respect to all variables
           corresponding to the model this FSDP unit is wrapping (that
           are not wrapped by other FSDP sub-units). The latter
           represents a pre-sized tensor to store a chunk of a sharded
           gradient after reduction. In both cases, callable performs
           all necessary processing and returns "None". Callables with
           signature 1 are expected to handle gradient communication
           for a NO_SHARD case. Callables with signature 2 are
           expected to handle gradient communication for sharded
           cases.
static rekey_optim_state_dict(optim_state_dict, optim_state_key_type, model, optim_input=None, optim=None)