more JET happiness

Author: lkdvos

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/fusiontrees/fusiontrees.jl

@@ -334,8 +334,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/pullbacks/tensoroperations.jl

@@ -1,5 +1,5 @@
 function blas_contract_pullback_ΔA!(
-        ΔA, ΔC, A, pA, B, pB, pAB, α, backend, allocator
+        ΔA, ΔC, A::AbstractTensorMap, pA, B::AbstractTensorMap, pB, pAB, α, backend, allocator
     )
     ipAB = invperm(linearize(pAB))
     pΔC = TO.repartition(ipAB, TO.numout(pA))
@@ -23,7 +23,7 @@ function blas_contract_pullback_ΔA!(
 end
 
 function blas_contract_pullback_ΔB!(
-        ΔB, ΔC, A, pA, B, pB, pAB, α, backend, allocator
+        ΔB, ΔC, A::AbstractTensorMap, pA, B::AbstractTensorMap, pB, pAB, α, backend, allocator
     )
     ipAB = invperm(linearize(pAB))
     pΔC = TO.repartition(ipAB, TO.numout(pA))
@@ -48,7 +48,7 @@ function blas_contract_pullback_ΔB!(
 end
 
 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/spaces/productspace.jl

@@ -146,7 +146,7 @@ that make up the `ProductSpace` instance.
 """
 function blocksectors(P::ProductSpace{S, N}) where {S, N}
     I = sectortype(S)
-    if I == Trivial
+    if I === Trivial
         return OneOrNoneIterator(dim(P) != 0, Trivial())
     end
     bs = Vector{I}()

src/tensors/abstracttensor.jl

@@ -237,7 +237,11 @@ 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
+# `typeintersect(Type{<:AbstractTensorMap}, Type{<:Union{FusionTreePair, FusionTreeBlock}})` is
+# `Type{Union{}}`, which leaves the parametric methods mutually ambiguous. Resolve it explicitly so
+# the intended `MethodError` is thrown instead of an ambiguity error.
+numout(::Type{Union{}}) = throw(MethodError(numout, (Union{},)))
 numout(::Type{<:AbstractTensorMap{T, S, N₁}}) where {T, S, N₁} = N₁
 
 @doc """
@@ -251,7 +255,8 @@ 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{Union{}}) = throw(MethodError(numin, (Union{},))) # see `numout(::Type{Union{}})`
 numin(::Type{<:AbstractTensorMap{T, S, N₁, N₂}}) where {T, S, N₁, N₂} = N₂
 
 """

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)

test/other/jet.jl

@@ -2,9 +2,23 @@ using TensorKit
 using Test
 using JET
 
+# Reports that are not TensorKit's to fix. Each entry needs an upstream issue.
+const IGNORED = (
+    # `schur_full`/`schur_vals` route `Diagonal` inputs to `DiagonalAlgorithm`, which schur does
+    # not implement, so `schur_full(::DiagonalTensorMap)` throws a `MethodError`. This reproduces
+    # with a plain `LinearAlgebra.Diagonal` and no TensorKit involved:
+    # https://github.com/QuantumKitHub/MatrixAlgebraKit.jl/issues/276
+    JET.LastFrameMethod(:schur_full!),
+    JET.LastFrameMethod(:schur_vals!),
+    # `local variable kwargs may be undefined` inside `GenericTreeTransformer`, coming entirely
+    # from Base's `@debug` expansion (`local msg, kwargs` bound in a short-circuit guard) rather
+    # than from any TensorKit code: https://github.com/aviatesk/JET.jl/issues/860
+    JET.LastFrameMethod(:GenericTreeTransformer),
+)
+
 # also gated in runtests.jl; JET < 0.12 does not define `JET_AVAILABLE`
 if isdefined(JET, :JET_AVAILABLE) && JET.JET_AVAILABLE
-    JET.test_package(TensorKit; target_modules = (TensorKit,))
+    JET.test_package(TensorKit; target_modules = (TensorKit,), ignored_modules = IGNORED)
 else
     @info "Full JET functionality is unavailable on Julia $VERSION; skipping JET analysis"
 end