diff --git a/src/indexnotation/contractiontrees.jl b/src/indexnotation/contractiontrees.jl index 77c4ccb8..7f442f4a 100644 --- a/src/indexnotation/contractiontrees.jl +++ b/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) diff --git a/src/indexnotation/parser.jl b/src/indexnotation/parser.jl index 6bdd55ca..6afc0729 100644 --- a/src/indexnotation/parser.jl +++ b/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 diff --git a/src/indexnotation/verifiers.jl b/src/indexnotation/verifiers.jl index 77bec1da..948c9445 100644 --- a/src/indexnotation/verifiers.jl +++ b/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) diff --git a/test/auxiliary.jl b/test/auxiliary.jl index eb20ffc8..7e11ede7 100644 --- a/test/auxiliary.jl +++ b/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 diff --git a/test/cutensor.jl b/test/cutensor.jl index 338cda37..dd36f6d8 100644 --- a/test/cutensor.jl +++ b/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 diff --git a/test/tensor.jl b/test/tensor.jl index 7e97c362..60b0b390 100644 --- a/test/tensor.jl +++ b/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