Update JET to 0.12 and run the package analysis in CI

ext/TensorKitGPUArraysExt.jl

@@ -195,7 +195,7 @@ struct DeviceGenericTreeTransformer{VO <: AbstractVector{Int}, DA <: DeviceUniqu
 end
 
 # strides of a dense array of shape `sz`
-_dense_strides(size::Dims) = (1, Base.front(cumprod(size))...)
+_dense_strides(size::Dims) = Base.front((1, cumprod(size)...))
 
 # `permute(Vsrc, p) == Vdst` is enforced when the transformer is built, so the permuted
 # source shape always matches `sz_dst` and the two views share Cartesian inds.

src/auxiliary/caches.jl

@@ -87,18 +87,16 @@ function _cached(mod::Module, ex)
     end
     fcall = ex.args[1]
     if Meta.isexpr(fcall, :where)
-        hasparams = true
         params = fcall.args[2:end]
         fcall = fcall.args[1]
     else
-        hasparams = false
+        params = nothing
     end
     if Meta.isexpr(fcall, :(::))
-        typed = true
         typeex = fcall.args[2]
         fcall = fcall.args[1]
     else
-        typed = false
+        typeex = nothing
     end
     Meta.isexpr(fcall, :call) || _cached_error("expected a function signature", fcall)
     fname = fcall.args[1]
@@ -117,14 +115,14 @@ function _cached(mod::Module, ex)
     # actual implenetation, with underscore name
     _fname = Symbol(:_, basename)
     _fcall = Expr(:call, _fname, fargs...)
-    if hasparams
+    if !isnothing(params)
         _fcall = Expr(:where, _fcall, params...)
     end
     _fex = Expr(:function, _fcall, _fbody)
 
     # implementation that chooses the cache style
     newfcall = fcall
-    if hasparams
+    if !isnothing(params)
         newfcall = Expr(:where, newfcall, params...)
     end
     cachestylevar = gensym(:cachestyle)
@@ -138,11 +136,11 @@ function _cached(mod::Module, ex)
 
     # nocache implementation
     fnocachecall = Expr(:call, fname, fargs..., :(::$NoCache))
-    if hasparams
+    if !isnothing(params)
         fnocachecall = Expr(:where, fnocachecall, params...)
     end
     fnocachebody = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...))
-    if typed
+    if !isnothing(typeex)
         T = gensym(:T)
         fnocachebody = Expr(:block, Expr(:(=), T, typeex), Expr(:(::), fnocachebody, T))
     end
@@ -151,7 +149,7 @@ function _cached(mod::Module, ex)
     # tasklocal cache implementation
     Dvar = gensym(:D)
     flocalcachecall = Expr(:call, fname, fargs..., :(::$TaskLocalCache{$Dvar}))
-    if hasparams
+    if !isnothing(params)
         flocalcachecall = Expr(:where, flocalcachecall, params..., Dvar)
     else
         flocalcachecall = Expr(:where, flocalcachecall, Dvar)
@@ -172,7 +170,7 @@ function _cached(mod::Module, ex)
     end
     missex = _timeit_expr(misslabel, Expr(:call, _fname, fargnames...))
     getvalex = _timeit_expr(lookuplabel, :(get!(() -> $missex, $cachevar, $key)))
-    if typed
+    if !isnothing(typeex)
         T = gensym(:T)
         flocalcachebody = Expr(
             :block,
@@ -193,11 +191,11 @@ function _cached(mod::Module, ex)
 
     # # global cache implementation
     fglobalcachecall = Expr(:call, fname, fargs..., :(::$GlobalLRUCache))
-    if hasparams
+    if !isnothing(params)
         fglobalcachecall = Expr(:where, fglobalcachecall, params...)
     end
     getglobalcachex = Expr(:(=), cachevar, globalcachename)
-    if typed
+    if !isnothing(typeex)
         T = gensym(:T)
         fglobalcachebody = Expr(
             :block,

src/auxiliary/iterators.jl

@@ -4,7 +4,7 @@ struct OneOrNoneIterator{T}
 end
 
 function Base.iterate(it::OneOrNoneIterator, state = true)
-    if state && it.cond
+    if isone(state) && it.cond
         return (it.first, false)
     else
         return nothing

src/factorizations/matrixalgebrakit.jl

@@ -103,7 +103,6 @@ MAK.zero!(t::AbstractTensorMap) = zerovector!(t)
 for f in [
         :lq_full, :lq_compact, :lq_null,
         :qr_full, :qr_compact, :qr_null,
-        :schur_full, :schur_vals,
         :eig_full, :eig_vals, :eig_trunc, :eig_trunc_no_error,
         :eigh_full, :eigh_vals, :eigh_trunc, :eigh_trunc_no_error,
         :svd_full, :svd_compact, :svd_trunc, :svd_trunc_no_error, :svd_vals,

src/factorizations/pullbacks.jl

@@ -1,5 +1,5 @@
 nothing_or_block(x, c) = isnothing(x) ? x : block(x, c)
-nothing_or_block(x::Diagonal, c) = block(MAK.diagonal(parent(x)), c)
+nothing_or_block(x::Diagonal{T, <:SectorVector{T}}, c) where {T} = block(MAK.diagonal(parent(x)), c)
 
 for pullback! in (
         :qr_pullback!, :lq_pullback!, :left_polar_pullback!, :right_polar_pullback!,

src/fusiontrees/braiding_manipulations.jl

@@ -332,7 +332,7 @@ end
         if N₂ == 0
             return dst => U
         else
-            dst, U_tmp = repartition(dst, N₁)
+            dst, U_tmp = repartition(dst::FusionTreeBlock{I}, N₁)
             U = U_tmp * U
             return dst => U
         end

src/fusiontrees/duality_manipulations.jl

@@ -657,7 +657,7 @@ function planar_trace(f::FusionTree, (q₁, q₂)::Index2Tuple)
     # We thus handle the total trace recursively, by first looking for and
     # tracing away neighbouring pairs.
     k = 1
-    local i, j
+    i = j = 0
     while k <= length(q₁)
         if mod1(q₁[k] + 1, length(f)) == q₂[k]
             i = q₁[k]

src/fusiontrees/fusiontrees.jl

@@ -338,8 +338,7 @@ function fusiontensor((f₁, f₂)::FusionTreePair)
     d1 = TupleTools.front(sz1)
     d2 = TupleTools.front(sz2)
     return reshape(
-        reshape(F₁, TupleTools.prod(d1), sz1[end]) *
-            reshape(F₂, TupleTools.prod(d2), sz2[end])', (d1..., d2...)
+        reshape(F₁, :, sz1[end]) * reshape(F₂, :, sz2[end])', (d1..., d2...)
     )
 end
 fusiontensor(src::FusionTreeBlock) = sum(fusiontensor, fusiontrees(src))

src/planar/macros.jl

@@ -56,7 +56,7 @@ function planarparser(planarexpr, kwargs...)
         if name == :order
             isexpr(val, :tuple) ||
                 throw(ArgumentError("Invalid use of `order`, should be `order=(...,)`"))
-            indexorder = map(normalizeindex, val.args)
+            indexorder = map(TO.normalizeindex, val.args)
             parser.contractiontreebuilder = network -> TO.indexordertree(
                 network, indexorder
             )

src/planar/preprocessors.jl

@@ -180,22 +180,18 @@ function _construct_braidingtensors!(ex, preargs, indexmap, non_braiding) # ex i
                 i1a, i2a, = rightind
             end
 
-            foundV1, foundV2 = false, false
+            V1 = V2 = nothing
             if haskey(indexmap, i1a)
                 V1 = indexmap[i1a]
-                foundV1 = true
             elseif haskey(indexmap, i1b)
                 V1 = Expr(:call, :dual, indexmap[i1b])
-                foundV1 = true
             end
             if haskey(indexmap, i2a)
                 V2 = indexmap[i2a]
-                foundV2 = true
             elseif haskey(indexmap, i2b)
                 V2 = Expr(:call, :dual, indexmap[i2b])
-                foundV2 = true
             end
-            if foundV1 && foundV2
+            if !isnothing(V1) && !isnothing(V2)
                 s = gensym(:τ)
                 scalarex = Expr(
                     :call, :promote_type,
@@ -609,7 +605,7 @@ function _extract_contraction_pairs(rhs, lhs, pre, temporaries)
         newarg = _extract_contraction_pairs(rhs.args[3], lhs, pre, temporaries)
         return Expr(:call, :\, rhs.args[2], newarg)
     else
-        throw(ArgumentError("unknown tensor expression $ex"))
+        throw(ArgumentError("unknown tensor expression $rhs"))
     end
 end
 

src/pullbacks/tensoroperations.jl

@@ -47,8 +47,11 @@ function blas_contract_pullback_ΔB!(
     return nothing
 end
 
+# Only the primal `A` is annotated: it fixes the element and space types of the workspace `E`
+# below, which `twist!` needs. The tangents stay untyped because their representation differs
+# per AD backend (Mooncake passes a tangent object, Enzyme the `.dval` shadow).
 function trace_permute_pullback_ΔA!(
-        ΔA, ΔC, A, p, q, α, backend
+        ΔA, ΔC, A::AbstractTensorMap, p, q, α, backend
     )
     ip = invperm((linearize(p)..., q[1]..., q[2]...))
     pdA = TO.repartition(ip, numout(A))

src/tensors/abstracttensor.jl

@@ -238,7 +238,7 @@ See also [`numin`](@ref) and [`numind`](@ref).
 """ numout
 
 numout(x) = numout(typeof(x))
-numout(T::Type) = throw(MethodError(numout, T)) # avoid infinite recursion
+numout(T::Type) = throw(MethodError(numout, (T,))) # avoid infinite recursion
 numout(::Type{<:AbstractTensorMap{T, S, N₁}}) where {T, S, N₁} = N₁
 
 @doc """
@@ -252,7 +252,7 @@ See also [`numout`](@ref) and [`numind`](@ref).
 """ numin
 
 numin(x) = numin(typeof(x))
-numin(T::Type) = throw(MethodError(numin, T)) # avoid infinite recursion
+numin(T::Type) = throw(MethodError(numin, (T,))) # avoid infinite recursion
 numin(::Type{<:AbstractTensorMap{T, S, N₁, N₂}}) where {T, S, N₁, N₂} = N₂
 
 """
@@ -314,7 +314,6 @@ InnerProductStyle(t::AbstractTensorMap) = InnerProductStyle(typeof(t))
 
 numout(t::AbstractTensorMap) = numout(typeof(t))
 numin(t::AbstractTensorMap) = numin(typeof(t))
-numind(t::AbstractTensorMap) = numind(typeof(t))
 
 # tensor characteristics: data structure and properties
 #------------------------------------------------------

src/tensors/adjoint.jl

@@ -13,6 +13,7 @@ end
 Base.parent(t::AdjointTensorMap) = t.parent
 parenttype(t::AdjointTensorMap) = parenttype(typeof(t))
 parenttype(::Type{AdjointTensorMap{T, S, N₁, N₂, TT}}) where {T, S, N₁, N₂, TT} = TT
+parenttype(::Type{Union{}}) = throw(MethodError(parenttype, (Union{},)))
 
 # Constructor: construct from taking adjoint of a tensor
 Base.adjoint(t::AdjointTensorMap) = parent(t)

src/tensors/indexmanipulations.jl

@@ -219,7 +219,7 @@ See also [`permute`](@ref) for creating a new tensor.
         α::Number = One(), β::Number = Zero(),
         backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator()
     )
-    @boundscheck spacecheck_transform(permute, tdst, tsrc, p)
+    @boundscheck spacecheck_transform(permute, space(tdst), space(tsrc), p)
     @timeit_debug GLOBAL_TIMER "permute!/braid!" begin
         tdst′, tsrc′, p′, _, conjsrc′, α′, β′ = unwrap_adjoints(tdst, tsrc, p, nothing, false, α, β)
         @inbounds _braid!(tdst′, tsrc′, p′, conjsrc′, allind(tsrc′), α′, β′, backend, allocator)
@@ -309,7 +309,7 @@ See also [`braid`](@ref) for creating a new tensor.
         α::Number = One(), β::Number = Zero(),
         backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator()
     )
-    @boundscheck spacecheck_transform(braid, tdst, tsrc, p, levels)
+    @boundscheck spacecheck_transform(braid, space(tdst), space(tsrc), p, levels)
     @timeit_debug GLOBAL_TIMER "permute!/braid!" begin
         tdst′, tsrc′, p′, levels′, conjsrc′, α′, β′ = unwrap_adjoints(tdst, tsrc, p, levels, false, α, β)
         @inbounds _braid!(tdst′, tsrc′, p′, conjsrc′, levels′, α′, β′, backend, allocator)
@@ -382,7 +382,7 @@ end
         α::Number = One(), β::Number = Zero(),
         backend::AbstractBackend = TO.DefaultBackend(), allocator = TO.DefaultAllocator()
     )
-    @boundscheck spacecheck_transform(transpose, tdst, tsrc, p)
+    @boundscheck spacecheck_transform(transpose, space(tdst), space(tsrc), p)
     @timeit_debug GLOBAL_TIMER "transpose!" begin
         tdst′, tsrc′, p′, _, conjsrc′, α′, β′ = unwrap_adjoints(tdst, tsrc, p, nothing, false, α, β)
         @inbounds _transpose!(tdst′, tsrc′, p′, conjsrc′, α′, β′, backend, allocator)
@@ -505,8 +505,6 @@ for (operation, manipulation) in (
     end
 end
 
-spacecheck_transform(f, tdst::AbstractTensorMap, tsrc::AbstractTensorMap, args...) =
-    spacecheck_transform(f, space(tdst), space(tsrc), args...)
 @noinline function spacecheck_transform(f, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple)
     check_spacetype(Vdst, Vsrc)
     f(Vsrc, p) == Vdst ||
@@ -604,19 +602,21 @@ end
 @propagate_inbounds function _braid!(
         tdst, tsrc, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple, α, β, backend, allocator
     )
-    @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc)
+    Vdst, Vsrc = space(tdst), space(tsrc)
+    @boundscheck spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc)
     all(has_array_view, (tdst, tsrc)) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator)
-    transformer = treebraider(tdst, tsrc, p, conjsrc, levels)
+    transformer = treebraider(storagetype(tdst), Vdst, Vsrc, p, conjsrc, levels)
     return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator)
 end
 
 # counterpart of `_braid!` for `transpose!`; the cyclicity of `p` is checked by the caller
 @propagate_inbounds function _transpose!(
         tdst, tsrc, p::Index2Tuple, conjsrc::Bool, α, β, backend, allocator
     )
-    @boundscheck spacecheck_transform(permute, tdst, tsrc, p, conjsrc)
+    Vdst, Vsrc = space(tdst), space(tsrc)
+    @boundscheck spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc)
     all(has_array_view, (tdst, tsrc)) && return _dense_transform!(tdst, tsrc, p, conjsrc, α, β, backend, allocator)
-    transformer = treetransposer(tdst, tsrc, p, conjsrc)
+    transformer = treetransposer(storagetype(tdst), Vdst, Vsrc, p, conjsrc)
     return add_transform!(tdst, tsrc, p, conjsrc, transformer, α, β, backend, allocator)
 end
 
@@ -737,7 +737,7 @@ function _recouple!(
         buffer::CPUStorage{Complex{R}}, buffer_dst::StridedView, buffer_src::StridedView, U::Matrix{R}
     ) where {R <: LinearAlgebra.BlasReal}
     rbuffer = reinterpret(R, buffer)
-    LinearAlgebra.BLAS.gemm!('N', 'T', one(R), _realview(rbuffer, buffer_src), U, zero(R), _realview(rbuffer, buffer_dst))
+    LinearAlgebra.BLAS.gemm!('N', 'T', true, _realview(rbuffer, buffer_src), U, false, _realview(rbuffer, buffer_dst))
     return nothing
 end
 # real view of the complex column-major view `b` into the reinterpreted buffer `rbuffer`

src/tensors/linalg.jl

@@ -586,7 +586,7 @@ tdst[sub_axes...] .= tsrc[sub_axes...]
 absorb(tdst::AbstractTensorMap, tsrc::AbstractTensorMap) = absorb!(copy(tdst), tsrc)
 function absorb!(tdst::AbstractTensorMap, tsrc::AbstractTensorMap)
     numin(tdst) == numin(tsrc) && numout(tdst) == numout(tsrc) ||
-        throw(DimensionError("Incompatible number of indices for source and destination"))
+        throw(IndexError("Incompatible number of indices for source and destination"))
     S = check_spacetype(tdst, tsrc)
     dom = mapreduce(infimum, ⊗, domain(tdst), domain(tsrc); init = one(S))
     cod = mapreduce(infimum, ⊗, codomain(tdst), codomain(tsrc); init = one(S))

src/tensors/tensoroperations.jl

@@ -236,7 +236,10 @@ function trace_permute!(
     return tdst
 end
 
-function _trace_permute!(::UniqueFusion, tdst, tsrc, (p₁, p₂), (q₁, q₂), α, β, backend)
+function _trace_permute!(
+        ::UniqueFusion, tdst::AbstractTensorMap, tsrc::AbstractTensorMap,
+        (p₁, p₂), (q₁, q₂), α, β, backend
+    )
     scale!(tdst, β)
     r₁, r₂ = (p₁..., q₁...), (p₂..., q₂...)
     N₁, N₂ = length(p₁), length(p₂)
@@ -261,7 +264,10 @@ function _trace_permute!(::UniqueFusion, tdst, tsrc, (p₁, p₂), (q₁, q₂),
     return tdst
 end
 
-function _trace_permute!(::FusionStyle, tdst, tsrc, (p₁, p₂), (q₁, q₂), α, β, backend)
+function _trace_permute!(
+        ::FusionStyle, tdst::AbstractTensorMap, tsrc::AbstractTensorMap,
+        (p₁, p₂), (q₁, q₂), α, β, backend
+    )
     scale!(tdst, β)
     r₁, r₂ = (p₁..., q₁...), (p₂..., q₂...)
     N₁, N₂ = length(p₁), length(p₂)

src/tensors/treetransformers.jl

@@ -68,7 +68,7 @@ struct GenericTreeTransformer{T, N} <: TreeTransformer
     structure_src::Vector{StridedStructure{N}}
 end
 
-function UniqueTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bool) where {A}
+function UniqueTreeTransformer(::Type{A}, transform, p, Vdst::HomSpace, Vsrc::HomSpace, conjsrc::Bool) where {A}
     t₀ = Base.time()
 
     spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc)
@@ -93,7 +93,7 @@ function UniqueTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Boo
     return transformer
 end
 
-function GenericTreeTransformer(::Type{A}, transform, p, Vdst, Vsrc, conjsrc::Bool) where {A}
+function GenericTreeTransformer(::Type{A}, transform, p, Vdst::HomSpace, Vsrc::HomSpace, conjsrc::Bool) where {A}
     t₀ = Base.time()
     spacecheck_transform(permute, Vdst, Vsrc, p, conjsrc)
     # the fusion blocks that are transformed are those of the adjoint space for a conjugated source
@@ -183,11 +183,6 @@ function TreeTransformer(
 end
 
 # braid is special because it has levels
-function treebraider(
-        tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple
-    )
-    return treebraider(storagetype(tdst), space(tdst), space(tsrc), p, conjsrc, levels)
-end
 @cached function treebraider(
         A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool, levels::IndexTuple
     )::treetransformertype(A, Vdst, Vsrc) where {TA}
@@ -199,9 +194,6 @@ end
     return TreeTransformer(A, fusiontreebraider, p, Vdst, Vsrc, conjsrc)
 end
 
-function treetransposer(tdst::AbstractTensorMap, tsrc::AbstractTensorMap, p::Index2Tuple, conjsrc::Bool)
-    return treetransposer(storagetype(tdst), space(tdst), space(tsrc), p, conjsrc)
-end
 @cached function treetransposer(
         A::Type{TA}, Vdst::TensorMapSpace, Vsrc::TensorMapSpace, p::Index2Tuple, conjsrc::Bool
     )::treetransformertype(A, Vdst, Vsrc) where {TA}

test/Project.toml

@@ -39,7 +39,7 @@ AllocCheck = "0.2"
 ChainRulesTestUtils = "1"
 Combinatorics = "1"
 cuTENSOR = "6"
-JET = "0.9, 0.10, 0.11"
+JET = "0.9, 0.12"
 ParallelTestRunner = "2"
 Test = "1"
 TestExtras = "0.2,0.3"