LogicalResult UnrollSplitV::matchAndRewrite(TFL::SplitVOp splitv_op,
                                            PatternRewriter& rewriter) const {
  // We need to make sure both splits & split dim are constants.
  auto splits = splitv_op.getSizeSplits().getDefiningOp();
  mlir::DenseIntElementsAttr splits_attr;
  if (!splits || !matchPattern(splits, m_Constant(&splits_attr)))
    return failure();

  auto split_dim = splitv_op.getSplitDim().getDefiningOp();

  auto num_split = rewriter.getI32IntegerAttr(tf_splitv_op.getNumSplit());

  rewriter.replaceOpWithNewOp<TFL::SplitVOp>(
      op, tf_splitv_op.getOutput().getTypes(), tf_splitv_op.getValue(),
      tf_splitv_op.getSizeSplits(), tf_splitv_op.getSplitDim(), num_split);
  return success();
}

LogicalResult ConvertTFUnpackOp::matchAndRewrite(
    Operation* op, PatternRewriter& rewriter) const {

// SplitVOp
//===----------------------------------------------------------------------===//

LogicalResult SplitVOp::verify() {
  SplitVOp op = *this;
  auto split_sizes_type =
      op.getSizeSplits().getType().dyn_cast<RankedTensorType>();
  if (!split_sizes_type) return success();

  if (split_sizes_type.getRank() != 1 ||
      (!ShapedType::isDynamic(split_sizes_type.getDimSize(0)) &&

  // If 'size_splits' is not a constant, there are no other checks.
  ElementsAttr size_splits_attr;
  if (!matchPattern(op.getSizeSplits(), m_Constant(&size_splits_attr)))
    return success();

  if (size_splits_attr.getNumElements() != num_splits) {
    auto size_splits_type =
        op.getSizeSplits().getType().cast<RankedTensorType>();
    RankedTensorType expected_size_splits_type =

      return failure();

    // We can only match when the split sizes is a constant int vector.
    DenseIntElementsAttr split_sizes_attr;
    if (!matchPattern(op.getSizeSplits(), m_Constant(&split_sizes_attr)))
      return failure();

    // Get each chunck's size along the dimension to split. It may contain
    // dynamic sizes and we need to update it if so.