code: def gather_optimization_pass(program, weight_name):
      op_idxs = []
      gather_grad_op = None
      momentum_op = None
      for idx, op in enumerate(program.global_block().ops):
          if (
              op.type == "gather_grad" or op.type == "momentum"
          ) and check_contains(weight_name, op.input_arg_names):
              op_idxs.append(idx)
              if op.type == "momentum":
                  momentum_op = op
              if op.type == "gather_grad":
                  gather_grad_op = op
  
      if gather_grad_op is not None and momentum_op is not None:
          inputs = {
              "Param": momentum_op.input("Param"),
              "Velocity": momentum_op.input("Velocity"),
              "LearningRate": momentum_op.input("LearningRate"),
              "Grad": gather_grad_op.input("Out@GRAD"),
              "Index": gather_grad_op.input("Index"),
              "Axis": gather_grad_op.input("Axis"),
          }
          outputs = {
              "ParamOut": momentum_op.output("ParamOut"),
              "VelocityOut": momentum_op.output("VelocityOut"),
          }
          if (
              "MasterParam" in momentum_op.input_names
              and len(momentum_op.input("MasterParam")) > 0
          ):
              inputs["MasterParam"] = momentum_op.input("MasterParam")
          if (
              "MasterParamOut" in momentum_op.output_names
              and len(momentum_op.output("MasterParamOut")) > 0
          ):
              outputs["MasterParamOut"] = momentum_op.output(
                  "MasterParamOut"
              )
  
          attrs = {
              "mu": momentum_op.attr("mu"),
              "use_nesterov": momentum_op.attr("use_nesterov"),
              "regularization_method": momentum_op.attr(
                  "regularization_method"
              ),
              "regularization_coeff": momentum_op.attr(
                  "regularization_coeff"
              ),
              "multi_precision": momentum_op.attr("multi_precision"),
              "rescale_grad": momentum_op.attr("rescale_grad"),
              "op_device": momentum_op.attr("op_device"),
              "op_namescope": momentum_op.attr("op_namescope"),
              "op_role": momentum_op.attr("op_role"),
              "op_role_var": momentum_op.input("Param"),
              "axis": gather_grad_op.attr("axis"),
          }
          program.global_block()._insert_op(
              op_idxs[-1] + 1,
              type="sparse_momentum",
              inputs=inputs,
              outputs=outputs,
              attrs=attrs,
          )
  
          for idx in reversed(op_idxs):
              program.global_block()._remove_op(idx, sync=False)
  
          var_names = []
          for idx, name in enumerate(program.global_block().vars):
              if "@GRAD" in name and weight_name in name:
                  var_names.append(name)
          for name in var_names:
              program.global_block()._remove_var(name, sync=False)
          program.global_block()._sync_with_cpp()
  
label_desc: Use image node_id 2 for calling a global function with example usage: gather_optimization_pass(program, weight_name) without return types