auto create_slice_op = [&](int32_t axis_i, int32_t begin_i, int32_t size_i,
                               Value input) {
      SmallVector<int64_t, 4> begin_values(input_rank, 0);
      begin_values[axis_i] = begin_i;
      auto begin_attr = DenseIntElementsAttr::get(axis_type, begin_values);
      auto begin =
          rewriter.create<ConstOp>(op->getLoc(), axis_type, begin_attr);

      SmallVector<int64_t, 4> output_shape;