Statistics kernels?

[Moelf]

We have sum but not mean and var which exist in GPUArrays.jl, are they in-scope for this project?

In addition, what about Combinatorics.jl (like combination and permutations)

[anicusan]

Sure, mean and var would be very easy to add and similar in scope to sum.
Wouldn't the Combinatorics.jl work as iterators to pass to e.g. foreachindex? I haven't tested it, but we could make sure within Combinatorics.jl that they are GPU-able - is there anything else we'd need on the AK side?

[Moelf]

good to know, right now looking at JuliaGPU/GPUArrays.jl#611 and wonder if AK.jl can provide a more locked-down version of mean and var and actually be useful

[anicusan]

Hi, I’m back on GitHub after a few months away working on my PhD thesis - I’m catching up with AK issues and PRs now.

I looked a bit into implementing mean, var and std following the signatures from Statistics.jl - this seems to be a bit of a rabbit hole that I need some external / end-user input for:

• Do we prioritise raw speed or numerical consistency here? There are multiple algorithms to compute all three, e.g. compensated summation, Welford variance - which have different trade-offs in terms of numbers of kernel launches, memory usage, and numerical robustness.
• The Statistics.jl functions always return Float64 / BigFloat - what output types should we support? For sum we could keep it as an Int for Int inputs, but not here.
• Do we add the interfaces with a function as the first argument, e.g. mean(f, A::AbstractArray; dims)?

I guess I found the reason there's sum in Julia Base, but mean is inside Statistics.jl...

I wrote this initial mean function in case someone off Google is looking for an implementation:

function mean(
    src::AbstractArray, backend::Backend=get_backend(src);
    init=zero(promote_type(eltype(src), Float32)),
    dims::Union{Nothing, Int}=nothing,
    kwargs...
)
    if isnothing(dims)
        return sum(src, backend; init, kwargs...) / length(src)
    elseif size(src, dims) == 0
        # mapreduce_nd returns arrays filled with init when the reduced dimension is zero;
        # Statistics.mean returns NaN in those entries.
        return reduce(+, src, backend; init=oftype(init, NaN), dims, kwargs...)
    else
        # There are two ways to implement a multidimensional mean:
        # 1. A reduce (sum) followed by a map (division): two kernel launches, but minimum amount
        #    of compute
        # 2. Unified mapreduce that does division in the map phase: one kernel launch, but more
        #    compute (division is done on each element instead of just the final result).
        # The latter is faster.
        inv_size = oftype(init, 1) / size(src, dims)
        return mapreduce(x -> x * inv_size, +, src, backend; init, dims, kwargs...)
    end
end