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

19
chunk_codegen.py
View File

@ -2003,16 +2003,19 @@ 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(meta_node)[chunk_dim] != 1:
source_node = meta_node.args[0].args[0]
if ( if (
_get_node_shape( source_node not in chunk_search[region_idx]["node_chunk_dim"]
chunk_region_search.index_tracer.node_list[node_idx] or chunk_search[region_idx]["node_chunk_dim"][source_node][
)[chunk_dim] "chunk_dim"
== 1 ]
is None
): ):
continue
chunk_slice = _gen_chunk_slice_dim( chunk_slice = _gen_chunk_slice_dim(
chunk_dim, "chunk_idx", _get_node_shape(node) chunk_dim, "chunk_idx", _get_node_shape(node)
) )