Skip to content

GitLab

Commit 9c5e028a authored by oahzxl's avatar oahzxl
Browse files

fix bug again

parent 7fd3b45a
Hide whitespace changes
Inline Side-by-side
Showing with 19 additions and 16 deletions
chunk_codegen.py
View file @ 9c5e028a
...@@ -2003,22 +2003,25 @@ def emit_code_with_chunk( ...@@ -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"]
) )
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment