}
    } else {
      funcs_to_trim.push_back(func);
    }
  }

  // remove all unexported functions from the module.
  for (auto func : funcs_to_trim) {
    func.erase();
  }
  return true;
}

// validate that all reachable functions from the remaining functions are
// also in the allowlist.
void TrimFunctionsPass::Verify() {

      VLOG(2) << "Deleting node " << node->DebugString();
      // Ensure that we do not attempt to add control edges to nodes that are
      // deleted.
      control_inputs.erase(node);
      control_outputs.erase(node);
      graph->RemoveNode(node);
    }

    TF_ASSIGN_OR_RETURN(Node * xla_launch, graph->AddNode(def));
    for (int i = 0, end = data_inputs.size(); i < end; ++i) {

    }

    // Copy all the unregisted attributes to the new op.
    if (failed(CopyAllowedUnregisteredAttrs(op, new_op, defined_attrs))) {
      return WalkResult::interrupt();
    }

    op->erase();
    changed |= true;
    return WalkResult::advance();
  });

  // If `changed` is false, it is considered as a failure, so the recursive
  // rewrite will stop.
  return success(changed);
}

                print("stage 2")
            """.replaceIndent()
        )

        val fragment = source.fragment(0..6, 8..29)
        val stage1 = Program.Plugins(fragment)
        val stage2 = source.map { it.erase(listOf(fragment.range)) }
        val stagedProgram = Dynamic(
            Static(ApplyPluginRequestsOf(stage1), ApplyBasePlugins),
            stage2
        )

        assertStagedTopLevelProjectProgram(

  // Remove unused message identification ops.
  if (communication_key_index == 0)
    for (auto op : compilation_key_ops) op->erase();
  if (communication_key_index == 0 || core_to_device_ordinal.empty())
    for (auto op : core_to_device_ordinal_op) op->erase();

  for (mlir::tf_device::LaunchOp tmp_host_launch_op : core_to_tmp_launch)
    RemoveOutsideCompilation(tmp_host_launch_op);

  absl::flat_hash_map<std::string, SignatureDef> signature_def_map(
      signatures.begin(), signatures.end());

  // Init op is only used during initialization and it's not a target for
  // quantization.
  signature_def_map.erase(kSavedModelInitOpSignatureKey);
  return signature_def_map;
}

// Retrieves the function name -> function alias mapping from the
// `SavedModelBundle`.
// TODO: b/314124142 - Remove the need for this function.

      // Build TFL custom op, replace composite with custom op.
      TFL::CustomOp tfl_custom_op =
          BuildCustomOp(composite, custom_option_buffer);
      composite->replaceAllUsesWith(tfl_custom_op);
      composite->erase();
    });
  }
};

static PassRegistration<LegalizeCompositeToCustomOpPass> pass;

}  // namespace odml

  for (auto result :
       llvm::zip(results_to_remap, new_parallel_execute.getResults()))
    std::get<0>(result).replaceAllUsesWith(std::get<1>(result));

  parallel_execute.erase();
}

void TPUParallelExecuteSinkResourceWrite::runOnOperation() {
  llvm::SmallVector<tf_device::ParallelExecuteOp, 4> parallel_executes;
  getOperation().walk([&](tf_device::ParallelExecuteOp parallel_execute) {

                                builder.getContext(), kSessionInitFuncName)));
  if (session_init_op) {
    session_init_op->replaceAllUsesWith(new_session_init_op);
    session_init_op->erase();
  }
  return func;
}

func::FuncOp GetOrCreateSessionInitFunc(ModuleOp module) {
  SessionInitializerOp session_init_op = GetSessionInitializerOp(module);

        op->getLoc(), source_op.getToken(), /*input=*/op->getOperand(0),
        /*controlInputs=*/new_operands);
    op->replaceAllUsesWith(
        ValueRange({source_op.getOutput(), source_op.getControl()}));
    op->erase();
  });
}

class ConvertGraphOp : public OpConversionPattern<tfg::GraphOp> {
 public:
  using OpConversionPattern::OpConversionPattern;

  LogicalResult matchAndRewrite(