ahwillia / Einsum.jl

Einstein summation notation in Julia

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

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.