parameters: (train_map, predict_map=None)
return_statement: return ret
docstring: infer checkpoint slices to be loaded
docstring_summary: infer checkpoint slices to be loaded
docstring_tokens: infer checkpoint slices to be loaded
function: def _infer_rank_list(train_map, predict_map=None):
      """infer checkpoint slices to be loaded"""
      ret = {}
      if _get_pipeline_stages() > 1:
          local_rank = int(_get_global_rank() % (_get_device_num() / _get_pipeline_stages()))
      else:
          local_rank = _get_global_rank()
      for param_name in train_map:
          train_layout = train_map[param_name]
          train_dev_mat = train_layout[0]
          dev_num = np.array(train_dev_mat).prod()
          new_train_layout = _remove_repeated_slices(train_layout)
          array = np.arange(dev_num).reshape(train_dev_mat)
          index = ()
          for i in new_train_layout[0]:
              if i == 1:
                  index = index + (0,)
              else:
                  index = index + (slice(None),)
          rank_list = array[index].flatten()
          if not predict_map:
              ret[param_name] = (rank_list, False)
              continue
          if param_name not in predict_map:
              logger.warning("predict_map does not contain %s", param_name)
              continue
          predict_layout = predict_map[param_name]
          dev_num = np.array(predict_layout[0]).prod()
          # optimization pass
          if _check_same_layout(train_layout, predict_layout):
              ret[param_name] = ([local_rank], True)
              continue
          if _check_similar_layout(train_layout, predict_layout):
              if len(rank_list) == 1:
                  ret[param_name] = (rank_list, True)
              elif len(rank_list) == dev_num:
                  ret[param_name] = ([rank_list[local_rank]], True)
              else:
                  ret[param_name] = (rank_list, False)
          else:
              ret[param_name] = (rank_list, False)
      return ret
function_tokens: def _infer_rank_list ( train_map , predict_map = None ) : ret = { } if _get_pipeline_stages ( ) > 1 : local_rank = int ( _get_global_rank ( ) % ( _get_device_num ( ) / _get_pipeline_stages ( ) ) ) else : local_rank = _get_global_rank ( ) for param_name in train_map : train_layout = train_map [ param_name ] train_dev_mat = train_layout [ 0 ] dev_num = np . array ( train_dev_mat ) . prod ( ) new_train_layout = _remove_repeated_slices ( train_layout ) array = np . arange ( dev_num ) . reshape ( train_dev_mat ) index = ( ) for i in new_train_layout [ 0 ] : if i == 1 : index = index + ( 0 , ) else : index = index + ( slice ( None ) , ) rank_list = array [ index ] . flatten ( ) if not predict_map : ret [ param_name ] = ( rank_list , False ) continue if param_name not in predict_map : logger . warning ( "predict_map does not contain %s" , param_name ) continue predict_layout = predict_map [ param_name ] dev_num = np . array ( predict_layout [ 0 ] ) . prod ( ) # optimization pass if _check_same_layout ( train_layout , predict_layout ) : ret [ param_name ] = ( [ local_rank ] , True ) continue if _check_similar_layout ( train_layout , predict_layout ) : if len ( rank_list ) == 1 : ret [ param_name ] = ( rank_list , True ) elif len ( rank_list ) == dev_num : ret [ param_name ] = ( [ rank_list [ local_rank ] ] , True ) else : ret [ param_name ] = ( rank_list , False ) else : ret [ param_name ] = ( rank_list , False ) return ret