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

fix bug again

pull/2364/head
oahzxl 2023-01-02 00:27:11 +08:00
parent 7fd3b45af2
commit 9c5e028a62
1 changed files with 19 additions and 16 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

35
chunk_codegen.py
View File

@ -2003,22 +2003,25 @@ def emit_code_with_chunk(
) )
# ones like # ones like
if "ones_like" in node.name: if "ones_like" in node.name:
chunk_dim = chunk_search[region_idx]["node_chunk_dim"][ meta_node = chunk_region_search.index_tracer.node_list[node_idx]
chunk_region_search.index_tracer.node_list[node_idx] chunk_dim = chunk_search[region_idx]["node_chunk_dim"][meta_node][
]["chunk_dim"] "chunk_dim"
if ( ]
_get_node_shape( if _get_node_shape(meta_node)[chunk_dim] != 1:
chunk_region_search.index_tracer.node_list[node_idx] source_node = meta_node.args[0].args[0]
)[chunk_dim] if (
== 1 source_node not in chunk_search[region_idx]["node_chunk_dim"]
): or chunk_search[region_idx]["node_chunk_dim"][source_node][
continue "chunk_dim"
chunk_slice = _gen_chunk_slice_dim( ]
chunk_dim, "chunk_idx", _get_node_shape(node) is None
) ):
body[-1] = _replace_name( chunk_slice = _gen_chunk_slice_dim(
body[-1], node.args[0].name, node.args[0].name + chunk_slice chunk_dim, "chunk_idx", _get_node_shape(node)
) )
body[-1] = _replace_name(
body[-1], node.args[0].name, node.args[0].name + chunk_slice
)
body[-1] = _replace_reshape_size( body[-1] = _replace_reshape_size(
body[-1], node.name, chunk_search[region_idx]["reshape_size"] body[-1], node.name, chunk_search[region_idx]["reshape_size"]
) )