remove type instabilities

Author: hakkelt

FFTWOperators/src/IRDFT.jl

@@ -82,4 +82,7 @@ is_invertible(L::IRDFT) = true
 is_full_row_rank(L::IRDFT) = true
 
 has_fast_opnorm(::IRDFT) = true
-LinearAlgebra.opnorm(L::IRDFT{T}) where {T} = sqrt(prod(L.dim_out) / 2)
+function LinearAlgebra.opnorm(L::IRDFT{T}) where {T}
+    @assert length(L.dim_out) > 0
+    return sqrt(prod(L.dim_out) / 2)
+end

NFFTOperators/src/NFFTOp.jl

@@ -169,6 +169,19 @@ function create_plan(trajectory, image_size, threaded; kwargs...)
     end
 end
 
+# Helper to create matched forward/backward FFT plans so that JET can track
+# that both plans have the same element type T and dimension D.
+struct _MatchedFFTPlans{T, D}
+    forward::FFTW.cFFTWPlan{Complex{T}, -1, true, D, UnitRange{Int64}}
+    backward::FFTW.cFFTWPlan{Complex{T}, 1, true, D, UnitRange{Int64}}
+end
+
+function _make_matched_fft_plans(tmpVec::Array{Complex{T}, D}, dims_; kwargs...) where {T, D}
+    FP = FFTW.plan_fft!(tmpVec, dims_; kwargs...)::FFTW.cFFTWPlan{Complex{T}, -1, true, D, UnitRange{Int64}}
+    BP = FFTW.plan_bfft!(tmpVec, dims_; kwargs...)::FFTW.cFFTWPlan{Complex{T}, 1, true, D, UnitRange{Int64}}
+    return _MatchedFFTPlans{T, D}(FP, BP)
+end
+
 function NFFTPlan(
         k::Matrix{T},
         N::NTuple{D, Int};
@@ -187,8 +200,9 @@ function NFFTPlan(
     tmpVec = Array{Complex{T}, D}(undef, Ñ)
 
     fftflags_ = (fftflags !== nothing) ? (flags = fftflags,) : NamedTuple()
-    FP = FFTW.plan_fft!(tmpVec, dims_; num_threads = FFTW.get_num_threads(), fftflags_...)
-    BP = FFTW.plan_bfft!(tmpVec, dims_; num_threads = FFTW.get_num_threads(), fftflags_...)
+    plans = _make_matched_fft_plans(tmpVec, dims_; num_threads = FFTW.get_num_threads(), fftflags_...)
+    FP = plans.forward
+    BP = plans.backward
 
     calcBlocks =
         (

src/calculus/Ax_mul_Bx.jl

@@ -33,7 +33,7 @@ struct Ax_mul_Bx{
     bufB::C
     bufC::C
     bufD::D
-    function Ax_mul_Bx(A::L1, B::L2, bufA::C, bufB::C, bufC::C, bufD::D) where {L1, L2, C, D}
+    function Ax_mul_Bx(A::L1, B::L2, bufA::C, bufB::C, bufC::C, bufD::D) where {L1 <: AbstractOperator, L2 <: AbstractOperator, C <: AbstractArray, D <: AbstractArray}
         if ndims(A, 1) != 2 || size(A, 2) != size(B, 2) || size(A, 1)[2] != size(B, 1)[1]
             throw(DimensionMismatch("Cannot compose operators"))
         end

src/calculus/Ax_mul_Bxt.jl

@@ -33,7 +33,7 @@ struct Ax_mul_Bxt{
     bufB::C
     bufC::C
     bufD::D
-    function Ax_mul_Bxt(A::L1, B::L2, bufA::C, bufB::C, bufC::C, bufD::D) where {L1, L2, C, D}
+    function Ax_mul_Bxt(A::L1, B::L2, bufA::C, bufB::C, bufC::C, bufD::D) where {L1 <: AbstractOperator, L2 <: AbstractOperator, C <: AbstractArray, D <: AbstractArray}
         if ndims(A, 1) == 1
             if size(A) != size(B)
                 throw(DimensionMismatch("Cannot compose operators"))

src/calculus/Axt_mul_Bx.jl

@@ -33,7 +33,7 @@ struct Axt_mul_Bx{
     bufB::C
     bufC::C
     bufD::D
-    function Axt_mul_Bx(A::L1, B::L2, bufA::C, bufB::C, bufC::C, bufD::D) where {L1, L2, C, D}
+    function Axt_mul_Bx(A::L1, B::L2, bufA::C, bufB::C, bufC::C, bufD::D) where {L1 <: AbstractOperator, L2 <: AbstractOperator, C <: AbstractArray, D <: AbstractArray}
         if ndims(A, 1) == 1
             if size(A) != size(B)
                 throw(DimensionMismatch("Cannot compose operators"))

src/calculus/Compose.jl

@@ -31,6 +31,7 @@ struct Compose{N, M, L <: Tuple, T <: Tuple} <: AbstractOperator
         end
         # check for adjacent operators that can be combined
         i = 1
+        local new_op::AbstractOperator = A[1]  # placeholder; set below when should_be_combined is true
         while i < length(A)
             should_be_combined = false
             triple_combination = false
@@ -292,6 +293,7 @@ end
 # utils
 function permute(C::Compose, p::AbstractVector{Int})
     i = findfirst(x -> ndoms(x, 2) > 1, C.A)
+    i === nothing && throw(ArgumentError("No operator with multiple domain dimensions found in Compose"))
     P = permute(C.A[i], p)
     AA = (C.A[1:(i - 1)]..., P, C.A[(i + 1):end]...)
     return Compose(AA, C.buf)

src/calculus/HCAT.jl

@@ -75,7 +75,7 @@ function HCAT(A::Vararg{AbstractOperator})
             end
         end
         # use buffer from HCAT in A
-        buf = A[findfirst((<:).(typeof.(A), HCAT))].buf
+        buf = A[findfirst((<:).(typeof.(A), HCAT))::Int].buf
     else
         AA = A
         # generate buffer

src/calculus/VCAT.jl

@@ -65,7 +65,7 @@ function VCAT(A::Vararg{AbstractOperator})
             end
         end
         # use buffer from VCAT in A
-        buf = A[findfirst((<:).(typeof.(A), VCAT))].buf
+        buf = A[findfirst((<:).(typeof.(A), VCAT))::Int].buf
     else
         AA = A
         # generate buffer