class_docstr: Helper class to simplify split-cat pattern. In simple cases, both split and cat node can be removed in a "split->cat"
  pattern. However, there are various cases where they can't and we need to simplify split/ add transforms before cat.
  Some such cases are:
      1. Final node has additional args (not coming from the initial split)
      2. Shuffling of args between split/cat
      3. Some final nodes are non-(cat/stack)
      4. Split-dim != cat-dim (but equal split)
  
  Note that any combination of the above cases can happen.
  
  To deal with 1, 2, & 3 - we iterate over all users of split. And figure out common "ranges" that can be merged.
  Then, we simplify the split accordingly. In the best case, split can be entirely removed.
  
  To deal with 4, we add some transformations (unflatten + movedim) (See `get_transform_params`).
  
  Finally, depending on final node being cat or stack, unsqueeze/flatten needs to be added.
class_docstr_tok: ["Helper","class","to","simplify","split-cat","pattern",".","In","simple","cases",",","both","split","and","cat","node","can","be","removed","in","a","``","split-",">","cat","''","pattern",".","However",",","there","are","various","cases","where","they","ca","n't","and","we","need","to","simplify","split\/","add","transforms","before","cat",".","Some","such","cases","are",":","1",".","Final","node","has","additional","args","(","not","coming","from","the","initial","split",")","2",".","Shuffling","of","args","between","split\/cat","3",".","Some","final","nodes","are","non-","(","cat\/stack",")","4",".","Split-dim","!","=","cat-dim","(","but","equal","split",")","Note","that","any","combination","of","the","above","cases","can","happen",".","To","deal","with","1",",","2",",","&","3","-","we","iterate","over","all","users","of","split",".","And","figure","out","common","``","ranges","''","that","can","be","merged",".","Then",",","we","simplify","the","split","accordingly",".","In","the","best","case",",","split","can","be","entirely","removed",".","To","deal","with","4",",","we","add","some","transformations","(","unflatten","+","movedim",")","(","See","`","get_transform_params","`",")",".","Finally",",","depending","on","final","node","being","cat","or","stack",",","unsqueeze\/flatten","needs","to","be","added","."]
code: def replace_cat(
      self,
      graph: torch.fx.GraphModule,
      split_node: torch.fx.Node,
      next_users: List[torch.fx.Node],
      user_inputs_list_new,
      transform_params_list: List[List[_TransformParam]],
  ):
      split_dim = split_node.kwargs["dim"]
  
      split_users = split_node.users.keys()
      new_cats = []
      for user_node, user_inputs_new, transform_params in zip(
          next_users, user_inputs_list_new, transform_params_list
      ):
          if user_node.target not in {torch.cat, torch.stack}:
              # Change the args and kwargs of non-cat/stack nodes. Replace old getitems (belonging to
              # the original split node) with the newer getitems
              next_cat_input = 0
              for input_node in user_node.all_input_nodes:
                  if input_node in split_users:
                      user_node.replace_input_with(
                          input_node, user_inputs_new[next_cat_input]
                      )
                      next_cat_input += 1
              continue
  
          # Handle cat/stack user nodes
          cat_dim = get_arg_value(user_node, 1, "dim")
          user_inputs_new_transformed = []
          # For `unsqueeze` transform, we will combine consecutive inputs with the same unsqueeze params, and stack them
          to_stack = []
          stack_dim = None
          with graph.inserting_before(user_node):
              for user_input_new, transform_param in zip(
                  user_inputs_new, transform_params
              ):
                  # Apply transforms
                  (
                      unflatten_params,
                      movedim_params,
                      unsqueeze_params,
                      flatten_params,
                  ) = transform_param
                  if unsqueeze_params and (
                      stack_dim is None
                      or stack_dim == unsqueeze_params[0]
                  ):
                      to_stack.append(user_input_new)
                      stack_dim = unsqueeze_params[0]
                      continue
                  elif to_stack:
                      stacked_input = graph.call_function(
                          torch.stack,
                          args=(to_stack,),
                          kwargs={"dim": stack_dim},
                      )
                      to_stack = []
                      stack_dim = None
                      user_inputs_new_transformed.append(stacked_input)
                      if unsqueeze_params:
                          to_stack.append(user_input_new)
                          stack_dim = unsqueeze_params[0]
                          continue
  
                  if unflatten_params:
                      user_input_new = graph.call_function(
                          torch.unflatten,
                          args=(user_input_new, *unflatten_params),
                      )
                  if movedim_params:
                      user_input_new = graph.call_function(
                          torch.movedim,
                          args=(user_input_new, *movedim_params),
                      )
                  if flatten_params:
                      user_input_new = graph.call_function(
                          torch.flatten,
                          args=(user_input_new, *flatten_params),
                      )
                  user_inputs_new_transformed.append(user_input_new)
              if to_stack:
                  stacked_input = graph.call_function(
                      torch.stack,
                      args=(to_stack,),
                      kwargs={"dim": stack_dim},
                  )
                  user_inputs_new_transformed.append(stacked_input)
  
          with graph.inserting_after(user_node):
              if len(user_inputs_new_transformed) > 1:
                  new_cat_node = graph.call_function(
                      torch.cat,
                      args=(user_inputs_new_transformed,),
                      kwargs={"dim": cat_dim},
                  )
                  new_cat_node.meta.update(user_node.meta)
                  counters["inductor"]["scmerge_cat_added"] += 1
              else:
                  new_cat_node = user_inputs_new_transformed[-1]
  
          if (
              user_node.target == torch.cat
              and split_dim != cat_dim
              and split_node.target == torch.split
          ):
              with graph.inserting_after(new_cat_node):
                  new_cat_node = graph.call_function(
                      torch.flatten,
                      args=(new_cat_node, cat_dim, cat_dim + 1),
                  )
          user_node.replace_all_uses_with(new_cat_node)
          new_cats.append(new_cat_node)
  
imports: import itertools
  import logging
  import operator
  from typing import Any, Callable, List, Optional, Sequence, Set, Tuple, Union
  from typing_extensions import TypeAlias
  import torch
  from torch._dynamo.utils import counters
  from ..pattern_matcher import Arg, CallFunction, CallFunctionVarArgs, CallMethodVarArgs, config_flag, FailedMatch, get_arg_value, Ignored, KeywordArg, ListOf, Match, MatchContext, MULTIPLE, PatternExpr, register_graph_pattern, RepeatedExpr
  from .pre_grad import merge_getitem_cat_pass, merge_splits_pass, normalization_pass, split_cat_pass, unbind_stack_pass
label_desc: Use image node_id 11 for calling the SplitCatSimplifier obj's underlying member method code with example usage: obj.replace_cat(graph, split_node, next_users, user_inputs_list_new, transform_params_list) without return types