Surprising results with expressions involving additions
Got bit by this (in a more complicated context)
julia> a = randn(2)
2-element Vector{Float64}:
-0.012727152255729035
-0.4625600586199899
julia> @tullio b := a[i]*a[i]
0.214123788235069
julia> @tullio b := a[i]*a[i]+1
2.214123788235069
I would have expected 1.21 (mathematically that's what you would get), although I understand how 2.21 could be a sensible answer also. Is this by design (because Tullio treats the RHS as a large expression and doesn't really go inside) or could this be changed?
edit: forgot to say thanks a lot for this package, it's awesome! Super flexible and wonderful syntax
Sorry I missed this. I agree this behaviour is a potential surprise.
It is intentional, though, in that Tullio always wants to write one set of nested loops, and doesn't think too hard about what the right hand side contains -- arbitrary functions are allowed.
Different packages differ on this. The example I put in this readme is:
julia> using Einsum, TensorOperations
julia> E = [5,10]; F = ones(Int, 2,3,3);
julia> @tensor D[i] := 2 * E[i] + F[i,k,k] # partial trace of F only, Dᵢ = 2Eᵢ + Σⱼ Fᵢⱼⱼ
2-element Vector{Int64}:
13
23
julia> @einsum G[i] := 2 * E[i] + F[i,k,k] # the sum includes everyting: Gᵢ = Σⱼ (2Eᵢ + Fᵢⱼⱼ)
2-element Vector{Int64}:
33
63
Einsum's logic it much like Tullio's, one loop nest. For TensorCast, you would write @reduce C[i] := sum(i) 2 * E[i] + F[i,k,k] and it would again reduce over the whole right hand side. (Except that this doesn't work thanks to repeated k.) But having sum(i) in front at least makes it clear what the intention is. The technical reason in that case is that it always broadcasts to make one large array, and then applies one reduction function.
Whether this could be changed is perhaps something to think about. It wouldn't be impossible to notice the +, and decompose to two steps, each of which makes its own loop nest, without intermediate allocations:
@tullio D[i] := 2 * E[i]
@tullio D[i] += F[i,k,k]
Trickier cases though are @tullio H[i] := 2(E[i] + F[i,k,k]), or @tullio J[i] := 2*E[i] + log(F[i,k,k]) |> exp (in which exp is after the sum, right now), should these also be split? And @tullio (*) P[i] := 2 * E[i] * F[i,k,k], should this similarly split on *?
The other tricky thing is that it wants to differentiate the right hand side. This two-stage thing would need to know about that, I think, to propagate dD back to both F and E.