rafaqz / DimensionalData.jl

Named dimensions and indexing for julia arrays and other data
https://rafaqz.github.io/DimensionalData.jl/stable/
MIT License
269 stars 37 forks source link

maximum and friends not type-inferrable #587

Closed sethaxen closed 6 months ago

sethaxen commented 7 months ago

Calling reductions like maximum on Julia v1.10.0 is not type-inferrable:

julia> using DimensionalData, Test

julia> x = DimArray(randn(2, 3, 4), (X, Y, Z));

julia> foo(x) = @inline maximum(x; dims=(1, 2))

julia> @inferred foo(x)
ERROR: return type DimArray{Float64, 3, Tuple{X{DimensionalData.Dimensions.LookupArrays.NoLookup{Base.OneTo{Int64}}}, Y{DimensionalData.Dimensions.LookupArrays.NoLookup{Base.OneTo{Int64}}}, Z{DimensionalData.Dimensions.LookupArrays.NoLookup{Base.OneTo{Int64}}}}, Tuple{}, Array{Float64, 3}, DimensionalData.NoName, DimensionalData.Dimensions.LookupArrays.NoMetadata} does not match inferred return type DimArray{Float64, 3, D, Tuple{}, Array{Float64, 3}, DimensionalData.NoName, DimensionalData.Dimensions.LookupArrays.NoMetadata} where D<:Tuple
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] top-level scope
   @ REPL[22]:1

Ultimately we're hitting the following line: https://github.com/rafaqz/DimensionalData.jl/blob/d927b773cd3243fc3c949371a3371938898fd811/src/array/methods.jl#L20-L21 so I thought inlining a call to that method directly might be type-stable, but nope.

julia> foo(x) = @inline DimensionalData._maximum(x, (1, 2));

julia> @inferred foo(x)
ERROR: return type DimArray{Float64, 3, Tuple{X{DimensionalData.Dimensions.LookupArrays.NoLookup{Base.OneTo{Int64}}}, Y{DimensionalData.Dimensions.LookupArrays.NoLookup{Base.OneTo{Int64}}}, Z{DimensionalData.Dimensions.LookupArrays.NoLookup{Base.OneTo{Int64}}}}, Tuple{}, Array{Float64, 3}, DimensionalData.NoName, DimensionalData.Dimensions.LookupArrays.NoMetadata} does not match inferred return type DimArray{Float64, 3, D, Tuple{}, Array{Float64, 3}, DimensionalData.NoName, DimensionalData.Dimensions.LookupArrays.NoMetadata} where D<:Tuple
Stacktrace:
 [1] error(s::String)
   @ Base ./error.jl:35
 [2] top-level scope
   @ REPL[19]:1

But if I just implement the identical function myself, it just works.

julia> function bar(x)
           dims = (1, 2)
           vals = maximum(parent(x); dims=Dimensions.dimnum(x, DimensionalData._astuple(dims)))
           newdims = Dimensions.reducedims(x, dims)
           DimensionalData.rebuild(x, vals, newdims)
       end;

julia> @inferred bar(x)
1×1×4 DimArray{Float64,3} with dimensions: X, Y, Z
[:, :, 1]
 0.882391
[and 3 more slices...]

So I checked, and decorating the following with @inline or Base.@constprop :aggressive causes maximum to be type-inferrable: https://github.com/rafaqz/DimensionalData.jl/blob/d927b773cd3243fc3c949371a3371938898fd811/src/array/methods.jl#L20-L21

rafaqz commented 7 months ago

If we just make everything @generated it will be stable on all versions !

But seriously I guess we can specialise to v1.10, or add a Compat.jl dep

sethaxen commented 7 months ago

I've done the Compat.jl dependency elsewhere and it's been fine. I'm not really certain the trade-offs of that vs inlining.

rafaqz commented 6 months ago

I missed that you said @inline worked as well as @constprop, we can just use that.