diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-07-10 01:53:40 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-07-10 01:56:38 -0700 |
| commit | 7f85c95f71b01f711c366942a7cd911b0743b72c (patch) | |
| tree | ad92abbdf8a4c5736e323a75c3c153dadf3ce5c6 /tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc | |
| parent | c59bb780ebd1674ab34dd96d193c71698682ed4d (diff) | |
Implementation of arg_min.
PiperOrigin-RevId: 203908601
Diffstat (limited to 'tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc')
| -rw-r--r-- | tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc b/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc index 00ab7cbaa9..670bcf64e7 100644 --- a/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc +++ b/tensorflow/contrib/lite/toco/graph_transformations/propagate_array_data_types.cc @@ -100,6 +100,13 @@ bool PropagateArrayDataTypes::Run(Model* model, std::size_t op_index) { model->GetArray(op->outputs[0]).data_type = argmax_op->output_data_type; break; } + case OperatorType::kArgMin: { + // Data type of the ArgMin op is specified. + CHECK_EQ(op->outputs.size(), 1); + auto* argmin_op = static_cast<ArgMinOperator*>(op); + model->GetArray(op->outputs[0]).data_type = argmin_op->output_data_type; + break; + } case OperatorType::kRange: { auto* range_op = static_cast<RangeOperator*>(op); // Output type of the Range op can be set via an attribute |