Value zero_matrix = rewriter.create<BroadcastOp>(
        op.getLoc(), reshaped_input.getType(), zero,
        GetI64ElementsAttr({new_size, new_size}, &rewriter));
    Value masked =
        rewriter.create<SelectOp>(op.getLoc(), reshaped_input.getType(),
                                  compare, reshaped_input, zero_matrix);
    auto reduce = rewriter.create<ReduceOp>(op.getLoc(), masked, zero,