auto keys =
            CreateRngUniform32(op.getLoc(), num_elements, /*lower_limit=*/0,
                               /*upper_limit=*/u32_max, &rewriter);
        auto sorted = createSortOp(
            &rewriter, op.getLoc(), {keys, current},
            {rewriter.getIntegerType(32), input_type.getElementType()},
            /*dimension=*/-1, /*isStable=*/false,