Enforce strict Einstein index rules in @tensor

src/indexnotation/contractiontrees.jl

@@ -94,10 +94,8 @@ function insertcontractiontrees!(
     end
     args = ex.args[2:end]
     network = map(getindices, args)
-    for a in getallindices(ex)
-        count(a in n for n in network) <= 2 ||
-            throw(ArgumentError("index $a appears more than twice in tensor contraction: $ex"))
-    end
+    # the index labels have already been verified by `verifyindices`, such that every label
+    # appears at most twice within this term
     tree = treebuilder(network)
     treeex = treesorter(args, tree, depth)
 

src/indexnotation/parser.jl

@@ -7,7 +7,8 @@ mutable struct TensorParser
     postprocessors::Vector{Any}
     function TensorParser()
         preprocessors = [
-            normalizeindices, expandconj, nconindexcompletion, extracttensorobjects,
+            normalizeindices, verifyindices, expandconj, nconindexcompletion,
+            extracttensorobjects,
         ]
         contractiontreebuilder = defaulttreebuilder
         contractiontreesorter = defaulttreesorter

src/indexnotation/verifiers.jl

@@ -181,6 +181,119 @@ function istensorexpr(ex)
     return false
 end
 
+"""
+    verifyindices(ex) -> ex
+
+Verify that all index labels in `ex` obey the strict Einstein summation convention, and throw an `ArgumentError` if not.
+This convention entails that within a single term, every index label should appear either once (an open index) or exactly twice (a contracted index, either between two different tensors or within a single tensor as a trace).
+Parentheses group the factors of a term and thus determine the contraction order, but do not introduce a new scope for the index labels.
+Different terms of a sum, different statements, and the argument of an explicit `tensorscalar` call do constitute separate scopes, in which labels can be reused freely.
+
+This routine expects the indices to be normalized, i.e. it should be called after [`normalizeindices`](@ref).
+The expression is returned unchanged, such that this can be used as a preprocessor.
+"""
+function verifyindices(ex)
+    if isexpr(ex, :macrocall) && ex.args[1] == Symbol("@notensor")
+        return ex
+    elseif istensor(ex) || istensorexpr(ex)
+        _indexscope(ex)
+    elseif isa(ex, Expr)
+        foreach(verifyindices, ex.args)
+    end
+    return ex
+end
+
+# analyze a single term and return its open indices, along with the indices that have been
+# contracted (closed) within that term; throws if the Einstein convention is violated
+function _indexscope(ex)
+    if istensor(ex)
+        _, leftind, rightind = decomposetensor(ex)
+        allind = vcat(leftind, rightind)
+        open, closed = Any[], Any[]
+        for label in unique(allind)
+            n = count(isequal(label), allind)
+            if n == 1
+                push!(open, label)
+            elseif n == 2
+                push!(closed, label)
+            else
+                throw(ArgumentError("@tensor: index $label appears $n times in tensor $ex"))
+            end
+        end
+        return open, closed
+    elseif isexpr(ex, :call)
+        # note: an explicit `tensorscalar(...)` call is not an `istensorexpr`, so it never
+        # reaches here; it is verified as a separate scope through the scalar factor
+        # fallback in `_indexscope_product`
+        if (ex.args[1] == :+ || ex.args[1] == :-) && length(ex.args) > 2
+            return _indexscope_sum(ex)
+        elseif ex.args[1] == :*
+            return _indexscope_product(ex)
+        elseif ex.args[1] == :/ && length(ex.args) == 3
+            verifyindices(ex.args[3])
+            return _indexscope(ex.args[2])
+        elseif ex.args[1] == :\ && length(ex.args) == 3
+            verifyindices(ex.args[2])
+            return _indexscope(ex.args[3])
+        elseif length(ex.args) == 2 # unary plus or minus, conj, adjoint, ...
+            return _indexscope(ex.args[2])
+        end
+    elseif isexpr(ex, prime) && length(ex.args) == 1
+        return _indexscope(ex.args[1])
+    end
+    return Any[], Any[]
+end
+
+function _indexscope_sum(ex)
+    # unlike a product, a sum cannot carry scalar terms, since `istensorexpr` requires all
+    # of its terms to be tensor expressions
+    open = nothing
+    for term in ex.args[2:end]
+        openterm, = _indexscope(term)
+        if isnothing(open)
+            open = openterm
+        elseif Set(openterm) != Set(open)
+            throw(
+                ArgumentError(
+                    "@tensor: non-matching indices $(tuple(open...)) and $(tuple(openterm...)) between terms of $ex"
+                )
+            )
+        end
+    end
+    # the contracted labels of the individual terms are not visible outside of that term
+    return open, Any[]
+end
+
+function _indexscope_product(ex)
+    opens, closeds = Any[], Any[]
+    for factor in ex.args[2:end]
+        if !istensorexpr(factor)
+            verifyindices(factor) # scalar factor: cannot contribute index labels
+            continue
+        end
+        openfactor, closedfactor = _indexscope(factor)
+        append!(opens, openfactor)
+        append!(closeds, closedfactor)
+    end
+    # labels that are already contracted cannot be reused within this term
+    for label in unique(closeds)
+        n = 2 * count(isequal(label), closeds) + count(isequal(label), opens)
+        n > 2 && throw(ArgumentError("@tensor: index $label appears $n times in $ex"))
+    end
+    open, closed = Any[], closeds
+    for label in unique(opens)
+        n = count(isequal(label), opens)
+        if n == 1
+            push!(open, label)
+        elseif n == 2
+            push!(closed, label)
+        else
+            throw(ArgumentError("@tensor: index $label appears $n times in $ex"))
+        end
+    end
+    return open, closed
+end
+
 """
     isassignment(ex)
 

test/auxiliary.jl

@@ -118,6 +118,82 @@
             :(promote_add(promote_contract(scalartype(a), scalartype(b)), scalartype(c)))
     end
 
+    @testset "index verification" begin
+        using TensorOperations: verifyindices, normalizeindices
+
+        # every label appears once or exactly twice within a single term; different terms of
+        # a sum, different statements and explicit `tensorscalar` calls are separate scopes
+        valid = (
+            :(a[-1, -2] := b[-1, 1] * c[1, -2]),
+            :(a[-1, -2] := b[-1, 1] * c[1, -2] + d[-1, 1] * e[1, -2]),
+            :(a[-1, -2] := b[-1, 1] * (c[1, 2] + d[1, 2]) * e[2, -2]),
+            :(a[-1, -2] := b[-1, 1] * c[1, -2] * (d[2, 2])),
+            :(a[x, y] := b[x, y] - 2 * c[y, x]),
+            :(a[-1] := b[-1, 1, 1]),
+            :(a[x, y] := b[x, z] * tensorscalar(c[z, z]) * d[z, y]),
+            :(a[x, y] := b[x′, x] * c[x′, y]),
+            # the divisor of `/` and `\` is a separate scope, but the tensor operand is not
+            :(a[x, y] := b[x, z] * c[z, y] / tensorscalar(d[z, z])),
+            :(a[x, y] := tensorscalar(d[z, z]) \ (b[x, z] * c[z, y])),
+            :(a[x, y] := b[x, z] * c[z, y] / alpha),
+            :(a[x, y] := 2 \ (b[x, z] * c[z, y])),
+            # conjugation, adjoint and unary minus pass the labels through unchanged
+            :(a[x, y] := conj(b[x, z]) * c[z, y]),
+            :(a[x, y] := (b[x, z])' * c[z, y]),
+            :(a[x, y] := -(b[x, z] * c[z, y])),
+            # `@notensor` is not index-verified at all
+            :(a[-1, -2] := (@notensor f(1)) * b[-1, -2]),
+            :(a[-1, -2] := (@notensor g[1, 1, 1]) * b[-1, -2]),
+            quote
+                @notensor b = c[1, 1, 1]
+                a[-1, -2] := d[-1, 1] * e[1, -2]
+            end,
+            quote
+                a[-1, -2] := b[-1, 1] * c[1, -2]
+                d[-1, -2] := e[-1, 1] * f[1, -2]
+            end,
+        )
+        for ex in valid
+            @test verifyindices(normalizeindices(ex)) isa Expr
+        end
+
+        # a label consumed by a trace or contraction cannot be reused within the same term,
+        # and all terms of a sum should have matching open indices
+        invalid = (
+            :(a[-1, -2, -3; -4, -5] := b[-2, 1] * c[1, -3; 1, -5] * d[-1, 1; -4]), # issue #288
+            :(a[-1, -2] := b[-1, 1; 1, -2] * c[1, 1]),
+            :(a[-1, -2] := b[-1, 1] * c[1, -2] * d[1, 1, -3, -3]),
+            :(a[-1, -2] := b[-1, 1] * c[1, -2] * (d[1, 1])),
+            # parentheses do not introduce a new scope for the labels
+            :(a[-1, -2] := (b[-1, 1] * c[1, -3]) * (d[-3, 1] * e[1, -2])),
+            :(a[-1, -2] := b[-1, 1] * c[1, -3] * d[-3, 1] * e[1, -2]),
+            :(a[-1] := b[1, 1, 1]),
+            :(a[-1, -2] := b[-1, -2] + c[-1, 1]),
+            :(a[x, y] := b[x′, x] * c[x′, y] * d[x', 1]), # primes are equal after normalizing
+            # the divisor of `/` and `\` is verified in its own scope
+            :(a[x, y] := b[x, y] / tensorscalar(c[z, z, z])),
+            :(a[x, y] := tensorscalar(c[z, z, z]) \ b[x, y]),
+            # ... but the labels of the tensor operand do escape to the enclosing term
+            :(a[-1] := (b[-1, 1, 1] / 2) * c[1, 1]),
+            :(a[-1] := (2 \ b[-1, 1, 1]) * c[1, 1]),
+            # the argument of an explicit `tensorscalar` is still verified
+            :(a[x, y] := b[x, z] * tensorscalar(c[z, z, z]) * d[z, y]),
+            :(a[x, y] := b[x, y] * tensorscalar(c[z, z] + d[z, w])),
+            :(alpha = tensorscalar(c[z, z, z])),
+        )
+        for ex in invalid
+            @test_throws ArgumentError verifyindices(normalizeindices(ex))
+        end
+
+        # errors are thrown at macro expansion time
+        @test_throws ArgumentError @macroexpand(
+            @tensor T[-1, -2, -3; -4, -5] := A[-2, 1] * B[1, -3; 1, -5] * C[-1, 1; -4]
+        )
+        @test_throws ArgumentError @macroexpand(
+            @tensoropt T[-1, -2] := A[-1, 1] * B[1, -2] * C[1, 1]
+        )
+    end
+
     @testset "parsecost" begin
         using TensorOperations: parsecost, Power
         @test parsecost(:(3 / 5)) === 3 / 5

test/cutensor.jl

@@ -280,15 +280,15 @@ if cuTENSOR.functional()
             @tensor C[3, 1, 4, 2] = β * C[3, 1, 4, 2] + α * A[1, 2, 3, 4]
             Ccopy = β * Ccopy + α * Acopy
             @test copy(C) ≈ Ccopy
-            @test_throws IndexError begin
+            @test_throws ArgumentError @macroexpand(
                 @tensor C[3, 1, 4, 2] = 0.5 * C[3, 1, 4, 2] + 1.2 * A[1, 2, 3]
-            end
+            )
             @test_throws CUTENSORError begin
                 @tensor C[3, 1, 4, 2] = 0.5 * C[3, 1, 4, 2] + 1.2 * A[3, 1, 2, 4]
             end
-            @test_throws IndexError begin
+            @test_throws ArgumentError @macroexpand(
                 @tensor C[1, 1, 2, 3] = 0.5 * C[1, 1, 2, 3] + 1.2 * A[1, 2, 3, 4]
-            end
+            )
         end
 
         @testset "views 3" begin
@@ -310,9 +310,7 @@ if cuTENSOR.functional()
             @test_throws CUTENSORError begin
                 @tensor B[c, b] += α * A[a, b, c, a]
             end
-            @test_throws IndexError begin
-                @tensor B[c, b] += α * A[a, b, a, a]
-            end
+            @test_throws ArgumentError @macroexpand(@tensor B[c, b] += α * A[a, b, a, a])
             @test_throws CUTENSORError begin
                 @tensor B[c, b] += α * A[a, b, a, c]
             end

test/tensor.jl

@@ -125,15 +125,15 @@ backendlist = (BaseCopy(), BaseView(), StridedNative(), StridedBLAS())
         @tensor backend = b C[3, 1, 4, 2] = β * C[3, 1, 4, 2] + α * A[1, 2, 3, 4]
         Ccopy = β * Ccopy + α * Acopy
         @test C ≈ Ccopy
-        @test_throws IndexError begin
+        @test_throws ArgumentError @macroexpand(
             @tensor C[3, 1, 4, 2] = 0.5 * C[3, 1, 4, 2] + 1.2 * A[1, 2, 3]
-        end
+        )
         @test_throws DimensionMismatch begin
             @tensor C[3, 1, 4, 2] = 0.5 * C[3, 1, 4, 2] + 1.2 * A[3, 1, 2, 4]
         end
-        @test_throws IndexError begin
+        @test_throws ArgumentError @macroexpand(
             @tensor C[1, 1, 2, 3] = 0.5 * C[1, 1, 2, 3] + 1.2 * A[1, 2, 3, 4]
-        end
+        )
     end
 
     @testset "views 3" begin
@@ -155,9 +155,7 @@ backendlist = (BaseCopy(), BaseView(), StridedNative(), StridedBLAS())
         @test_throws DimensionMismatch begin
             @tensor B[c, b] += α * A[a, b, c, a]
         end
-        @test_throws IndexError begin
-            @tensor B[c, b] += α * A[a, b, a, a]
-        end
+        @test_throws ArgumentError @macroexpand(@tensor B[c, b] += α * A[a, b, a, a])
         @test_throws DimensionMismatch begin
             @tensor B[c, b] += α * A[a, b, a, c]
         end
@@ -593,3 +591,33 @@ end
         @test E ≈ c * (A * B + C * D)
     end
 end
+
+# these are independent of the backend, so they are only tested once
+@testset "notensor" begin
+    list = randn(Float64, (4, 4, 2))
+    A = randn(Float64, (4, 4))
+    B = randn(Float64, (4, 4))
+    D = randn(Float64, (3, 3))
+
+    @tensor begin
+        @notensor X = list[:, :, 1]
+        @notensor Y = list[:, :, 2]
+        C[a, c] := X[a, b] * Y[b, c]
+    end
+    @test C ≈ list[:, :, 1] * list[:, :, 2]
+
+    @tensor G[a, c] := (@notensor tr(D)) * A[a, b] * B[b, c]
+    @test G ≈ tr(D) * A * B
+end
+
+@testset "scalar division" begin
+    A = randn(Float64, (4, 4))
+    B = randn(Float64, (4, 4))
+    D = randn(Float64, (3, 3))
+
+    # the label `b` is reused within the divisor, which is a separate scope
+    @tensor E[a, c] := A[a, b] * B[b, c] / tensorscalar(D[b, b])
+    @test E ≈ (A * B) / tr(D)
+    @tensor F[a, c] := tensorscalar(D[b, b]) \ (A[a, b] * B[b, c])
+    @test F ≈ (A * B) / tr(D)
+end