parameters: (sharded_states_original_dims, state_dict, mixed_precision_enabled)
docstring: Reshape model and optimizer states in the state_dict according to dimensions in sharded_states_original_dims
docstring_summary: Reshape model and optimizer states in the state_dict according to dimensions in sharded_states_original_dims
docstring_tokens: Reshape model and optimizer states in the state_dict according to dimensions in sharded_states_original_dims
function: def _reshape_states(sharded_states_original_dims, state_dict, mixed_precision_enabled):
      """Reshape model and optimizer states in the state_dict according to dimensions in sharded_states_original_dims"""
  
      model = _utils.state_dict_model_key()
      full_precision = _utils.state_dict_full_precision_key()
      optimizer = _utils.state_dict_optimizer_key()
      sharded_optimizer_keys = _utils.state_dict_sharded_optimizer_keys()
  
      for sharded_state_key, original_dim in sharded_states_original_dims.items():
          # reshape model states to original_dim only when mixed precision is enabled
          if mixed_precision_enabled and (model in state_dict):
              state_dict[model][full_precision][sharded_state_key] = \
                  state_dict[model][full_precision][sharded_state_key].reshape(original_dim)
  
          # reshape optimizer states to original_dim
          if optimizer in state_dict:
              for optimizer_key, optimizer_value in state_dict[optimizer][sharded_state_key].items():
                  if optimizer_key in sharded_optimizer_keys:
                      state_dict[optimizer][sharded_state_key][optimizer_key] = optimizer_value.reshape(original_dim)
function_tokens: def _reshape_states ( sharded_states_original_dims , state_dict , mixed_precision_enabled ) : model = _utils . state_dict_model_key ( ) full_precision = _utils . state_dict_full_precision_key ( ) optimizer = _utils . state_dict_optimizer_key ( ) sharded_optimizer_keys = _utils . state_dict_sharded_optimizer_keys ( ) for sharded_state_key , original_dim in sharded_states_original_dims . items ( ) : # reshape model states to original_dim only when mixed precision is enabled if mixed_precision_enabled and ( model in state_dict ) : state_dict [ model ] [ full_precision ] [ sharded_state_key ] = state_dict [ model ] [ full_precision ] [ sharded_state_key ] . reshape ( original_dim ) # reshape optimizer states to original_dim if optimizer in state_dict : for optimizer_key , optimizer_value in state_dict [ optimizer ] [ sharded_state_key ] . items ( ) : if optimizer_key in sharded_optimizer_keys : state_dict [ optimizer ] [ sharded_state_key ] [ optimizer_key ] = optimizer_value . reshape ( original_dim )