MOInput Type Stability (JuliaGaussianProcesses#465)

* Parametrise type of MOInput

* Bump patch version

* Formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* MOInputByFeatures and fix up doctests

* Add inference tests

* Uncomment tests

* Formatting

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Drop 1.3 CI

* Change Zygote bound that we weren't checking

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: willtebbutt

.github/workflows/ci.yml

@@ -22,7 +22,7 @@ jobs:
       matrix:
         version:
           - '1'
-          - '1.3'
+          - '1.6'
           - 'nightly'
         os:
           - ubuntu-latest

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.40"
+version = "0.10.41"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/mokernels/moinput.jl

@@ -7,7 +7,7 @@
 julia> x = [1, 2, 3];
 
 julia> KernelFunctions.MOInputIsotopicByFeatures(x, 2)
-6-element KernelFunctions.MOInputIsotopicByFeatures{Int64, Vector{Int64}}:
+6-element KernelFunctions.MOInputIsotopicByFeatures{Int64, Vector{Int64}, Int64}:
  (1, 1)
  (1, 2)
  (2, 1)
@@ -24,9 +24,10 @@ The first `out_dim` elements represent all outputs for the first input, the seco
 
 See [Inputs for Multiple Outputs](@ref) in the docs for more info.
 """
-struct MOInputIsotopicByFeatures{S,T<:AbstractVector{S}} <: AbstractVector{Tuple{S,Int}}
+struct MOInputIsotopicByFeatures{S,T<:AbstractVector{S},Tout_dim<:Integer} <:
+       AbstractVector{Tuple{S,Int}}
     x::T
-    out_dim::Integer
+    out_dim::Tout_dim
 end
 
 """
@@ -38,7 +39,7 @@ end
 julia> x = [1, 2, 3];
 
 julia> KernelFunctions.MOInputIsotopicByOutputs(x, 2)
-6-element KernelFunctions.MOInputIsotopicByOutputs{Int64, Vector{Int64}}:
+6-element KernelFunctions.MOInputIsotopicByOutputs{Int64, Vector{Int64}, Int64}:
  (1, 1)
  (2, 1)
  (3, 1)
@@ -53,9 +54,10 @@ As shown above, an `MOInputIsotopicByOutputs` represents a vector of tuples.
 The first `length(x)` elements represent the inputs for the first output, the second
 `length(x)` elements represent the inputs for the second output, etc.
 """
-struct MOInputIsotopicByOutputs{S,T<:AbstractVector{S}} <: AbstractVector{Tuple{S,Int}}
+struct MOInputIsotopicByOutputs{S,T<:AbstractVector{S},Tout_dim<:Integer} <:
+       AbstractVector{Tuple{S,Int}}
     x::T
-    out_dim::Integer
+    out_dim::Tout_dim
 end
 
 const IsotopicMOInputsUnion = Union{MOInputIsotopicByFeatures,MOInputIsotopicByOutputs}
@@ -96,7 +98,7 @@ A data type to accommodate modelling multi-dimensional output data.
 julia> x = [1, 2, 3];
 
 julia> MOInput(x, 2)
-6-element KernelFunctions.MOInputIsotopicByOutputs{Int64, Vector{Int64}}:
+6-element KernelFunctions.MOInputIsotopicByOutputs{Int64, Vector{Int64}, Int64}:
  (1, 1)
  (2, 1)
  (3, 1)
@@ -136,7 +138,7 @@ julia> Y = [1.1 2.1 3.1; 1.2 2.2 3.2]
 julia> inputs, outputs = prepare_isotopic_multi_output_data(x, ColVecs(Y));
 
 julia> inputs
-6-element KernelFunctions.MOInputIsotopicByFeatures{Float64, Vector{Float64}}:
+6-element KernelFunctions.MOInputIsotopicByFeatures{Float64, Vector{Float64}, Int64}:
  (1.0, 1)
  (1.0, 2)
  (2.0, 1)
@@ -184,7 +186,7 @@ julia> Y = [1.1 1.2; 2.1 2.2; 3.1 3.2]
 julia> inputs, outputs = prepare_isotopic_multi_output_data(x, RowVecs(Y));
 
 julia> inputs
-6-element KernelFunctions.MOInputIsotopicByOutputs{Float64, Vector{Float64}}:
+6-element KernelFunctions.MOInputIsotopicByOutputs{Float64, Vector{Float64}, Int64}:
  (1.0, 1)
  (2.0, 1)
  (3.0, 1)

test/mokernels/moinput.jl

@@ -26,6 +26,7 @@
         @test ibo[5] == (x[1], 2)
         @test ibo[7] == (x[3], 2)
         @test all([(x_, i) for i in 1:3 for x_ in x] .== ibo)
+        @inferred getindex(ibo, 1)
     end
 
     @testset "isotopicbyfeatures" begin
@@ -46,6 +47,7 @@
         @test ibf[5] == (x[2], 2)
         @test ibf[7] == (x[3], 1)
         @test all([(x_, i) for x_ in x for i in 1:3] .== ibf)
+        @inferred getindex(ibf, 1)
     end
 
     @testset "prepare_isotopic_multi_output_data" begin

test/utils.jl

@@ -63,7 +63,7 @@
             @test back(ones(size(X)))[1].X == ones(size(X))
         end
 
-        if VERSION >= v"1.6"
+        if VERSION >= v"1.7"
             @testset "Zygote type-inference" begin
                 ctx = NoContext()
                 x = ColVecs(randn(2, 4))