Define type of aggregate since not all types are closed under addition or multiplication
freemin7 opened this issue · comments
MWE:
using DynamicPolynomials
julia> using Einsum
julia> @ncpolyvar a[1:2,1:2]
(PolyVar{false}[a₁₋₁ a₁₋₂; a₂₋₁ a₂₋₂],)
julia> @ncpolyvar b[1:2,1:2]
(PolyVar{false}[b₁₋₁ b₁₋₂; b₂₋₁ b₂₋₂],)
julia> @einsum c[x,y] := a[x,z]*b[z,y]
ERROR: InexactError: convert(PolyVar{false}, a[1,1]*b[1,1] + a[1,2]*b[2,1])
Stacktrace:
[1] convert(#unused#::Type{PolyVar{false}}, p::Polynomial{false, Int64})
@ MultivariatePolynomials ~/.julia/packages/MultivariatePolynomials/KMk2o/src/conversion.jl:58
[2] setindex!(::Matrix{PolyVar{false}}, ::Polynomial{false, Int64}, ::Int64, ::Int64)
@ Base ./array.jl:841
[3] top-level scope
@ ~/.julia/packages/Einsum/AVMOj/src/Einsum.jl:178
If it was possible to define the aggregation type to Polynomial{false,Int64}
this would be solveable.
The unsatisfying answer is that Einsum isn't really maintained anymore. Have you tried Tullio.jl?
Other than that (if I understand the problem right): have you tried allocating c
first with the right type, and then use an einsum expression with =
instead of :=
?
Unfortunately Tullio gets this wrong too -- it infers the type as being that of the product without the sum. But perhaps it can learn to do better.
julia> @tullio c[x,y] := a[x,z]*b[z,y]
ERROR: InexactError: convert(Monomial{false}, a[1,1]*b[1,1] + a[1,2]*b[2,1])
Stacktrace:
[1] convert(#unused#::Type{Monomial{false}}, p::Polynomial{false, Int64})
@ MultivariatePolynomials ~/.julia/packages/MultivariatePolynomials/JG6mC/src/conversion.jl:58
[2] setindex!
@ ./array.jl:965 [inlined]
[3] 𝒜𝒸𝓉!
@ ~/.julia/dev/Tullio/src/macro.jl:1037 [inlined]
[4] tile_halves(fun!::var"#𝒜𝒸𝓉!#5", ::Type{Matrix{Monomial{false}}}, As::Tuple{Matrix{Monomial{false}}, Matrix{PolyVar{false}}, Matrix{PolyVar{false}}}, Is::Tuple{UnitRange{Int64}, UnitRange{Int64}}, Js::Tuple{UnitRange{Int64}}, keep::Nothing, final::Bool)
@ Tullio ~/.julia/dev/Tullio/src/threads.jl:139
...
julia> promote_type(eltype(a), eltype(b)) # @einsum
PolyVar{false}
julia> typeof(a[1] * b[1]) # @tullio
Monomial{false}
julia> eltype(a * b) # desired
Polynomial{false, Int64}
julia> typeof(a[1] * b[1] + a[1] * b[1])
Polynomial{false, Int64}
You can trick it by making the type of the right hand side wider:
julia> @tullio c[x,y] := Polynomial(a[x,z] * b[z,y])
2×2 Matrix{Polynomial{false, Int64}}:
a₁₋₁b₁₋₁ + a₁₋₂b₂₋₁ a₁₋₁b₁₋₂ + a₁₋₂b₂₋₂
a₂₋₁b₁₋₁ + a₂₋₂b₂₋₁ a₂₋₁b₁₋₂ + a₂₋₂b₂₋₂
julia> @tullio c[x,y] := Polynomial <| a[x,z] * b[z,y]
2×2 Matrix{Polynomial{false, Int64}}:
a₁₋₁b₁₋₁ + a₁₋₂b₂₋₁ a₁₋₁b₁₋₂ + a₁₋₂b₂₋₂
a₂₋₁b₁₋₁ + a₂₋₂b₂₋₁ a₂₋₁b₁₋₂ + a₂₋₂b₂₋₂
How does the constructor on the rhs affect allocations and vectorization?
I think to solve this we would need a promote type that takes account what operations are performed ... or just specify it manually.
Not sure, I know nothing about this package. I doubt that this can vectorize, given that it wraps a string.
julia> a[1] |> isbits
false
julia> a[1] |> dump
PolyVar{false}
id: Int64 279
name: String "a[1,1]"
julia> @btime $a[1] * $b[1]
min 50.499 ns, mean 64.375 ns (2 allocations, 176 bytes. GC mean 16.14%)
a₁₋₁b₁₋₁
But yes, currently it infers type of acting with *
on the arguments, when there is a reduction, it should really infer the result of then acting with +
too. I think this is the first example I've seen where these differ.