Improve static analysis and add JET tests

src/kernels.jl

@@ -3,7 +3,7 @@ abstract type AbstractKernel end
 """SquaredExponentialKernel
 
 """
-struct SquaredExponentialKernel{T} <: AbstractKernel
+struct SquaredExponentialKernel{T <: Real} <: AbstractKernel
     σ::T
 end
 
@@ -27,7 +27,7 @@ end
 """Matern52Kernel
 
 """
-struct Matern52Kernel{T} <: AbstractKernel
+struct Matern52Kernel{T <: Real} <: AbstractKernel
     σ::T
     ρ::T
 end

src/problems.jl

@@ -12,3 +12,7 @@ end
 
 struct IVPSolver <: AbstractFastSolver
 end
+
+# Type aliases for convenience
+const BVP = BVPSolver
+const IVP = IVPSolver

test/Project.toml

@@ -1,6 +1,7 @@
 [deps]
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
 QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

test/jet.jl

@@ -0,0 +1,30 @@
+using JET
+
+@testset "JET static analysis" begin
+    @testset "SquaredExponentialKernel type stability" begin
+        k = SquaredExponentialKernel(1.0)
+
+        # Scalar inputs
+        @test_opt target_modules = (FastSolvers,) k(1.0, 2.0)
+        @test_opt target_modules = (FastSolvers,) k(1.0, 2.0, ∇())
+        @test_opt target_modules = (FastSolvers,) k(1.0, 2.0, Δ())
+
+        # Vector inputs
+        X = [1.0, 2.0, 3.0]
+        Y = [4.0, 5.0, 6.0]
+        @test_opt target_modules = (FastSolvers,) k(X, Y)
+        @test_opt target_modules = (FastSolvers,) k(X, Y, ∇())
+        @test_opt target_modules = (FastSolvers,) k(X, Y, Δ())
+    end
+
+    @testset "No errors in FastSolvers module" begin
+        # Use @test_call which is more focused on actual errors vs false positives
+        k = SquaredExponentialKernel(1.0)
+        X = [1.0, 2.0]
+        Y = [3.0, 4.0]
+
+        @test_call target_modules = (FastSolvers,) k(X, Y)
+        @test_call target_modules = (FastSolvers,) k(X, Y, ∇())
+        @test_call target_modules = (FastSolvers,) k(X, Y, Δ())
+    end
+end

test/runtests.jl

@@ -26,3 +26,7 @@ end
 @testset "Explicit Imports" begin
     include("explicit_imports.jl")
 end
+
+@testset "JET static analysis" begin
+    include("jet.jl")
+end