| 113 | }; |
| 114 | |
| 115 | TEST_F(LstmUtilsTest, ConvertLSTMCellSimple) { |
| 116 | mlir::TFL::ConvertLSTMCellSimpleToFusedLSTM convert(fused_lstm_func_); |
| 117 | |
| 118 | auto result = convert.RewriteFunc(); |
| 119 | EXPECT_FALSE(failed(result)); |
| 120 | fused_lstm_func_.dump(); |
| 121 | |
| 122 | // verify transpose |
| 123 | EXPECT_EQ( |
| 124 | fused_lstm_func_.getAttrOfType<StringAttr>(kTFImplements).getValue(), |
| 125 | convert.GetCompositeOpName()); |
| 126 | EXPECT_EQ(fused_lstm_func_.getNumArguments(), 5); |
| 127 | EXPECT_EQ(fused_lstm_func_.getType().getNumResults(), 1); |
| 128 | |
| 129 | auto transpose_op = fused_lstm_func_.getBody().front().begin(); |
| 130 | transpose_op++; |
| 131 | EXPECT_EQ( |
| 132 | transpose_op->getOperand(0).getType().cast<RankedTensorType>().getDimSize( |
| 133 | 0), |
| 134 | 3); |
| 135 | EXPECT_EQ( |
| 136 | transpose_op->getOperand(0).getType().cast<RankedTensorType>().getDimSize( |
| 137 | 1), |
| 138 | 12); |
| 139 | EXPECT_EQ( |
| 140 | transpose_op->getResult(0).getType().cast<RankedTensorType>().getDimSize( |
| 141 | 0), |
| 142 | 12); |
| 143 | EXPECT_EQ( |
| 144 | transpose_op->getResult(0).getType().cast<RankedTensorType>().getDimSize( |
| 145 | 1), |
| 146 | 3); |
| 147 | |
| 148 | auto it = fused_lstm_func_.getBody().back().rbegin(); |
| 149 | EXPECT_EQ(it->getName().getStringRef(), mlir::ReturnOp::getOperationName()); |
| 150 | it++; // tensor_cast |
| 151 | it++; // lstm |
| 152 | EXPECT_EQ(it->getName().getStringRef(), |
| 153 | mlir::TFL::LSTMOp::getOperationName()); |
| 154 | EXPECT_EQ(it->getNumOperands(), 24); |
| 155 | EXPECT_EQ(it->getNumResults(), 1); |
| 156 | // cifg = false, so input2input is not None. |
| 157 | EXPECT_FALSE(it->getOperand(1).getType().isa<NoneType>()); |
| 158 | // input layer norm is None |
| 159 | EXPECT_TRUE(it->getOperand(20).getType().isa<NoneType>()); |
| 160 | // proj_bias is F32 |
| 161 | EXPECT_TRUE(it->getOperand(17) |
| 162 | .getType() |
| 163 | .cast<RankedTensorType>() |
| 164 | .getElementType() |
| 165 | .isF32()); |
| 166 | |
| 167 | // output gate bias is 0 since it is out of bounds of the bias tensor, so |
| 168 | // we set its value as a const tensor of specified size and value 0. |
| 169 | EXPECT_TRUE( |
| 170 | mlir::cast<mlir::ConstantOp>(it->getOpOperand(15).get().getDefiningOp()) |
| 171 | .getValue() |
| 172 | .cast<ElementsAttr>() |