TFL_StatefulTensor:$fw_input_cell_state,
    TFL_StatefulTensor:$bw_input_activation_state,
    TFL_StatefulTensor:$bw_input_cell_state,

    // Auxiliary input & weights.
    TFL_TensorOfOrNone<[F32, I8]>:$aux_input,
    // Auxiliary fw weights.
    TFL_TensorOfOrNone<[F32, I8]>:$fw_aux_input_to_input_weights,
    TFL_TensorOfOrNone<[F32, I8]>:$fw_aux_input_to_forget_weights,

             llvm::enumerate(block.getArgumentTypes())) {
          TensorType updated_ty =
              UpdateElementTypeTo(original_ty, element_types[block_idx]);
          signature.addInputs(block_idx, {updated_ty});
        }
        rewriter.applySignatureConversion(&region.front(), signature);
      }
    }

    // Replace all uses of `op` results with the newly created op.