Implementation of ctc beam search decoder op in custom op fashion.

PiperOrigin-RevId: 207210333

1 files changed, 23 insertions, 0 deletions

diff --git a/tensorflow/contrib/lite/toco/model.h b/tensorflow/contrib/lite/toco/model.h
index 7d0dbfcc05..cd263930f5 100644
--- a/tensorflow/contrib/lite/toco/model.h
+++ b/tensorflow/contrib/lite/toco/model.h
@@ -148,6 +148,7 @@ enum class OperatorType : uint8 {
   kLogicalAnd,
   kLogicalNot,
   kLogicalOr,
+  kCTCBeamSearchDecoder,
 };
 
 // Helper to deal with TensorFlow arrays using a different ordering of
@@ -438,6 +439,28 @@ struct ConvOperator : Operator {
   int dilation_height_factor = 1;
 };
 
+// CTCBeamSearchDecoder operator:
+//
+// Inputs:
+//   inputs[0]: required: the logits.
+//   inputs[1]: required: sequence length.
+//   inputs[2]: optional: beam width.
+//   inputs[3]: optional: top paths.
+//   inputs[4]: optional: merge repeated.
+//
+//  Outputs:
+//    outputs[0]: deocoded.
+//    outputs[1]: log probability.
+//
+// TensorFlow equivalent: CTCBeamSearchDecoder
+struct CTCBeamSearchDecoderOperator : Operator {
+  CTCBeamSearchDecoderOperator()
+      : Operator(OperatorType::kCTCBeamSearchDecoder) {}
+  int beam_width;
+  int top_paths;
+  bool merge_repeated = true;
+};
+
 // Depthwise-separable convolution operator.
 //
 // Inputs: