github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-05 11:02:05 +00:00
Code Issues Packages Projects Releases Wiki Activity

rename function from index to indice

Browse Source
This commit is contained in:
oahzxl
2023-01-09 17:34:30 +08:00
parent 0ea903b94e
commit cb9817f75d
4 changed files with 91 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
colossalai/autochunk/reorder_graph.py
View File

@@ -6,7 +6,7 @@ class ReorderGraph(object):
def __init__(self, trace_indice: TraceIndice) -> None:
self.trace_indice = trace_indice
self.all_reorder_map = {
i: i for i in range(len(self.trace_indice.idx_trace_list))
i: i for i in range(len(self.trace_indice.indice_trace_list))
}
def _get_reorder_map(self, chunk_info):
@@ -60,18 +60,18 @@ class ReorderGraph(object):
def _reorder_idx_trace(self, reorder_map):
# reorder list
new_idx_trace_list = [None for _ in range(len(self.trace_indice.idx_trace_list))]
new_idx_trace_list = [None for _ in range(len(self.trace_indice.indice_trace_list))]
for old_idx, new_idx in reorder_map.items():
new_idx_trace_list[new_idx] = self.trace_indice.idx_trace_list[old_idx]
self.trace_indice.idx_trace_list = new_idx_trace_list
new_idx_trace_list[new_idx] = self.trace_indice.indice_trace_list[old_idx]
self.trace_indice.indice_trace_list = new_idx_trace_list
# update compute
for idx_trace in self.trace_indice.idx_trace_list:
for idx_trace in self.trace_indice.indice_trace_list:
compute = idx_trace["compute"]
for dim_compute in compute:
for idx, i in enumerate(dim_compute):
dim_compute[idx] = reorder_map[i]
# update source
for idx_trace in self.trace_indice.idx_trace_list:
for idx_trace in self.trace_indice.indice_trace_list:
source = idx_trace["source"]
for dim_idx, dim_source in enumerate(source):
new_dim_source = {}