Skip to content

Commit

Permalink
move IA declaration inside NodeDef
Browse files Browse the repository at this point in the history
  • Loading branch information
prathikr committed Mar 13, 2024
1 parent f9c961f commit 7840801
Showing 1 changed file with 3 additions and 4 deletions.
7 changes: 3 additions & 4 deletions orttraining/orttraining/core/graph/gradient_builder.cc
Original file line number Diff line number Diff line change
Expand Up @@ -1113,18 +1113,17 @@ IMPLEMENT_GRADIENT_BUILDER(GetReduceMeanGradient) {
ArgDef grad = GO(0);
if (!keepdims) {
size_t numInputs = GetSrcNodeInputSize();
grad = IA("Unqueezed_Grad");
if (attributes.find("axes") != attributes.end()) {
std::vector<int64_t> axes_values = RetrieveValues<int64_t>(attributes.at("axes"));
if (SrcNodeOpsetVersion() < 13) { // axes is attribute for unsqueeze
result.push_back(NodeDef("Unsqueeze", {GO(0)}, {grad}, {MakeAttribute("axes", axes_values)}));
result.push_back(NodeDef("Unsqueeze", {GO(0)}, {IA("Unqueezed_Grad")}, {MakeAttribute("axes", axes_values)}));
} else {
NodeDef axes_values_node = ConstantVectorNode(axes_values, Name("axes_values"));
result.push_back(axes_values_node);
result.push_back(NodeDef(OpDef{"Unsqueeze", kOnnxDomain, 13}, {GO(0), axes_values_node.output_args[0]}, {grad}));
result.push_back(NodeDef(OpDef{"Unsqueeze", kOnnxDomain, 13}, {GO(0), axes_values_node.output_args[0]}, {IA("Unqueezed_Grad")}));

Check warning on line 1123 in orttraining/orttraining/core/graph/gradient_builder.cc

View workflow job for this annotation

GitHub Actions / Lint C++

[cpplint] reported by reviewdog 🐶 Lines should be <= 120 characters long [whitespace/line_length] [2] Raw Output: orttraining/orttraining/core/graph/gradient_builder.cc:1123: Lines should be <= 120 characters long [whitespace/line_length] [2]
}
} else if (numInputs == 2) { // optional input 'axes' is available as input I(1)
result.push_back(NodeDef("Unsqueeze", {GO(0), I(1)}, {grad}));
result.push_back(NodeDef("Unsqueeze", {GO(0), I(1)}, {IA("Unqueezed_Grad")}));
}
}

Expand Down

0 comments on commit 7840801

Please sign in to comment.