important bug fixes

Author: Jutho

Project.toml

@@ -1,7 +1,7 @@
 name = "TensorOperations"
 uuid = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
 authors = ["Lukas Devos <[email protected]>", "Maarten Van Damme <[email protected]>", "Jutho Haegeman <[email protected]>"]
-version = "4.0.1"
+version = "4.0.2"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/indexnotation/parser.jl

@@ -113,7 +113,17 @@ function tensorify(ex::Expr)
                                        ExistingTensor)
                 end
             else
-                return instantiate(dst, false, rhs, true, leftind, rightind, NewTensor)
+                # deal with the case that dst can be an existing variable while also a new value is assigned to it
+                if dst ∈ getinputtensorobjects(rhs)
+                    dst2 = gensym(dst)
+                    return quote
+                        $dst2 = $dst
+                        $(instantiate(dst2, false, rhs, true, leftind, rightind, NewTensor))
+                        $dst = $dst2
+                    end
+                else
+                    return instantiate(dst, false, rhs, true, leftind, rightind, NewTensor)
+                end
             end
         elseif isassignment(ex) && isscalarexpr(lhs)
             if istensorexpr(rhs) && isempty(getindices(rhs))
@@ -136,12 +146,6 @@ function tensorify(ex::Expr)
     if ex.head == :block # @tensor begin ... end
         return Expr(ex.head, map(tensorify, ex.args)...)
     end
-    if ex.head == :for # @tensor for ... end
-        return Expr(ex.head, ex.args[1], tensorify(ex.args[2]))
-    end
-    if ex.head == :function # @tensor function ... end
-        return Expr(ex.head, ex.args[1], tensorify(ex.args[2]))
-    end
     # constructions of the form: a = @tensor ...
     if isscalarexpr(ex)
         return instantiate_scalar(ex)

src/indexnotation/preprocessors.jl

@@ -119,10 +119,9 @@ end
 Group all scalar factors of a tensor expression into a single scalar factor at the start of the expression.
 """
 function groupscalarfactors(ex)
-    if isa(ex, Expr) # prewalk
-        ex = Expr(ex.head, map(groupscalarfactors, ex.args)...)
-    end
-    if istensorexpr(ex) && ex.args[1] == :*
+    if istensor(ex) || (isexpr(ex, :macrocall) && ex.args[1] == Symbol("@notensor"))
+        return ex
+    elseif istensorexpr(ex) && ex.args[1] == :*
         args = ex.args[2:end]
         scalarpos = findall(isscalarexpr, args)
         length(scalarpos) == 0 && return ex
@@ -133,6 +132,8 @@ function groupscalarfactors(ex)
             scalar = Expr(:call, :*, args[scalarpos]...)
         end
         return Expr(:call, :*, scalar, args[tensorpos]...)
+    elseif isa(ex, Expr)
+        return Expr(ex.head, map(groupscalarfactors, ex.args)...)
     end
     return ex
 end

src/indexnotation/verifiers.jl

@@ -125,7 +125,7 @@ function isscalarexpr(ex)
     elseif isexpr(ex, :call)
         return all(isscalarexpr, ex.args[2:end])
     else
-        return false
+        return true # assume everything else is valid scalar code
     end
 end
 

test/tensor.jl

@@ -418,6 +418,15 @@ using LinearAlgebra
         @test res ≈ -(mat1 * vec + mat2 * vec)
     end
 
+    @testset "Handling definitions" begin
+        vec = rand(2)
+        veccopy = copy(vec)
+        mat1 = rand(2, 2)
+        mat2 = rand(2, 2)
+        @tensor vec[a] := (mat1[a, b] * mat2[b, c]) * vec[c]
+        @test vec ≈ mat1 * mat2 * veccopy
+    end
+
     # @testset "Issue 136" begin
     #     A = rand(2, 2)
     #     B = rand(2, 2)