Repository: finch-tensor/Finch.jl
Pull request #694: Adding stats functions
Author: login: no-dpr
Created: 2025-01-17T18:24:31Z | Merged: 2025-03-25T17:07:44Z

fixes #574 

Implementing size property to lazy tensors to use with stats functions.

Linked issue #574: `mean`, `std`, and `var` functions
Hi @willow-ahrens,

This is a tracking issue for [mean](https://data-apis.org/array-api/latest/API_specification/generated/array_api.mean.html), [std](https://data-apis.org/array-api/latest/API_specification/generated/array_api.std.html), and [var](https://data-apis.org/array-api/latest/API_specification/generated/array_api.var.html) functions for `SwizzleArray`, `LazyArray`, and `Tensor` objects (similar to `max`, `min` that we already have).


Diff:
diff --git a/Project.toml b/Project.toml
index 23cbe6ee2..5919a8409 100644
--- a/Project.toml
+++ b/Project.toml
@@ -74,6 +74,7 @@ TensorMarket = "8b7d4fe7-0b45-4d0d-9dd8-5cc9b23b4b77"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 FileWatching = "7b1f6079-737a-58dc-b8bc-7a2ca5c1b5ee"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [targets]
-test = ["ReTestItems", "Test", "ArgParse", "LinearAlgebra", "Random", "SparseArrays", "Graphs", "SimpleWeightedGraphs", "HDF5", "NPZ", "Pkg", "TensorMarket", "Documenter", "FileWatching", "JuliaFormatter"]
+test = ["ReTestItems", "Test", "ArgParse", "LinearAlgebra", "Random", "SparseArrays", "Graphs", "SimpleWeightedGraphs", "HDF5", "NPZ", "Pkg", "TensorMarket", "Documenter", "FileWatching", "JuliaFormatter", "Statistics"]
diff --git a/docs/src/getting_started.md b/docs/src/getting_started.md
index 251f1815e..59c4eabec 100644
--- a/docs/src/getting_started.md
+++ b/docs/src/getting_started.md
@@ -103,13 +103,13 @@ julia> B = A .* 2
  0.0  0.0  0.0  0.0  0.0  0.0
 
 julia> C = lazy(A)
-?×?-LazyTensor{Float64}
+3×6-LazyTensor{Float64}
 
 julia> D = lazy(B)
-?×?-LazyTensor{Float64}
+3×6-LazyTensor{Float64}
 
 julia> E = (C .+ D) ./ 2
-?×?-LazyTensor{Float64}
+3×6-LazyTensor{Float64}
 
 julia> compute(E)
 3×6 Tensor{SparseDictLevel{Int64, Vector{Int64}, Vector{Int64}, Vector{Int64}, Dict{Tuple{Int64, Int64}, Int64}, Vector{Int64}, SparseDictLevel{Int64, Vector{Int64}, Vector{Int64}, Vector{Int64}, Dict{Tuple{Int64, Int64}, Int64}, Vector{Int64}, ElementLevel{0.0, Float64, Int64, Vector{Float64}}}}}:
diff --git a/src/Finch.jl b/src/Finch.jl
index 62f75aa76..25df496df 100644
--- a/src/Finch.jl
+++ b/src/Finch.jl
@@ -14,6 +14,7 @@ using Random: AbstractRNG, default_rng, randexp, randperm
 using PrecompileTools
 using Compat
 using DataStructures
+using Statistics
 using JSON
 using Distributions: Binomial, Normal, Poisson
 using TOML
diff --git a/src/Galley/FinchCompat/translate.jl b/src/Galley/FinchCompat/translate.jl
index 2a7a7d6f3..9694e9ebd 100644
--- a/src/Galley/FinchCompat/translate.jl
+++ b/src/Galley/FinchCompat/translate.jl
@@ -53,10 +53,8 @@ function aggs_to_mapjoins(prgm)
         Fixpoint(
             Prewalk(
                 Chain([
-                    (@rule aggregate(~op, ~init, ~arg) => if isnothing(init)
-                        arg
-                    end),
                     (@rule aggregate(~op, ~init, ~arg) => mapjoin(op, init, arg)),
+                    (@rule mapjoin(overwrite, nothing, ~arg) => arg),
                 ]),
             ),
         ),
@@ -111,7 +109,8 @@ function remove_reorders(prgm::LogicNode)
                         if length(idxs2) < length(getfields(arg))
                             reorder(
                                 aggregate(
-                                    nothing, nothing, arg, setdiff(getfields(arg), idxs2)...
+                                    overwrite, nothing, arg,
+                                    setdiff(getfields(arg), idxs2)...,
                                 ),
                                 idxs2...,
                             )
diff --git a/src/interface/eager.jl b/src/interface/eager.jl
index 4c3d76ab4..7fe152781 100644
--- a/src/interface/eager.jl
+++ b/src/interface/eager.jl
@@ -178,3 +178,63 @@ Expand the dimensions of an array by inserting a new singleton axis or axes that
 will appear at the `dims` position in the expanded array shape.
 """
 expanddims(arr::AbstractTensor, dims) = compute(expanddims(lazy(arr), dims))
+
+function Statistics.mean(tns::AbstractTensorOrBroadcast; dims=:)
+    res = compute(mean(lazy(tns); dims=dims))
+    if dims === Colon()
+        return res[]
+    else
+        return res
+    end
+end
+
+function Statistics.mean(f, tns::AbstractTensorOrBroadcast; dims=:)
+    res = compute(mean(f, lazy(tns); dims=dims))
+    if dims === Colon()
+        return res[]
+    else
+        return res
+    end
+end
+
+function Statistics.varm(tns::AbstractTensorOrBroadcast, m; corrected=true, dims=:)
+    res = compute(varm(lazy(tns), m; corrected=corrected, dims=dims))
+    if dims === Colon()
+        return res[]
+    else
+        return res
+    end
+end
+
+function Statistics.var(
+    tns::AbstractTensorOrBroadcast; corrected=true, mean=nothing, dims=:
+)
+    res = compute(var(lazy(tns); corrected=corrected, mean=mean, dims=dims))
+    if dims === Colon()
+        return res[]
+    else
+        return res
+    end
+end
+
+function Statistics.stdm(
+    tns::AbstractTensorOrBroadcast, m; corrected=true, dims=:
+)
+    res = compute(stdm(lazy(tns), lazy(m); corrected=corrected, dims=dims))
+    if dims === Colon()
+        return res[]
+    else
+        return res
+    end
+end
+
+function Statistics.std(
+    tns::AbstractTensorOrBroadcast; corrected=true, mean=nothing, dims=:
+)
+    res = compute(std(lazy(tns); corrected=corrected, mean=mean, dims=dims))
+    if dims === Colon()
+        return res[]
+    else
+        return res
+    end
+end
diff --git a/src/interface/einsum.jl b/src/interface/einsum.jl
index 158d55933..b8c1375e9 100644
--- a/src/interface/einsum.jl
+++ b/src/interface/einsum.jl
@@ -14,39 +14,40 @@ end
 
 einsum_tensor(tns) = EinsumTensor(einsum_style(tns), lazy(tns))
 
-struct EinsumArgument{T,Style}
+struct EinsumArgument{Vf,Tv,Style}
     style::Style
     data::LogicNode
-    extrude::Dict{Symbol,Bool}
-    fill_value::T
+    shape::Dict{Symbol,Any}
 end
 
-function EinsumArgument{T}(style::Style, data, extrude, fill_value) where {T,Style}
-    EinsumArgument{T,Style}(style, data, extrude, fill_value)
+function EinsumArgument{Vf,Tv}(style::Style, data, shape) where {Vf,Tv,Style}
+    EinsumArgument{Vf,Tv,Style}(style, data, shape)
 end
 
-Base.eltype(::EinsumArgument{T}) where {T} = T
+Base.eltype(::EinsumArgument{Vf,Tv}) where {Vf,Tv} = Tv
+fill_value(::EinsumArgument{Vf}) where {Vf} = Vf
 
 function einsum_access(tns::EinsumTensor, idxs...)
-    EinsumArgument{eltype(tns.arg)}(
+    EinsumArgument{fill_value(tns.arg),eltype(tns.arg)}(
         tns.style,
         relabel(tns.arg.data, map(field, idxs)...),
-        Dict(idx => idx_extrude for (idx, idx_extrude) in zip(idxs, tns.arg.extrude)),
-        tns.arg.fill_value,
+        Dict(idx => idx_shape for (idx, idx_shape) in zip(idxs, tns.arg.shape)),
     )
 end
 
 function einsum_op(op, args::EinsumArgument...)
-    EinsumArgument{return_type(DefaultAlgebra(), op, map(eltype, args)...)}(
+    EinsumArgument{
+        op((fill_value(arg) for arg in args)...),
+        return_type(DefaultAlgebra(), op, map(eltype, args)...),
+    }(
         reduce(result_style, [arg.style for arg in args]; init=EinsumEagerStyle()),
         mapjoin(op, (arg.data for arg in args)...),
-        mergewith(&, (arg.extrude for arg in args)...),
-        op((arg.fill_value for arg in args)...),
+        merge((arg.shape for arg in reverse(args))...),
     )
 end
 
 function einsum_immediate(val)
-    EinsumArgument{typeof(val)}(EinsumEagerStyle(), immediate(val), Dict(), val)
+    EinsumArgument{val,typeof(val)}(EinsumEagerStyle(), immediate(val), Dict())
 end
 
 struct EinsumProgram{Style,Arg<:LazyTensor}
@@ -55,23 +56,25 @@ struct EinsumProgram{Style,Arg<:LazyTensor}
 end
 
 function einsum(
-    ::typeof(overwrite), arg::EinsumArgument{T}, idxs...; init=nothing
-) where {T}
-    einsum(initwrite(arg.fill_value), arg, idxs...; init=arg.fill_value)
+    ::typeof(overwrite), arg::EinsumArgument{Vf}, idxs...; init=nothing
+) where {Vf}
+    einsum(initwrite(Vf), arg, idxs...; init=Vf)
 end
 
-function einsum(op, arg::EinsumArgument{T}, idxs...; init=initial_value(op, T)) where {T}
-    extrude = ntuple(n -> arg.extrude[idxs[n]], length(idxs))
+function einsum(
+    op, arg::EinsumArgument{Vf,Tv}, idxs...; init=initial_value(op, Tv)
+) where {Vf,Tv}
+    shape = ntuple(n -> arg.shape[idxs[n]], length(idxs))
     data = reorder(
         aggregate(
             immediate(op),
             immediate(init),
             arg.data,
-            map(field, setdiff(collect(keys(arg.extrude)), idxs))...,
+            map(field, setdiff(collect(keys(arg.shape)), idxs))...,
         ),
         map(field, idxs)...,
     )
-    einsum_execute(arg.style, LazyTensor{typeof(init)}(data, extrude, init))
+    einsum_execute(arg.style, LazyTensor{init,typeof(init)}(data, shape))
 end
 
 function einsum_execute(::EinsumEagerStyle, arg)
diff --git a/src/interface/lazy.jl b/src/interface/lazy.jl
index 51bcd68fd..e3b7e00c2 100644
--- a/src/interface/lazy.jl
+++ b/src/interface/lazy.jl
@@ -4,31 +4,36 @@ using Base: broadcasted
 using LinearAlgebra
 const AbstractArrayOrBroadcasted = Union{AbstractArray,Broadcasted}
 
-mutable struct LazyTensor{T,N}
+"""
+    LazyTensor{Vf, [Tv=typeof(Vf)], N, TI<:Tuple}()
+"""
+mutable struct LazyTensor{Vf,Tv,N,TI<:Tuple}
     data
-    extrude::NTuple{N,Bool}
-    fill_value::T
+    shape::TI
 end
-function LazyTensor{T}(data, extrude::NTuple{N,Bool}, fill_value) where {T,N}
-    LazyTensor{T,N}(data, extrude, fill_value)
+function LazyTensor{Vf}(data, shape::TI
+) where {Vf,TI<:Tuple}
+    LazyTensor{Vf,typeof(Vf),length(shape),TI}(data, shape)
+end
+
+function LazyTensor{Vf,Tv}(data, shape::TI
+) where {Vf,Tv,TI<:Tuple}
+    LazyTensor{Vf,Tv,length(shape),TI}(data, shape)
 end
 
 function Base.show(io::IO, tns::LazyTensor)
-    join(io, [x ? "1" : "?" for x in tns.extrude], "×")
+    join(io, tns.shape, "×")
     print(io, "-LazyTensor{", eltype(tns), "}")
 end
 
-Base.ndims(::Type{LazyTensor{T,N}}) where {T,N} = N
+Base.ndims(::Type{<:LazyTensor{Vf,Tv,N}}) where {Vf,Tv,N} = N
 Base.ndims(tns::LazyTensor) = ndims(typeof(tns))
-Base.eltype(::Type{<:LazyTensor{T}}) where {T} = T
+Base.eltype(::Type{<:LazyTensor{Vf,Tv}}) where {Vf,Tv} = Tv
 Base.eltype(tns::LazyTensor) = eltype(typeof(tns))
-fill_value(tns::LazyTensor) = tns.fill_value
+fill_value(::Type{<:LazyTensor{Vf}}) where {Vf} = Vf
+fill_value(tns::LazyTensor) = fill_value(typeof(tns))
 
-function Base.size(::LazyTensor)
-    throw(
-        ErrorException("Base.size is not supported for LazyTensor. Call `compute()` first.")
-    )
-end
+Base.size(tns::LazyTensor) = tns.shape
 
 function Base.getindex(::LazyTensor, i...)
     throw(
@@ -38,8 +43,8 @@ function Base.getindex(::LazyTensor, i...)
     )
 end
 
-function Base.getindex(arr::LazyTensor{T,N}, idxs::Vararg{Union{Nothing,Colon}}) where {T,N}
-    if length(idxs) - count(isnothing, idxs) != N
+function Base.getindex(arr::LazyTensor, idxs::Vararg{Union{Nothing,Colon}})
+    if length(idxs) - count(isnothing, idxs) != ndims(arr)
         throw(
             ArgumentError(
                 "Cannot index a lazy tensor with more or fewer `:` dims than it had original dims."
@@ -49,19 +54,23 @@ function Base.getindex(arr::LazyTensor{T,N}, idxs::Vararg{Union{Nothing,Colon}})
     return expanddims(arr, findall(isnothing, idxs))
 end
 
-function expanddims(arr::LazyTensor{T}, dims) where {T}
+function expanddims(arr::LazyTensor{Vf,Tv}, dims) where {Vf,Tv}
     @assert allunique(dims)
     @assert issubset(dims, 1:(ndims(arr) + length(dims)))
-    antidims = setdiff(1:(ndims(arr) + length(dims)), dims)
+    offset = zeros(Int, ndims(arr) + length(dims))
+    offset[dims] .= 1
+    offset = cumsum(offset)
     idxs_1 = [field(gensym(:i)) for _ in 1:ndims(arr)]
-    idxs_2 = [field(gensym(:i)) for _ in 1:(ndims(arr) + length(dims))]
-    idxs_2[antidims] .= idxs_1
+    idxs_2 = [
+        n in dims ? field(gensym(:i)) : idxs_1[n - offset[n]] for
+        n in 1:(ndims(arr) + length(dims))
+    ]
     data_2 = reorder(relabel(arr.data, idxs_1...), idxs_2...)
-    extrude_2 = [false for _ in 1:(ndims(arr) + length(dims))]
-    extrude_2[antidims] .= arr.extrude
-    return LazyTensor{T,ndims(arr) + length(dims)}(
-        data_2, tuple(extrude_2...), fill_value(arr)
+    shape_2 = ntuple(
+        n -> n in dims ? 1 : arr.shape[n - offset[n]], ndims(arr) + length(dims)
     )
+
+    return LazyTensor{Vf,Tv}(data_2, shape_2)
 end
 
 function identify(data)
@@ -69,29 +78,29 @@ function identify(data)
     subquery(lhs, data)
 end
 
-LazyTensor(data::Number) = LazyTensor{typeof(data),0}(immediate(data), (), data)
-LazyTensor{T}(data::Number) where {T} = LazyTensor{T,0}(immediate(data), (), data)
-LazyTensor(arr::Base.AbstractArrayOrBroadcasted) = LazyTensor{eltype(arr)}(arr)
-function LazyTensor{T}(arr::Base.AbstractArrayOrBroadcasted) where {T}
+LazyTensor(data::Number) = LazyTensor{data,typeof(data)}(immediate(data), ())
+LazyTensor{Vf}(data::Number) where {Vf} = LazyTensor{data,typeof(data)}(immediate(data), ())
+LazyTensor{Vf,Tv}(data::Number) where {Vf,Tv} = LazyTensor{Vf,Tv}(immediate(data), ())
+function LazyTensor(arr::Base.AbstractArrayOrBroadcasted)
+    LazyTensor{fill_value(arr),eltype(arr)}(arr)
+end
+function LazyTensor{Vf,Tv}(arr::Base.AbstractArrayOrBroadcasted) where {Vf,Tv}
     name = alias(gensym(:A))
     idxs = [field(gensym(:i)) for _ in 1:ndims(arr)]
-    extrude = ntuple(n -> size(arr, n) == 1, ndims(arr))
+    shape = size(arr)
     tns = subquery(name, table(immediate(arr), idxs...))
-    LazyTensor{eltype(arr),ndims(arr)}(tns, extrude, fill_value(arr))
+    LazyTensor{Vf,Tv}(tns, shape)
 end
-LazyTensor(arr::AbstractTensor) = LazyTensor{eltype(arr)}(arr)
+LazyTensor(arr::AbstractTensor) = LazyTensor{fill_value(arr),eltype(arr)}(arr)
 function LazyTensor(swizzle_arr::SwizzleArray{dims,<:Tensor}) where {dims}
     permutedims(LazyTensor(swizzle_arr.body), dims)
 end
-function LazyTensor{T}(arr::AbstractTensor) where {T}
+function LazyTensor{Vf,Tv}(arr::AbstractTensor) where {Vf,Tv}
     name = alias(gensym(:A))
     idxs = [field(gensym(:i)) for _ in 1:ndims(arr)]
-    extrude = ntuple(n -> size(arr)[n] == 1, ndims(arr))
+    shape = size(arr)
     tns = subquery(name, table(immediate(arr), idxs...))
-    LazyTensor{eltype(arr),ndims(arr)}(tns, extrude, fill_value(arr))
-end
-function LazyTensor{T}(swizzle_arr::SwizzleArray{dims,<:Tensor}) where {T,dims}
-    permutedims(LazyTensor{T}(swizzle_arr.body), dims)
+    LazyTensor{Vf,Tv}(tns, shape)
 end
 LazyTensor(data::LazyTensor) = data
 
@@ -114,26 +123,26 @@ end
 
 function Base.map(f, src::LazyTensor, args...)
     largs = map(LazyTensor, (src, args...))
-    extrude = largs[something(findfirst(arg -> length(arg.extrude) > 0, largs), 1)].extrude
-    idxs = [field(gensym(:i)) for _ in src.extrude]
+    shape = largs[something(findfirst(arg -> length(arg.shape) > 0, largs), 1)].shape
+    idxs = [field(gensym(:i)) for _ in shape]
     ldatas = map(largs) do larg
-        if larg.extrude == extrude
+        if larg.shape == shape
             return relabel(larg.data, idxs...)
-        elseif larg.extrude == ()
+        elseif larg.shape == ()
             return relabel(larg.data)
         else
             throw(DimensionMismatch("Cannot map across arrays with different sizes."))
         end
     end
-    T = return_type(DefaultAlgebra(), f, eltype(src), map(eltype, args)...)
-    new_default = f(map(fill_value, largs)...)
+    Tv_2 = return_type(DefaultAlgebra(), f, eltype(src), map(eltype, args)...)
+    Vf_2 = f(map(fill_value, largs)...)
     data = mapjoin(immediate(f), ldatas...)
-    return LazyTensor{T}(identify(data), src.extrude, new_default)
+    return LazyTensor{Vf_2,Tv_2}(identify(data), shape)
 end
 
 function Base.map!(dst, f, src::LazyTensor, args...)
     res = map(f, src, args...)
-    return LazyTensor(identify(reformat(dst, res.data)), res.extrude, res.fill_value)
+    return LazyTensor(identify(reformat(dst, res.data)), res.shape)
 end
 
 function initial_value(op, T)
@@ -158,16 +167,16 @@ function fixpoint_type(op, z, T)
 end
 
 function Base.reduce(
-    op, arg::LazyTensor{T,N}; dims=:, init=initial_value(op, T)
-) where {T,N}
+    op, arg::LazyTensor{Vf,Tv,N}; dims=:, init=initial_value(op, Tv)
+) where {Vf,Tv,N}
     dims = dims == Colon() ? (1:N) : collect(dims)
-    extrude = ((arg.extrude[n] for n in 1:N if !(n in dims))...,)
+    shape = ((arg.shape[n] for n in 1:N if !(n in dims))...,)
     fields = [field(gensym(:i)) for _ in 1:N]
-    S = fixpoint_type(op, init, eltype(arg))
+    Sv = fixpoint_type(op, init, Tv)
     data = aggregate(
         immediate(op), immediate(init), relabel(arg.data, fields), fields[dims]...
     )
-    LazyTensor{S}(identify(data), extrude, init)
+    LazyTensor{init,Sv}(identify(data), shape)
 end
 
 function tensordot(A::LazyTensor, B::Union{AbstractTensor,AbstractArray}, idxs; kwargs...)
@@ -179,13 +188,13 @@ end
 
 # tensordot takes in two tensors `A` and `B` and performs a product and contraction
 function tensordot(
-    A::LazyTensor{T1,N1},
-    B::LazyTensor{T2,N2},
+    A::LazyTensor{Vf1,Tv1,N1},
+    B::LazyTensor{Vf2,Tv2,N2},
     idxs;
     mult_op=*,
     add_op=+,
-    init=initial_value(add_op, return_type(DefaultAlgebra(), mult_op, T1, T2)),
-) where {T1,T2,N1,N2}
+    init=initial_value(add_op, return_type(DefaultAlgebra(), mult_op, Tv1, Tv2)),
+) where {Vf1,Vf2,Tv1,Tv2,N1,N2}
     if idxs isa Number
         idxs = ([i for i in 1:idxs], [i for i in 1:idxs])
     end
@@ -206,8 +215,8 @@ function tensordot(
         )
     end
 
-    extrude = ((A.extrude[n] for n in 1:N1 if !(n in A_idxs))...,
-        (B.extrude[n] for n in 1:N2 if !(n in B_idxs))...)
+    shape = ((A.shape[n] for n in 1:N1 if !(n in A_idxs))...,
+        (B.shape[n] for n in 1:N2 if !(n in B_idxs))...)
     A_fields = [field(gensym(:i)) for _ in 1:N1]
     B_fields = [field(gensym(:i)) for _ in 1:N2]
     reduce_fields = []
@@ -217,13 +226,15 @@ function tensordot(
     end
     AB = mapjoin(immediate(mult_op), relabel(A.data, A_fields), relabel(B.data, B_fields))
     AB_reduce = aggregate(immediate(add_op), immediate(init), AB, reduce_fields...)
-    T = return_type(DefaultAlgebra(), mult_op, T1, T2)
-    S = fixpoint_type(add_op, init, T)
-    return LazyTensor{S}(identify(AB_reduce), extrude, init)
+    Tv = return_type(DefaultAlgebra(), mult_op, Tv1, Tv2)
+    Sv = fixpoint_type(add_op, init, Tv)
+    return LazyTensor{init,Sv}(identify(AB_reduce), shape)
 end
 
 struct LazyStyle{N} <: BroadcastStyle end
-Base.Broadcast.BroadcastStyle(F::Type{<:LazyTensor{T,N}}) where {T,N} = LazyStyle{N}()
+function Base.Broadcast.BroadcastStyle(F::Type{<:LazyTensor{Vf,Tv,N}}) where {Vf,Tv,N}
+    LazyStyle{N}()
+end
 Base.Broadcast.broadcastable(tns::LazyTensor) = tns
 function Base.Broadcast.BroadcastStyle(a::LazyStyle{M}, b::LazyStyle{N}) where {M,N}
     LazyStyle{max(M, N)}()
@@ -250,26 +261,26 @@ function broadcast_to_query(bc::Broadcast.Broadcasted, idxs)
     mapjoin(immediate(bc.f), map(arg -> broadcast_to_query(arg, idxs), bc.args)...)
 end
 
-function broadcast_to_query(tns::LazyTensor{T,N}, idxs) where {T,N}
-    idxs_2 = [tns.extrude[i] ? field(gensym(:idx)) : idxs[i] for i in 1:N]
+function broadcast_to_query(tns::LazyTensor{Vf,Tv,N}, idxs) where {Vf,Tv,N}
+    idxs_2 = [isone(tns.shape[i]) ? field(gensym(:idx)) : idxs[i] for i in 1:N]
     data_2 = relabel(tns.data, idxs_2...)
-    reorder(data_2, idxs[findall(!, tns.extrude)]...)
+    reorder(data_2, idxs[[i for i in 1:N if !isone(tns.shape[i])]]...)
 end
 
-function broadcast_to_extrude(bc::Broadcast.Broadcasted, n)
-    all(map(arg -> broadcast_to_extrude(arg, n), bc.args))
+function broadcast_to_shape(bc::Broadcast.Broadcasted, n)
+    maximum(map(arg -> broadcast_to_shape(arg, n), bc.args))
 end
 
-function broadcast_to_extrude(tns::LazyTensor, n)
-    get(tns.extrude, n, false)
+function broadcast_to_shape(tns::LazyTensor, n)
+    get(tns.shape, n, 0)
 end
 
-function broadcast_to_default(bc::Broadcast.Broadcasted)
-    bc.f(map(arg -> broadcast_to_default(arg), bc.args)...)
+function broadcast_to_fill_value(bc::Broadcast.Broadcasted)
+    bc.f(map(arg -> broadcast_to_fill_value(arg), bc.args)...)
 end
 
-function broadcast_to_default(tns::LazyTensor)
-    tns.fill_value
+function broadcast_to_fill_value(tns::LazyTensor)
+    fill_value(tns)
 end
 
 function broadcast_to_eltype(bc::Broadcast.Broadcasted)
@@ -288,30 +299,30 @@ function Base.copy(bc::Broadcasted{LazyStyle{N}}) where {N}
     bc_lgc = broadcast_to_logic(bc)
     idxs = [field(gensym(:i)) for _ in 1:N]
     data = reorder(broadcast_to_query(bc_lgc, idxs), idxs)
-    extrude = ntuple(n -> broadcast_to_extrude(bc_lgc, n), N)
-    def = broadcast_to_default(bc_lgc)
-    return LazyTensor{broadcast_to_eltype(bc)}(identify(data), extrude, def)
+    shape = ntuple(n -> broadcast_to_shape(bc_lgc, n), N)
+    return LazyTensor{broadcast_to_fill_value(bc_lgc),broadcast_to_eltype(bc)}(
+        identify(data), shape
+    )
 end
 
 function Base.copyto!(::LazyTensor, ::Any)
     throw(ArgumentError("cannot materialize into a LazyTensor"))
 end
 
-function Base.copyto!(dst::AbstractArray, src::LazyTensor{T,N}) where {T,N}
-    return LazyTensor{T,N}(reformat(immediate(dst), src.data), src.extrude, src.fill_value)
+function Base.copyto!(dst::AbstractArray, src::LazyTensor{Vf,Tv}) where {Vf,Tv}
+    return LazyTensor{Vf,Tv}(reformat(immediate(dst), src.data), src.shape)
 end
 
-Base.permutedims(arg::LazyTensor{T,2}) where {T} = permutedims(arg, [2, 1])
-function Base.permutedims(arg::LazyTensor{T,N}, perm) where {T,N}
+Base.permutedims(arg::LazyTensor{Vf,Tv,2}) where {Vf,Tv} = permutedims(arg, [2, 1])
+function Base.permutedims(arg::LazyTensor{Vf,Tv,N}, perm) where {Vf,Tv,N}
     length(perm) == N ||
         throw(ArgumentError("permutedims given wrong number of dimensions"))
     isperm(perm) || throw(ArgumentError("permutedims given invalid permutation"))
     perm = collect(perm)
     idxs = [field(gensym(:i)) for _ in 1:N]
-    return LazyTensor{T,N}(
+    return LazyTensor{Vf,Tv}(
         reorder(relabel(arg.data, idxs...), idxs[perm]...),
-        arg.extrude[perm],
-        arg.fill_value,
+        (arg.shape[perm]...,),
     )
 end
 Base.permutedims(arr::SwizzleArray, perm) = swizzle(arr, perm...)
@@ -539,6 +550,41 @@ function LinearAlgebra.norm(arr::LazyTensor, p::Real=2)
     end
 end
 
+function Statistics.mean(tns::LazyTensor; dims=:)
+    dims = dims == Colon() ? (1:ndims(tns)) : collect(dims)
+    n = prod(collect(size(tns))[dims])
+    return sum(tns; dims=dims) ./ n
+end
+
+function Statistics.mean(f, tns::LazyTensor; dims=:)
+    dims = dims == Colon() ? (1:ndims(tns)) : collect(dims)
+    n = prod(collect(size(tns))[dims])
+    return sum(map(f, tns); dims=dims) ./ n
+end
+
+function Statistics.var(tns::LazyTensor; mean=nothing, corrected=true, dims=:)
+    dims = dims == Colon() ? (1:ndims(tns)) : collect(dims)
+    if mean === nothing
+        mean = expanddims(Statistics.mean(tns; dims=dims), dims)
+    end
+    n = prod(collect(size(tns))[dims])
+    return sum(abs2.(tns .- mean); dims=dims) ./ (n - corrected)
+end
+
+function Statistics.varm(tns::LazyTensor, mean; corrected=true, dims=:)
+    var(tns; mean=mean, corrected=corrected, dims=dims)
+end
+
+function Statistics.std(
+    tns::LazyTensor; corrected=true, mean=nothing, dims=:
+)
+    sqrt.(var(tns; corrected=corrected, mean=mean, dims=dims))
+end
+
+function Statistics.stdm(tns::LazyTensor, m; corrected=true, dims=:)
+    std(tns; corrected=corrected, mean=m, dims=dims)
+end
+
 """
     lazy(arg)
 
@@ -637,5 +683,13 @@ function compute_parse(ctx, args::Tuple)
     bodies = map((arg, var) -> query(var, arg.data), args, vars)
     prgm = plan(bodies, produces(vars))
 
-    return ctx(prgm)
+    ress = ctx(prgm)
+
+    @debug begin
+        for (arg, res) in zip(args, ress)
+            @assert size(arg) == size(res)
+        end
+    end
+
+    return ress
 end
diff --git a/test/Project.toml b/test/Project.toml
index c361ea1f1..a0f90d97b 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -30,3 +30,4 @@ SyntaxInterface = "b33eeca9-aacb-4496-a840-e75f1646a4fb"
 TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76"
 TensorMarket = "8b7d4fe7-0b45-4d0d-9dd8-5cc9b23b4b77"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
diff --git a/test/suites/interface_tests.jl b/test/suites/interface_tests.jl
index 293919d43..ad6c48240 100644
--- a/test/suites/interface_tests.jl
+++ b/test/suites/interface_tests.jl
@@ -161,6 +161,7 @@ end
 @testitem "interface_asmd" setup = [CheckOutput] begin
     using Finch: AsArray
     using SparseArrays
+    using Statistics
     using LinearAlgebra
 
     A = Tensor(
@@ -491,6 +492,81 @@ end
                     c = compute(sum(A[:, :, nothing] .* B[nothing, :, :]; dims=[2]))
                     @test c == c_correct
                 end
+
+                let
+                    A_ref = [0 0 44; 11 0 0; 22 00 55; 33 0 0]
+                    A = Tensor(A_ref)
+
+                    expected = mean(A_ref)
+                    result = mean(A)
+                    @test result == expected
+
+                    expected = mean(A_ref; dims=1:1)[1, :]
+                    result = mean(A; dims=1:1)
+                    @test all(isapprox.(result, expected))
+
+                    expected = mean(A_ref; dims=2:2)[:, 1]
+                    result = mean(A; dims=2:2)
+                    @test all(isapprox.(result, expected))
+                end
+
+                let
+                    A_ref = [0 0 44; 11 0 0; 22 00 55; 33 0 0]
+                    A = Tensor(A_ref)
+
+                    expected = var(A_ref; corrected=false)
+                    result = var(A; corrected=false)
+                    @test result == expected
+
+                    expected = var(A_ref)
+                    result = var(A)
+                    @test result == expected
+
+                    expected = var(A_ref; dims=1, corrected=false)[1, :]
+                    result = var(A; dims=1, corrected=false)
+                    @test all(isapprox.(result, expected))
+
+                    expected = var(A_ref; dims=1)[1, :]
+                    result = var(A; dims=1)
+                    @test all(isapprox.(result, expected))
+
+                    expected = var(A_ref; dims=2, corrected=false)[:, 1]
+                    result = var(A; dims=2, corrected=false)
+                    @test all(isapprox.(result, expected))
+
+                    expected = var(A_ref; dims=2)[:, 1]
+                    result = var(A; dims=2)
+                    @test all(isapprox.(result, expected))
+                end
+
+                let
+                    A_ref = [0 0 44; 11 0 0; 22 00 55; 33 0 0]
+                    A = Tensor(A_ref)
+
+                    expected = std(A_ref; corrected=false)
+                    result = std(A; corrected=false)
+                    @test result == expected
+
+                    expected = std(A_ref)
+                    result = std(A)
+                    @test result == expected
+
+                    expected = std(A_ref; dims=1, corrected=false)[1, :]
+                    result = std(A; dims=1, corrected=false)
+                    @test all(isapprox.(result, expected))
+
+                    expected = std(A_ref; dims=1)[1, :]
+                    result = std(A; dims=1)
+                    @test all(isapprox.(result, expected))
+
+                    expected = std(A_ref; dims=2, corrected=false)[:, 1]
+                    result = std(A; dims=2, corrected=false)
+                    @test all(isapprox.(result, expected))
+
+                    expected = std(A_ref; dims=2)[:, 1]
+                    result = std(A; dims=2)
+                    @test all(isapprox.(result, expected))
+                end
             end
         end
     end


Reviews:
user: login: willow-ahrens
body: please rename the field
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-01-17T18:25:33Z
user: login: willow-ahrens
body: Great work! This is looking good! I've added some comments about packaging. This PR needs some tests as well. You could add tests to interface_tests.jl. Be sure to leave a link to the github issue to document what precisely the tests are testing.
state: CHANGES_REQUESTED
author_association: MEMBER
submitted_at: 2025-01-22T15:44:21Z
user: login: willow-ahrens
state: APPROVED
author_association: MEMBER
submitted_at: 2025-03-25T15:02:59Z
user: login: willow-ahrens
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-25T15:37:27Z
user: login: kylebd99
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-25T16:13:28Z

Review comments:
diff_hunk: @@ -7,12 +7,13 @@ const AbstractArrayOrBroadcasted = Union{AbstractArray,Broadcasted}
 mutable struct LazyTensor{T, N}
     data
     extrude::NTuple{N, Bool}
+    size::NTuple{N, UInt64}
path: src/interface/lazy.jl
user: login: willow-ahrens
body: This field should be renamed `shape` for consistency with other naming choices in the codebase
created_at: 2025-01-17T18:25:22Z
author_association: MEMBER
diff_hunk: @@ -54,6 +54,7 @@ export lazy, compute, fused, tensordot, @einsum
 export choose, minby, maxby, overwrite, initwrite, filterop, d
 
 export fill_value, AsArray, expanddims, tensor_tree
+export mean
path: src/Finch.jl
user: login: willow-ahrens
body: instead of exporting our own mean, we should use the `mean` from the `Statistics` standard library. We also don't want to depend on `Statistics` explicitly, instead we should use an "optional dependency", where `mean`, `std`, and `var` definitions are loaded only when the user loads both `Finch` and `Statistics`. `Finch` currently takes this approach with `SparseArrays`. You can see the approach in [`ext/SparseArraysExt.jl`](https://github.com/finch-tensor/Finch.jl/blob/main/ext/SparseArraysExt.jl). We would also need to handle `Statistics` the same way we handle `SparseArrays` in the [`Project.toml`](https://github.com/finch-tensor/Finch.jl/blob/main/Project.toml). More on this concept here: https://pkgdocs.julialang.org/v1/creating-packages/#Conditional-loading-of-code-in-packages-(Extensions)
created_at: 2025-01-22T15:33:44Z
author_association: MEMBER
diff_hunk: @@ -112,14 +112,25 @@ Base.maximum(arr::AbstractTensorOrBroadcast; kwargs...) = reduce(max, arr; init
 
 Base.extrema(arr::AbstractTensorOrBroadcast; kwargs...) = mapreduce(plex, min1max2, arr; init = (typemax(broadcast_to_eltype(arr)), typemin(broadcast_to_eltype(arr))), kwargs...)
 
+
 function LinearAlgebra.norm(arr::AbstractTensorOrBroadcast, p::Real = 2)
     compute(norm(lazy(arr), p))[]
 end
 
 """
-    expanddims(arr::AbstractTensor, dims)
+expanddims(arr::AbstractTensor, dims)
path: src/interface/eager.jl
user: login: willow-ahrens
body: This line should stay indented
created_at: 2025-01-22T15:34:00Z
author_association: MEMBER
diff_hunk: @@ -7,12 +7,13 @@ const AbstractArrayOrBroadcasted = Union{AbstractArray,Broadcasted}
 mutable struct LazyTensor{T, N}
     data
     extrude::NTuple{N, Bool}
+    shape::NTuple{N, Int}
path: src/interface/lazy.jl
user: login: willow-ahrens
body: For better or worse, Finch array dimensions may use different types independently. So you could have a matrix with row dimension `UInt8` and column dimension `BigInt`. The way we handle this is to parameterize the type of LazyTensor based on the type of the dimensions: 
```
mutable struct LazyTensor{T, N, Shape <: Tuple}
    data
    extrude::NTuple{N, Bool}
    shape::Shape
    fill_value::T
end
```
Note that the dimension tuple is not constrained to be length `N` with the type annotations. That's okay, we can check that it's correct when we construct the LazyTensor.
created_at: 2025-01-22T15:38:44Z
author_association: MEMBER
diff_hunk: @@ -7,12 +7,13 @@ const AbstractArrayOrBroadcasted = Union{AbstractArray,Broadcasted}
 mutable struct LazyTensor{T, N}
     data
     extrude::NTuple{N, Bool}
+    shape::NTuple{N, Int}
     fill_value::T
 end
-LazyTensor{T}(data, extrude::NTuple{N, Bool}, fill_value) where {T, N} = LazyTensor{T, N}(data, extrude, fill_value)
+LazyTensor{T}(data, extrude::NTuple{N, Bool}, shape::NTuple{N, Int}, fill_value) where {T, N} = LazyTensor{T, N}(data, extrude, shape, fill_value)
 
 function Base.show(io::IO, tns::LazyTensor)
-    join(io, [x ? "1" : "?" for x in tns.extrude], "×")
+    join(io, [string(x) for x in tns.shape], "×")
path: src/interface/lazy.jl
user: login: willow-ahrens
body: I think you may also be able to get away with `join(io, tns.shape, "×")`, but I'm not sure, try it?
created_at: 2025-01-22T15:39:21Z
author_association: MEMBER
diff_hunk: @@ -103,12 +113,12 @@ function Base.map(f, src::LazyTensor, args...)
     T = return_type(DefaultAlgebra(), f, eltype(src), map(eltype, args)...)
     new_default = f(map(fill_value, largs)...)
     data = mapjoin(immediate(f), ldatas...)
-    return LazyTensor{T}(identify(data), src.extrude, new_default)
+    return LazyTensor{T}(identify(data), src.extrude, src.shape, new_default)
 end
 
 function Base.map!(dst, f, src::LazyTensor, args...)
     res = map(f, src, args...)
-    return LazyTensor(identify(reformat(dst, res.data)), res.extrude, res.fill_value)
+    return LazyTensor(identify(reformat(dst, res.data)), res.extrude, res.size, res.fill_value)
path: src/interface/lazy.jl
user: login: willow-ahrens
body: Almost certainly, this should be `size(res)` or `res.shape`, depending on whether or not we know that `res` is a `LazyTensor`.
created_at: 2025-01-22T15:40:06Z
author_association: MEMBER
diff_hunk: @@ -111,7 +109,8 @@ function remove_reorders(prgm::LogicNode)
                         if length(idxs2) < length(getfields(arg))
                             reorder(
                                 aggregate(
-                                    nothing, nothing, arg, setdiff(getfields(arg), idxs2)...
+                                    overwrite, nothing, arg,
path: src/Galley/FinchCompat/translate.jl
user: login: willow-ahrens
body: @kylebd99, I found this was necessary to get things to work, and I'm not sure why. The issue was that merge_stats_def would try to evaluate the aggregate function because the original rule to remove `isnothing(init)` aggregates never matched. I'm just confused about how we never hit this before.
created_at: 2025-03-25T15:37:26Z
author_association: MEMBER
diff_hunk: @@ -111,7 +109,8 @@ function remove_reorders(prgm::LogicNode)
                         if length(idxs2) < length(getfields(arg))
                             reorder(
                                 aggregate(
-                                    nothing, nothing, arg, setdiff(getfields(arg), idxs2)...
+                                    overwrite, nothing, arg,
path: src/Galley/FinchCompat/translate.jl
user: login: kylebd99
body: Hmm, yeah I don't have a lot of intuition on this one. It seems like the original version should have checked isnothing on `op` not `init`. But, I'm quite surprised that it was never firing.
created_at: 2025-03-25T16:13:27Z
author_association: COLLABORATOR

PR comments:
user: login: willow-ahrens
created_at: 2025-03-07T21:24:23Z
body: Let's also check that the shape of the lazy tensors is always equal to the shape of the returned tensors when in debug mode using an `@assert` statement in `@debug` mode [here](https://github.com/finch-tensor/Finch.jl/blob/b5a1a87d49365ab6c694d472f72e22586cbd9ce0/src/interface/lazy.jl#L637)
author_association: MEMBER
user: login: codecov[bot]
created_at: 2025-03-25T16:46:52Z
body: ## [Codecov](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor) Report
Attention: Patch coverage is `74.62687%` with `34 lines` in your changes missing coverage. Please review.

| [Files with missing lines](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor) | Patch % | Lines |
|---|---|---|
| [src/interface/lazy.jl](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&filepath=src%2Finterface%2Flazy.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor#diff-c3JjL2ludGVyZmFjZS9sYXp5Lmps) | 79.56% | [19 Missing :warning: ](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor) |
| [src/interface/eager.jl](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&filepath=src%2Finterface%2Feager.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor#diff-c3JjL2ludGVyZmFjZS9lYWdlci5qbA==) | 50.00% | [15 Missing :warning: ](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor) |

| [Files with missing lines](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor) | Coverage Δ | |
|---|---|---|
| [src/Finch.jl](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&filepath=src%2FFinch.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor#diff-c3JjL0ZpbmNoLmps) | `94.68% <ø> (ø)` | |
| [src/Galley/FinchCompat/translate.jl](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&filepath=src%2FGalley%2FFinchCompat%2Ftranslate.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor#diff-c3JjL0dhbGxleS9GaW5jaENvbXBhdC90cmFuc2xhdGUuamw=) | `67.47% <ø> (+0.54%)` | :arrow_up: |
| [src/interface/einsum.jl](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&filepath=src%2Finterface%2Feinsum.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor#diff-c3JjL2ludGVyZmFjZS9laW5zdW0uamw=) | `78.57% <100.00%> (+0.25%)` | :arrow_up: |
| [src/interface/eager.jl](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&filepath=src%2Finterface%2Feager.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor#diff-c3JjL2ludGVyZmFjZS9lYWdlci5qbA==) | `70.65% <50.00%> (-8.39%)` | :arrow_down: |
| [src/interface/lazy.jl](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694?src=pr&el=tree&filepath=src%2Finterface%2Flazy.jl&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor#diff-c3JjL2ludGVyZmFjZS9sYXp5Lmps) | `77.61% <79.56%> (-1.22%)` | :arrow_down: |

... and [4 files with indirect coverage changes](https://app.codecov.io/gh/finch-tensor/Finch.jl/pull/694/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=finch-tensor)

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE