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 2023-01-09 16:30:16 +08:00
parent 212b5b1b5f
commit 1951f7fa87
1 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
colossalai/autochunk/autochunk_codegen.py
View File

@ -137,9 +137,9 @@ def emit_code_with_chunk(
):
"""
Emit code with chunk according to chunk_infos.
It will generate a for loop in chunk regions, and replace inputs
and outputs of regions with chunked variables.
It will generate a for loop in chunk regions, and
replace inputs and outputs of regions with chunked variables.
Args:
body: forward code
@ -156,9 +156,11 @@ def emit_code_with_chunk(
chunk_ends = [i["region"][1] for i in chunk_infos]
# chunk inputs
chunk_inputs = [i["inputs"] for i in chunk_infos] # input with chunk
chunk_inputs_non_chunk = [i["inputs_non_chunk"] for i in chunk_infos] # input without chunk
chunk_inputs_dim = [i["inputs_dim"] for i in chunk_infos] # input chunk dim
chunk_inputs = [i["inputs"] for i in chunk_infos] # input with chunk
chunk_inputs_non_chunk = [
i["inputs_non_chunk"] for i in chunk_infos
] # input without chunk
chunk_inputs_dim = [i["inputs_dim"] for i in chunk_infos] # input chunk dim
chunk_inputs_names = [j.name for i in chunk_inputs for j in i] + [
j.name for i in chunk_inputs_non_chunk for j in i
]