Add comprehensive type preservation tests for GPU unit tests

GeneralisedFilters/test/batch_kalman_test.jl

@@ -9,20 +9,21 @@
 
     using CUDA
 
+    T_elem = Float32  # Use Float32 for type preservation
     rng = StableRNG(1234)
     K = 10
     Dx = 2
     Dy = 2
-    μ0s = [rand(rng, Dx) for _ in 1:K]
-    Σ0s = [rand(rng, Dx, Dx) for _ in 1:K]
+    μ0s = [rand(rng, T_elem, Dx) for _ in 1:K]
+    Σ0s = [rand(rng, T_elem, Dx, Dx) for _ in 1:K]
     Σ0s .= Σ0s .* transpose.(Σ0s)
-    As = [rand(rng, Dx, Dx) for _ in 1:K]
-    bs = [rand(rng, Dx) for _ in 1:K]
-    Qs = [rand(rng, Dx, Dx) for _ in 1:K]
+    As = [rand(rng, T_elem, Dx, Dx) for _ in 1:K]
+    bs = [rand(rng, T_elem, Dx) for _ in 1:K]
+    Qs = [rand(rng, T_elem, Dx, Dx) for _ in 1:K]
     Qs .= Qs .* transpose.(Qs)
-    Hs = [rand(rng, Dy, Dx) for _ in 1:K]
-    cs = [rand(rng, Dy) for _ in 1:K]
-    Rs = [rand(rng, Dy, Dy) for _ in 1:K]
+    Hs = [rand(rng, T_elem, Dy, Dx) for _ in 1:K]
+    cs = [rand(rng, T_elem, Dy) for _ in 1:K]
+    Rs = [rand(rng, T_elem, Dy, Dy) for _ in 1:K]
     Rs .= Rs .* transpose.(Rs)
 
     models = [
@@ -32,7 +33,7 @@
     ]
 
     T = 5
-    Ys = [[rand(rng, Dy) for _ in 1:T] for _ in 1:K]
+    Ys = [[rand(rng, T_elem, Dy) for _ in 1:T] for _ in 1:K]
 
     outputs = [
         GeneralisedFilters.filter(rng, models[k], KalmanFilter(), Ys[k]) for k in 1:K
@@ -131,7 +132,7 @@
         BatchLinearGaussianObservations(Hs, cs, Rs),
     )
 
-    Ys_batch = Vector{Matrix{Float64}}(undef, T)
+    Ys_batch = Vector{Matrix{T_elem}}(undef, T)
     for t in 1:T
         Ys_batch[t] = stack(Ys[k][t] for k in 1:K)
     end
@@ -147,4 +148,18 @@
 
     @test Array(batch_output[2])[end] .≈ log_likelihoods[end] rtol = 1e-5
     @test Array(batch_output[1].μs) ≈ stack(getproperty.(states, :μ)) rtol = 1e-5
+    # Type preservation tests
+    @test eltype(batch_output[1].μs) == T_elem
+    @test eltype(batch_output[1].Σs) == T_elem
+    @test eltype(batch_output[2]) == T_elem
+    @test all(eltype(μ0) == T_elem for μ0 in μ0s)
+    @test all(eltype(Σ0) == T_elem for Σ0 in Σ0s)
+    @test all(eltype(A) == T_elem for A in As)
+    @test all(eltype(b) == T_elem for b in bs)
+    @test all(eltype(Q) == T_elem for Q in Qs)
+    @test all(eltype(H) == T_elem for H in Hs)
+    @test all(eltype(c) == T_elem for c in cs)
+    @test all(eltype(R) == T_elem for R in Rs)
+    @test all(all(eltype(y) == T_elem for y in Y) for Y in Ys)
+    @test all(eltype(Y_batch) == T_elem for Y_batch in Ys_batch)
 end

GeneralisedFilters/test/resamplers.jl

@@ -42,11 +42,12 @@ end
 
     SEED = 1234
     N = 10^6
+    T = Float32
 
     rng = CUDA.RNG(SEED)
 
-    xs = randn(rng, N)
-    ws = map(x -> pdf(Normal(1, 1), x) / pdf(Normal(0, 1), x), xs)
+    xs = randn(rng, T, N)
+    ws = map(x -> T(pdf(Normal(1, 1), x) / pdf(Normal(0, 1), x)), xs)
     ws ./= sum(ws)
 
     μ0 = sum(ws .* xs)
@@ -57,20 +58,38 @@ end
     @test length(idxs) == N
     μ1 = sum(xs[idxs]) / N
     @test μ0 ≈ μ1 rtol = 1e-1
+    # Type preservation tests
+    @test eltype(xs) == T
+    @test eltype(ws) == T
+    @test typeof(μ0) == T
+    @test typeof(μ1) == T
+    @test eltype(idxs) == Int
 end
 
 @testitem "Test GPU systematic resampling" setup = [GPUResamplingTestSetup] tags = [:gpu] begin
     idxs = GeneralisedFilters.sample_ancestors(rng, Systematic(), ws)
     @test length(idxs) == N
     μ1 = sum(xs[idxs]) / N
     @test μ0 ≈ μ1 rtol = 1e-1
+    # Type preservation tests
+    @test eltype(xs) == T
+    @test eltype(ws) == T
+    @test typeof(μ0) == T
+    @test typeof(μ1) == T
+    @test eltype(idxs) == Int
 end
 
 @testitem "Test GPU stratified resampling" setup = [GPUResamplingTestSetup] tags = [:gpu] begin
     idxs = GeneralisedFilters.sample_ancestors(rng, Stratified(), ws)
     @test length(idxs) == N
     μ1 = sum(xs[idxs]) / N
     @test μ0 ≈ μ1 rtol = 1e-1
+    # Type preservation tests
+    @test eltype(xs) == T
+    @test eltype(ws) == T
+    @test typeof(μ0) == T
+    @test typeof(μ1) == T
+    @test eltype(idxs) == Int
 end
 
 @testitem "Test GPU offspring-to-ancestors" tags = [:gpu] begin
@@ -79,6 +98,11 @@ end
     true_ancestors = CuVector{Int}([2, 2, 4, 5, 5])
     ancestors = GeneralisedFilters.offspring_to_ancestors(offspring)
     @test ancestors == true_ancestors
+    # Type preservation tests
+    @test eltype(offspring) == Int
+    @test eltype(true_ancestors) == Int
+    @test eltype(ancestors) == Int
+    @test ancestors isa CuVector{Int}
 end
 
 @testitem "Test GPU ancestors-to-offspring" tags = [:gpu] begin
@@ -87,4 +111,9 @@ end
     true_offspring = CuVector{Int}([1, 2, 1, 1, 0])
     offspring = GeneralisedFilters.ancestors_to_offspring(ancestors)
     @test offspring == true_offspring
+    # Type preservation tests
+    @test eltype(ancestors) == Int
+    @test eltype(true_offspring) == Int
+    @test eltype(offspring) == Int
+    @test offspring isa CuVector{Int}
 end

GeneralisedFilters/test/runtests.jl

@@ -304,7 +304,13 @@ end
     @test kf_ll ≈ ll rtol = 1e-2
     @test first(kf_state.μ) ≈ sum(xs[1, :] .* ws) rtol = 1e-1
     @test last(kf_state.μ) ≈ sum(zs.μs[end, :] .* ws) rtol = 1e-2
+    # Type preservation tests
     @test eltype(xs) == ET
+    @test eltype(zs.μs) == ET
+    @test eltype(zs.Σs) == ET
+    @test typeof(ll) == ET
+    @test typeof(kf_ll) == ET
+    @test eltype(ws) == ET
 end
 
 @testitem "RBPF ancestory test" begin
@@ -581,8 +587,14 @@ end
     rbpf = BatchRBPF(BatchKalmanFilter(N_particles), N_particles)
     states, ll = GeneralisedFilters.filter(hier_model, rbpf, ys; ref_state=ref_trajectory)
 
-    # Check returned type
+    # Type preservation tests
     @test typeof(ll) == T
+    @test eltype(states.particles.xs) == T
+    @test eltype(states.particles.zs.μs) == T
+    @test eltype(states.particles.zs.Σs) == T
+    @test eltype(states.log_weights) == T
+    @test all(eltype(ref_traj) == T for ref_traj in ref_trajectory)
+    @test all(eltype(y) == T for y in ys)
 end
 
 @testitem "GPU-RBPF ancestory test" tags = [:gpu] begin
@@ -624,6 +636,17 @@ end
     cb = GeneralisedFilters.ParallelAncestorCallback(tree)
     states, ll = GeneralisedFilters.filter(hier_model, rbpf, ys; callback=cb)
 
+    # Type preservation tests
+    @test typeof(ll) == T
+    @test eltype(states.particles.xs) == T
+    @test eltype(states.particles.zs.μs) == T
+    @test eltype(states.particles.zs.Σs) == T
+    @test eltype(states.log_weights) == T
+    @test eltype(tree.particles.xs) == T
+    @test eltype(tree.particles.zs.μs) == T
+    @test eltype(tree.particles.zs.Σs) == T
+    @test all(eltype(y) == T for y in ys)
+
     # TODO: add proper test comparing to dense storage
     ancestry = GeneralisedFilters.get_ancestry(tree, K)
 end
@@ -702,4 +725,19 @@ end
         @test state.μ[1] ≈ only(mean(x_trajectories)) rtol = 1e-1
         @test state.μ[2] ≈ only(mean(getproperty.(z_trajectories, :μs))) rtol = 1e-1
     end
+
+    # Type preservation tests
+    @test eltype(particle_template.xs) == T
+    @test eltype(particle_template.zs.μs) == T
+    @test eltype(particle_template.zs.Σs) == T
+    @test all(eltype(sample.xs) == T for sample in trajectory_samples if !isnothing(sample))
+    @test all(
+        eltype(getproperty(sample, :zs).μs) == T for
+        sample in trajectory_samples if !isnothing(sample)
+    )
+    @test all(
+        eltype(getproperty(sample, :zs).Σs) == T for
+        sample in trajectory_samples if !isnothing(sample)
+    )
+    @test all(eltype(y) == T for y in ys)
 end