use lazy broadcasting for var

Author: mcabbott

src/host/statistics.jl

@@ -2,13 +2,15 @@ using Statistics
 
 function Statistics.varm(A::AbstractGPUArray{<:Real}, M::AbstractArray{<:Real};
                          dims, corrected::Bool=true)
+    T = float(eltype(A))
+    λ = convert(T, inv(_mean_denom(A, dims) - corrected))
     #B = (A .- M).^2
     # NOTE: the above broadcast promotes to Float64 and uses power_by_squaring...
-    B = broadcast(A, M) do a, m
+    B = Broadcast.broadcasted(A, M) do a, m
         x = (a - m)
-        x*x
+        λ * x * x
     end
-    sum(B, dims=dims)/(prod(size(A)[[dims...]])::Int-corrected)
+    sum(Broadcast.instantiate(B); dims)
 end
 
 Statistics.stdm(A::AbstractGPUArray{<:Real},m::AbstractArray{<:Real}, dim::Int; corrected::Bool=true) =
@@ -49,3 +51,7 @@ function Statistics.corzm(x::AbstractGPUMatrix, vardim::Int=1)
     c = Statistics.unscaled_covzm(x, vardim)
     return Statistics.cov2cor!(c, sqrt.(diag(c)))
 end
+
+_mean_denom(x::AbstractArray, dims::Integer) = size(x, dims)
+_mean_denom(x::AbstractArray, dims::Colon) = length(x)
+_mean_denom(x::AbstractArray, dims) = prod(size(x,d) for d in unique(dims); init=1)