for (auto& [op_key, cost] : expected_op_cost_map) {
    cost = rand_r(&seed) % 1000;
  }
  tensorflow::tfrt_stub::CostRecorder cost_recorder;
  for (const auto& [op_key, cost] : expected_op_cost_map) {
    cost_recorder.RecordCost(op_key, cost);
  }

  // Update the TFRT MLIR with the cost recorder.
  tfrt_compiler::UpdateOpCostInTfrtMlir(module.get(), cost_recorder);

  // Check the updated costs.