const auto input_to_output_weights =
      op.getInputToOutputWeights().getType().cast<ShapedType>();
  const auto recurrent_to_output_weights =
      op.getRecurrentToOutputWeights().getType().cast<ShapedType>();
  if (activation_state.hasStaticShape() && cell_state.hasStaticShape() &&
      input_to_output_weights.hasStaticShape() &&
      recurrent_to_output_weights.hasStaticShape()) {

      /*input_to_input_weights=*/weights_array->getResult(0),
      /*input_to_forget_weights=*/weights_array->getResult(1),
      /*input_to_cell_weights=*/weights_array->getResult(2),
      /*input_to_output_weights=*/weights_array->getResult(3),
      /*recurrent_to_input_weights=*/recurrent_to_input_weights,
      /*recurrent_to_forget_weights=*/recurrent_to_forget_weights,

  func.return %24 : tensor<4xf32>

}

// -----

// 'input_to_output_weights' input for lstm op has unmatched rank with `input`.

    // Weights
    TFL_TensorOfOrNone<[F32, QI8]>:$input_to_input_weights,
    TFL_TensorOf<[F32, QI8]>:$input_to_forget_weights,
    TFL_TensorOf<[F32, QI8]>:$input_to_cell_weights,
    TFL_TensorOf<[F32, QI8]>:$input_to_output_weights,

    // Recurrent weights
    TFL_TensorOfOrNone<[F32, QI8]>:$recurrent_to_input_weights,
    TFL_TensorOf<[F32, QI8]>:$recurrent_to_forget_weights,
    TFL_TensorOf<[F32, QI8]>:$recurrent_to_cell_weights,