MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / MakeReductionAxes

Function MakeReductionAxes

tensorflow/core/graph/quantize_training.cc:333–364  ·  view source on GitHub ↗

Sets output to the Node that computes reduction axes corresponding to all dimensions of input and return.

Source from the content-addressed store, hash-verified

331// Sets output to the Node that computes reduction axes corresponding to all
332// dimensions of input and return.
333Status MakeReductionAxes(Graph* graph, string name_prefix, Node* input,
334 Node** output) {
335 name_prefix = strings::StrCat(name_prefix, "/ReductionAxes");
336 Node* start;
337 Tensor zero_tensor(DT_INT32, TensorShape());
338 zero_tensor.flat<int32>()(0) = 0;
339 TF_RETURN_IF_ERROR(
340 NodeBuilder(strings::StrCat(name_prefix, "/RangeStart"), "Const")
341 .Attr("dtype", DT_INT32)
342 .Attr("value", zero_tensor)
343 .Finalize(graph, &start));
344 Node* delta;
345 Tensor one_tensor(DT_INT32, TensorShape());
346 one_tensor.flat<int32>()(0) = 1;
347 TF_RETURN_IF_ERROR(
348 NodeBuilder(strings::StrCat(name_prefix, "/RangeDelta"), "Const")
349 .Attr("dtype", DT_INT32)
350 .Attr("value", one_tensor)
351 .Finalize(graph, &delta));
352 Node* rank;
353 TF_RETURN_IF_ERROR(
354 NodeBuilder(strings::StrCat(name_prefix, "/InputRank"), "Rank")
355 .Input(input)
356 .Finalize(graph, &rank));
357 TF_RETURN_IF_ERROR(
358 NodeBuilder(strings::StrCat(name_prefix, "/ReductionAxes"), "Range")
359 .Input(start)
360 .Input(rank)
361 .Input(delta)
362 .Finalize(graph, output));
363 return Status::OK();
364}
365
366// Computes the exponential moving average of input, updated in update_variable.
367Status MakeExponentialMovingAverage(Graph* graph, string name_prefix,

Callers 1

MakeEMAMinMaxVarsFunction · 0.85

Calls 6

NodeBuilderClass · 0.70
StrCatFunction · 0.50
TensorShapeClass · 0.50
FinalizeMethod · 0.45
AttrMethod · 0.45
InputMethod · 0.45

Tested by

no test coverage detected