github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

code style

pull/2364/head
oahzxl 2022-12-23 14:26:43 +08:00
parent d309e9338b
commit 49ba619085
1 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
chunk_codegen.py
View File

@ -1899,7 +1899,7 @@ def emit_code_with_chunk(
chunk_outputs = [i["outputs"][0] for i in chunk_search] chunk_outputs = [i["outputs"][0] for i in chunk_search]
chunk_outputs_dim = [i["outputs_dim"] for i in chunk_search] chunk_outputs_dim = [i["outputs_dim"] for i in chunk_search]
chunk_prepose_nodes = [i["args"]["prepose_nodes"] for i in chunk_search] chunk_prepose_nodes = [i["args"]["prepose_nodes"] for i in chunk_search]
node_idx = 0 node_idx = 0
@ -1933,7 +1933,9 @@ def emit_code_with_chunk(
emit_node_func(node, body) emit_node_func(node, body)
# replace input var with chunk var # replace input var with chunk var
for input_node_idx, input_node in enumerate(chunk_inputs[region_idx]): for input_node_idx, input_node in enumerate(chunk_inputs[region_idx]):
for idx, dim in chunk_inputs_dim[region_idx][input_node_idx].items(): for idx, dim in chunk_inputs_dim[region_idx][
input_node_idx
].items():
if idx == node_idx: if idx == node_idx:
chunk_slice = _gen_chunk_slice_dim( chunk_slice = _gen_chunk_slice_dim(
dim, "chunk_idx", _get_node_shape(input_node) dim, "chunk_idx", _get_node_shape(input_node)