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