Operation* op, PatternRewriter& rewriter) const {
  auto tf_split_op = cast<TF::SplitOp>(op);

  // Number of splits cannot be negative.
  auto num_split = rewriter.getI32IntegerAttr(tf_split_op.getNumSplit());

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