Add support for RNN state array of type tf.identity.

PiperOrigin-RevId: 192689747

diff --git a/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc b/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc
index aa2c293..8e6aaf5 100644 (file)
--- a/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc
+++ b/tensorflow/contrib/lite/toco/graph_transformations/remove_unused_op.cc
@@ -47,7 +47,8 @@ bool RemoveUnusedOp::Run(Model* model, std::size_t op_index) {
     bool found_output_as_rnn_state_array = false;
     for (const auto& rnn_state : model->flags.rnn_states()) {
       if (output == rnn_state.state_array()) {
-        CHECK(op->type == OperatorType::kFill);
+        CHECK(op->type == OperatorType::kFill ||
+              op->type == OperatorType::kTensorFlowIdentity);
         found_output_as_rnn_state_array = true;
         break;
       }