Fix CI

Project.toml

@@ -17,13 +17,13 @@ StructArrays = "09ab397b-f2b6-538f-b94a-2f83cf4a842a"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
-AbstractGPs = "0.5.15"
+AbstractGPs = "0.5.17"
 Bessels = "0.2.8"
 BlockDiagonals = "0.1.7"
 ChainRulesCore = "1"
 FillArrays = "0.13.0 - 0.13.7"
 KernelFunctions = "0.9, 0.10.1"
 StaticArrays = "1"
 StructArrays = "0.5, 0.6"
-Zygote = "0.6"
+Zygote = "0.6.65"
 julia = "1.6"

README.md

@@ -8,10 +8,6 @@ TemporalGPs.jl is a tool to make Gaussian processes (GPs) defined using [Abstrac
 
 [JuliaCon 2020 Talk](https://www.youtube.com/watch?v=dysmEpX1QoE)
 
-# Dependency Status
-
-In the interest of managing expectations, please note that TemporalGPs does not currently operate with the most current version of AbstractGPs / Zygote / ChainRules. I (Will) am aware of this problem, and will sort it out as soon as I have the time!
-
 # Installation
 
 TemporalGPs.jl is registered, so simply type the following at the REPL:

examples/approx_space_time_learning.jl

@@ -41,7 +41,7 @@ end
 # Also works with RectilinearGrids of inputs.
 # Times must be increasing, points in space can be anywhere.
 N = 50;
-T = 1000;
+T = 100;
 points_in_space = [randn(N) for _ in 1:T];
 points_in_time = RegularSpacing(0.0, 0.1, T);
 x = RegularInTime(points_in_time, points_in_space);

examples/exact_space_time_learning.jl

@@ -39,7 +39,7 @@ end
 # Exact inference only works for such grids.
 # Times must be increasing, points in space can be anywhere.
 N = 50;
-T = 1_000;
+T = 100;
 points_in_space = collect(range(-3.0, 3.0; length=N));
 points_in_time = RegularSpacing(0.0, 0.01, T);
 x = RectilinearGrid(points_in_space, points_in_time);
@@ -73,7 +73,7 @@ final_params = unpack(training_results.minimizer)
 f_post = posterior(build_gp(final_params)(x, final_params.var_noise), y);
 
 # Specify some locations at which to make predictions.
-T_pr = 1200;
+T_pr = 120;
 points_in_time_pr = RegularSpacing(0.0, 0.01, T_pr);
 x_pr = RectilinearGrid(points_in_space, points_in_time_pr);
 

src/models/linear_gaussian_conditionals.jl

@@ -46,8 +46,10 @@ be equivalent to
 function predict(x::Gaussian, f::AbstractLGC)
     A, a, Q = get_fields(f)
     m, P = get_fields(x)
-    # Symmetric wrapper needed for numerical stability. Do not unwrap.
-    return Gaussian(A * m + a, (A * symmetric(P)) * A' + Q)
+    # Previously, there was a symmetric wrapper around `P` in the following line.
+    # It was removed to solve a performance problem with `Symmetric` around static
+    # arrays, with unclear implications on numerical stability.
+    return Gaussian(A * m + a, (A * P) * A' + Q)
 end
 
 """

src/space_time/pseudo_point.jl

@@ -103,10 +103,11 @@ function kernel_diagonals(k::DTCSeparable, x::RegularInTime)
     space_kernel = k.k.l
     time_kernel = k.k.r
     time_vars = kernelmatrix_diag(time_kernel, get_times(x))
-    return map(
-        (s_t, x_r) -> Diagonal(kernelmatrix_diag(space_kernel, x_r) * s_t),
-        time_vars,
-        x.vs,
+    return Diagonal.(
+        kernelmatrix_diag.(
+            Ref(space_kernel),
+            x.vs
+        ) .* time_vars
     )
 end
 
@@ -185,7 +186,7 @@ function lgssm_components(k_dtc::DTCSeparable, x::RegularInTime, storage::Storag
     C = \(K_space_z_chol, C__)
     Cs = partition(ChainRulesCore.ignore_derivatives(map(length, x.vs)), C)
 
-    cs = _map((h, v) -> fill(h, length(v)), hs_t, x.vs) # This should currently be zero.
+    cs = fill.(hs_t, length.(x.vs)) # This should currently be zero.
     Hs = _map(
         ((I, H_t), ) -> kron(I, H_t),
         zip(Fill(ident_M, N), Hs_t),

src/space_time/regular_in_time.jl

@@ -24,7 +24,16 @@ function Base.collect(x::RegularInTime)
     return [(x, t) for (x, t) in zip(space_inputs, time_inputs)]
 end
 
-Base.getindex(x::RegularInTime, n::Int) = collect(x)[n]
+function Base.getindex(x::RegularInTime, n::Int)
+    n < 0 || throw(BoundsError(x, n))
+    sum_of_lengths = 0
+    for (i, v) in enumerate(x.vs)
+        temp = sum_of_lengths + length(v)
+        temp ≥ n && return (v[n - sum_of_lengths], x.ts[i])
+        sum_of_lengths = temp
+    end
+    throw(BoundsError(x, n))
+end
 
 Base.show(io::IO, x::RegularInTime) = Base.show(io::IO, collect(x))
 

src/util/chainrules.jl

@@ -28,13 +28,6 @@ Zygote.accum(a::Tuple, b::Tuple, c::Tuple) = map(Zygote.accum, a, b, c)
 #                                 StaticArrays                                 #
 # ---------------------------------------------------------------------------- #
 
-function ProjectTo(x::SArray{S,T}) where {S, T}
-    return ProjectTo{SArray}(; element=_eltype_projectto(T), axes=axes(x), static_size=S)
-end
-
-(proj::ProjectTo{SArray})(dx::SArray) = SArray{proj.static_size}(dx.data)
-(proj::ProjectTo{SArray})(dx::AbstractArray) = SArray{proj.static_size}(Tuple(dx))
-
 function rrule(::Type{T}, x::Tuple) where {T<:SArray}
     SArray_rrule(Δ) = begin
         (NoTangent(), Tangent{typeof(x)}(unthunk(Δ).data...))

test/space_time/pseudo_point.jl

@@ -96,7 +96,7 @@ include("../models/model_test_utils.jl")
         validate_dims(lgssm)
 
         # The two approaches to DTC computation should be equivalent up to roundoff error.
-        dtc_naive = dtc(VFE(f_naive(z_naive)), fx_naive, y)
+        dtc_naive = approx_log_evidence(DTC(f_naive(z_naive)), fx_naive, y)
         dtc_sde = dtc(fx, y, z_r)
         @test dtc_naive ≈ dtc_sde rtol=1e-6
 
@@ -150,7 +150,7 @@ include("../models/model_test_utils.jl")
             fx_naive = f_naive(naive_inputs_missings, 0.1)
 
             # Compute DTC using both approaches.
-            dtc_naive = dtc(VFE(f_naive(z_naive)), fx_naive, naive_y_missings)
+            dtc_naive = approx_log_evidence(DTC(f_naive(z_naive)), fx_naive, naive_y_missings)
             dtc_sde = dtc(fx, y_missing, z_r)
             @test dtc_naive ≈ dtc_sde rtol=1e-7 atol=1e-7