auto slice_i2c_size =
      CreateI64DenseConst(builder, size_shape, size_values, location);

  return builder->create<TF::SliceOp>(
      location,
      RankedTensorType::get(
          size_values,
          mlir::cast<RankedTensorType>(input.getType()).getElementType()),
      input, slice_i2c_begin, slice_i2c_size);
}

      auto size_value = rewriter.create<tensor::ExtractOp>(loc, sizes, indices);
      Value minus_one = rewriter.create<arith::IndexCastOp>(
          loc, shape_scalar_type,
          rewriter.create<arith::ConstantIndexOp>(loc, -1));
      auto is_minus_one = rewriter.create<arith::CmpIOp>(
          loc, arith::CmpIPredicate::eq, size_value, minus_one);
      Value end_value =