llvm::StringMap<Value> grads_map;
    for (const string& source : grad_it->getSecond()) {
      auto g = func.front().addArgument(t, func.getLoc());
      (*stats)[g].accumulate_on_write = true;
      grads_map[source] = g;
    }
    auto& stat = (*stats)[arg];
    stat.accumulate_on_write = ta_accumulate_on_write(argnum);
    stat.grads = std::move(grads_map);
  }
  UpdateFuncType(func);
}