docstr: Some shape args in reshape may have changed due to chunk
  reassign those changed shape
docstr_tok: ["Some","shape","args","in","reshape","may","have","changed","due","to","chunk","reassign","those","changed","shape"]
code: def _reassign_reshape_size(self, chunk_info):
      """
      Some shape args in reshape may have changed due to chunk
      reassign those changed shape
      """
      chunk_region = chunk_info["region"]
      reshape_size = {}
      chunk_shape = get_node_shape(chunk_info["outputs"][0])[
          chunk_info["outputs_dim"][0]
      ]
      for node in self.node_mgr.get_node_slice_by_idx(
          chunk_region[0], chunk_region[1] + 1
      ):
          if any(i == get_node_name(node) for i in ["reshape", "view"]):
              if node in chunk_info["args"]["prepose_nodes"]:
                  continue
              if node.args[0] in chunk_info["inputs_non_chunk"]:
                  continue
              reshape_args = flat_list(node.args[1:])
              if (
                  len(reshape_args) == 1
                  and get_node_shape(reshape_args[0]) is None
                  and len(reshape_args[0].meta["fwd_out"]) > 1
              ):
                  continue
              chunk_dim = chunk_info["node_chunk_dim"][node][
                  "chunk_dim"
              ]
              new_shape = ""
              for reshape_arg_dim, reshape_arg in enumerate(
                  reshape_args
              ):
                  if reshape_arg_dim == chunk_dim:
                      new_shape += (
                          "min(chunk_size, %d - chunk_idx), "
                          % chunk_shape
                      )
                  else:
                      if isinstance(reshape_arg, int):
                          new_shape += "%s, " % str(reshape_arg)
                      else:
                          new_shape += "%s, " % reshape_arg.name
              new_shape = new_shape[:-2]
              origin_shape = str(reshape_args)[1:-1]
              reshape_size[node.name] = [origin_shape, new_shape]
      chunk_info["reshape_size"] = reshape_size
      return chunk_info
  
file_path: colossalai/colossalai/autochunk/trace_flow.py
imports: from typing import Dict, List, Tuple
  from torch.fx.node import Node
  from .trace_indice import TraceIndice
  from .utils import NodeMgr, find_chunk_all_input_nodes, find_chunk_compute_input_and_output_nodes, find_tensor_shape_node, flat_list, get_node_name, get_node_shape, is_non_compute_node
label_desc: Use image node_id 12 for calling the TraceFlow obj's underlying member method code with example usage: obj._reassign_reshape_size(chunk_info) and returns: chunk_info