simplifying fix for GradientKernelElement

Author: SebastianAment

Manifest.toml

@@ -31,10 +31,16 @@ version = "2.3.0"
 uuid = "0dad84c5-d112-42e6-8d28-ef12dabb789f"
 
 [[deps.ArrayInterface]]
-deps = ["Compat", "IfElse", "LinearAlgebra", "Requires", "SparseArrays", "Static"]
-git-tree-sha1 = "81f0cb60dc994ca17f68d9fb7c942a5ae70d9ee4"
+deps = ["ArrayInterfaceCore", "Compat", "IfElse", "LinearAlgebra", "Static"]
+git-tree-sha1 = "a24db3a330d0ff64789abd52a26c732805619a53"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "5.0.8"
+version = "6.0.5"
+
+[[deps.ArrayInterfaceCore]]
+deps = ["LinearAlgebra", "SparseArrays", "SuiteSparse"]
+git-tree-sha1 = "d3a275e927d411e054c4192e5aca03998c233e94"
+uuid = "30b0a656-2188-435a-8636-2ec0e6a096e2"
+version = "0.1.7"
 
 [[deps.ArrayLayouts]]
 deps = ["FillArrays", "LinearAlgebra", "SparseArrays"]
@@ -84,21 +90,21 @@ version = "0.4.2"
 
 [[deps.CUDA]]
 deps = ["AbstractFFTs", "Adapt", "BFloat16s", "CEnum", "CompilerSupportLibraries_jll", "ExprTools", "GPUArrays", "GPUCompiler", "LLVM", "LazyArtifacts", "Libdl", "LinearAlgebra", "Logging", "Printf", "Random", "Random123", "RandomNumbers", "Reexport", "Requires", "SparseArrays", "SpecialFunctions", "TimerOutputs"]
-git-tree-sha1 = "19fb33957a5f85efb3cc10e70cf4dd4e30174ac9"
+git-tree-sha1 = "925a16b909fdae16920c1319feadecffb6695b9d"
 uuid = "052768ef-5323-5732-b1bb-66c8b64840ba"
-version = "3.10.0"
+version = "3.10.1"
 
 [[deps.ChainRules]]
 deps = ["ChainRulesCore", "Compat", "IrrationalConstants", "LinearAlgebra", "Random", "RealDot", "SparseArrays", "Statistics"]
-git-tree-sha1 = "de68815ccf15c7d3e3e3338f0bd3a8a0528f9b9f"
+git-tree-sha1 = "e9023f88b1655ffc6a4aaef2502878e8116151ef"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.33.0"
+version = "1.35.1"
 
 [[deps.ChainRulesCore]]
 deps = ["Compat", "LinearAlgebra", "SparseArrays"]
-git-tree-sha1 = "9950387274246d08af38f6eef8cb5480862a435f"
+git-tree-sha1 = "9489214b993cd42d17f44c36e359bf6a7c919abf"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.14.0"
+version = "1.15.0"
 
 [[deps.ChangesOfVariables]]
 deps = ["ChainRulesCore", "LinearAlgebra", "Test"]
@@ -114,9 +120,9 @@ version = "0.3.0"
 
 [[deps.Compat]]
 deps = ["Base64", "Dates", "DelimitedFiles", "Distributed", "InteractiveUtils", "LibGit2", "Libdl", "LinearAlgebra", "Markdown", "Mmap", "Pkg", "Printf", "REPL", "Random", "SHA", "Serialization", "SharedArrays", "Sockets", "SparseArrays", "Statistics", "Test", "UUIDs", "Unicode"]
-git-tree-sha1 = "b153278a25dd42c65abbf4e62344f9d22e59191b"
+git-tree-sha1 = "87e84b2293559571802f97dd9c94cfd6be52c5e5"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
-version = "3.43.0"
+version = "3.44.0"
 
 [[deps.CompilerSupportLibraries_jll]]
 deps = ["Artifacts", "Libdl"]
@@ -146,9 +152,9 @@ version = "1.10.0"
 
 [[deps.DataStructures]]
 deps = ["Compat", "InteractiveUtils", "OrderedCollections"]
-git-tree-sha1 = "cc1a8e22627f33c789ab60b36a9132ac050bbf75"
+git-tree-sha1 = "d1fff3a548102f48987a52a2e0d114fa97d730f0"
 uuid = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
-version = "0.18.12"
+version = "0.18.13"
 
 [[deps.DataValueInterfaces]]
 git-tree-sha1 = "bfc1187b79289637fa0ef6d4436ebdfe6905cbd6"
@@ -237,9 +243,9 @@ version = "0.13.2"
 
 [[deps.Flux]]
 deps = ["Adapt", "ArrayInterface", "CUDA", "ChainRulesCore", "Functors", "LinearAlgebra", "MLUtils", "MacroTools", "NNlib", "NNlibCUDA", "Optimisers", "ProgressLogging", "Random", "Reexport", "SparseArrays", "SpecialFunctions", "Statistics", "StatsBase", "Test", "Zygote"]
-git-tree-sha1 = "f84e50845ab88702c721dc7c6129a85cbc1de332"
+git-tree-sha1 = "62350a872545e1369b1d8f11358a21681aa73929"
 uuid = "587475ba-b771-5e3f-ad9e-33799f191a9c"
-version = "0.13.1"
+version = "0.13.3"
 
 [[deps.FoldsThreads]]
 deps = ["Accessors", "FunctionWrappers", "InitialValues", "SplittablesBase", "Transducers"]
@@ -317,9 +323,9 @@ uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 
 [[deps.Intervals]]
 deps = ["Dates", "Printf", "RecipesBase", "Serialization", "TimeZones"]
-git-tree-sha1 = "b993074580045d1551d30990dc0fa5ba6feef92b"
+git-tree-sha1 = "1fd6fccdbdccee5997fb245289d98386c8996180"
 uuid = "d8418881-c3e1-53bb-8760-2df7ec849ed5"
-version = "1.6.0"
+version = "1.7.0"
 
 [[deps.InverseFunctions]]
 deps = ["Test"]
@@ -363,9 +369,9 @@ version = "1.1.1"
 
 [[deps.LLVM]]
 deps = ["CEnum", "LLVMExtra_jll", "Libdl", "Printf", "Unicode"]
-git-tree-sha1 = "c8d47589611803a0f3b4813d9e267cd4e3dbcefb"
+git-tree-sha1 = "10a20c556107dc5833d3bb7c5e45c4a6e191bd28"
 uuid = "929cbde3-209d-540e-8aea-75f648917ca0"
-version = "4.11.1"
+version = "4.13.0"
 
 [[deps.LLVMExtra_jll]]
 deps = ["Artifacts", "JLLWrappers", "LazyArtifacts", "Libdl", "Pkg", "TOML"]
@@ -428,9 +434,9 @@ uuid = "856f044c-d86e-5d09-b602-aeab76dc8ba7"
 version = "2022.0.0+0"
 
 [[deps.MLStyle]]
-git-tree-sha1 = "e49789e5eb7b2d5577aaea395bfcac769df64bb8"
+git-tree-sha1 = "2041c1fd6833b3720d363c3ea8140bffaf86d9c4"
 uuid = "d8e11817-5142-5d16-987a-aa16d5891078"
-version = "0.4.11"
+version = "0.4.12"
 
 [[deps.MLUtils]]
 deps = ["ChainRulesCore", "DelimitedFiles", "FLoops", "FoldsThreads", "Random", "ShowCases", "Statistics", "StatsBase"]
@@ -506,9 +512,9 @@ version = "0.8.5"
 
 [[deps.NNlibCUDA]]
 deps = ["CUDA", "LinearAlgebra", "NNlib", "Random", "Statistics"]
-git-tree-sha1 = "0d18b4c80a92a00d3d96e8f9677511a7422a946e"
+git-tree-sha1 = "e161b835c6aa9e2339c1e72c3d4e39891eac7a4f"
 uuid = "a00861dc-f156-4864-bf3c-e6376f28a68d"
-version = "0.2.2"
+version = "0.2.3"
 
 [[deps.NaNMath]]
 git-tree-sha1 = "737a5957f387b17e74d4ad2f440eb330b39a62c5"
@@ -546,9 +552,9 @@ version = "0.5.5+0"
 
 [[deps.Optimisers]]
 deps = ["ChainRulesCore", "Functors", "LinearAlgebra", "Random", "Statistics"]
-git-tree-sha1 = "2442c3ddbda547c80e8b6451a103719d6a3593dd"
+git-tree-sha1 = "26f58049054343c8103d67a5530284a35f1186cb"
 uuid = "3bd65402-5787-11e9-1adc-39752487f4e2"
-version = "0.2.4"
+version = "0.2.5"
 
 [[deps.OrderedCollections]]
 git-tree-sha1 = "85f8e6578bf1f9ee0d11e7bb1b1456435479d47c"
@@ -682,9 +688,9 @@ version = "0.1.14"
 
 [[deps.Static]]
 deps = ["IfElse"]
-git-tree-sha1 = "3a2a99b067090deb096edecec1dc291c5b4b31cb"
+git-tree-sha1 = "5d2c08cef80c7a3a8ba9ca023031a85c263012c5"
 uuid = "aedffcd0-7271-4cad-89d0-dc628f76c6d3"
-version = "0.6.5"
+version = "0.6.6"
 
 [[deps.StaticArrays]]
 deps = ["LinearAlgebra", "Random", "Statistics"]
@@ -708,6 +714,10 @@ git-tree-sha1 = "8977b17906b0a1cc74ab2e3a05faa16cf08a8291"
 uuid = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 version = "0.33.16"
 
+[[deps.SuiteSparse]]
+deps = ["Libdl", "LinearAlgebra", "Serialization", "SparseArrays"]
+uuid = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
+
 [[deps.SymEngine]]
 deps = ["Compat", "Libdl", "LinearAlgebra", "RecipesBase", "SpecialFunctions", "SymEngine_jll"]
 git-tree-sha1 = "6cf88a0b98c758a36e6e978a41e8a12f6f5cdacc"

Project.toml

@@ -1,7 +1,7 @@
 name = "CovarianceFunctions"
 uuid = "b3329135-7958-41d4-bb02-e084c5a526bf"
 authors = ["sebastianament"]
-version = "0.3.5"
+version = "0.3.6"
 
 [deps]
 BesselK = "432ab697-7a72-484f-bc4a-bc531f5c819b"

src/gradient.jl

@@ -60,18 +60,27 @@ Base.eltype(K::GradientKernelElement{T}) where T = T
 # gradient kernel element only used for sparsely representable elements
 Base.Matrix(K::GradientKernelElement) = K * I(size(K, 1))
 
+function GradientKernelElement{T}(k, x, y, it::InputTrait) where T
+    GradientKernelElement{T, typeof(k), typeof(x), typeof(y), typeof(it)}(k, x, y, it)
+end
+
+function gradient_kernel(k, x, y, it::InputTrait)
+    T = gramian_eltype(k, x, y)
+    GradientKernelElement{T}(k, x, y, it)
+end
+
+function gradient_kernel!(K::GradientKernelElement, k, x, y, it::InputTrait)
+    GradientKernelElement{eltype(K)}(k, x, y, it)
+end
+
 function Base.:*(G::GradientKernelElement, a)
     T = promote_type(eltype(G), eltype(a))
     b  = zeros(T, size(a))
     mul!(b, G, a)
 end
 
 const GenericGradientKernelElement{T, K, X, Y} = GradientKernelElement{T, K, X, Y, <:GenericInput}
-
 const IsotropicGradientKernelElement{T, K, X, Y} = GradientKernelElement{T, K, X, Y, IsotropicInput}
-function IsotropicGradientKernelElement{T}(k, x, y) where T
-    IsotropicGradientKernelElement{T, typeof(k), typeof(x), typeof(y)}(k, x, y, IsotropicInput())
-end
 
 # isotropic kernel
 function LinearAlgebra.mul!(b, G::IsotropicGradientKernelElement, a, α::Number = 1, β::Number = 0) #, ::IsotropicInput = G.input_trait)
@@ -95,19 +104,8 @@ function WoodburyFactorizations.Woodbury(K::IsotropicGradientKernelElement)
     return K = Woodbury(D, r, C, r')
 end
 
-function gradient_kernel!(K::IsotropicGradientKernelElement, k, x, y, ::IsotropicInput)
-    typeof(K)(k, x, y, IsotropicInput())
-end
-
-function gradient_kernel(k, x, y, ::IsotropicInput)
-    T = gramian_eltype(k, x, y)
-    IsotropicGradientKernelElement{T}(k, x, y)
-end
-
 const DotProductGradientKernelElement{T, K, X, Y} = GradientKernelElement{T, K, X, Y, DotProductInput}
-function DotProductGradientKernelElement{T}(k, x, y) where T
-    DotProductGradientKernelElement{T, typeof(k), typeof(x), typeof(y)}(k, x, y, DotProductInput())
-end
+
 function LinearAlgebra.mul!(b, K::DotProductGradientKernelElement, a, α::Number = 1, β::Number = 0)
     k, x, y = K.k, K.x, K.y
     d² = dot(x, y)
@@ -126,19 +124,7 @@ function WoodburyFactorizations.Woodbury(K::DotProductGradientKernelElement)
     return K = Woodbury(D, copy(y), C, copy(x)')
 end
 
-function gradient_kernel!(K::DotProductGradientKernelElement, k, x, y, ::DotProductInput)
-    typeof(K)(k, x, y, DotProductInput())
-end
-
-function gradient_kernel(k, x, y, ::DotProductInput)
-    T = gramian_eltype(k, x, y)
-    DotProductGradientKernelElement{T}(k, x, y)
-end
-
 const LinearFunctionalGradientKernelElement{T, K, X, Y} = GradientKernelElement{T, K, X, Y, StationaryLinearFunctionalInput}
-function LinearFunctionalGradientKernelElement{T}(k, x, y) where T
-    LinearFunctionalGradientKernelElement{T, typeof(k), typeof(x), typeof(y)}(k, x, y, StationaryLinearFunctionalInput())
-end
 
 function LinearAlgebra.mul!(b, K::LinearFunctionalGradientKernelElement, a, α::Number = 1, β::Number = 0)
     k, x, y = K.k, K.x, K.y
@@ -162,16 +148,6 @@ function LazyMatrixProduct(K::LinearFunctionalGradientKernelElement)
     return LazyMatrixProduct(c, c2')
 end
 
-# is this necessary?
-function gradient_kernel!(K::LinearFunctionalGradientKernelElement, k, x, y, ::StationaryLinearFunctionalInput)
-    typeof(K)(k, x, y, StationaryLinearFunctionalInput())
-end
-
-function gradient_kernel(k, x, y, ::StationaryLinearFunctionalInput)
-    T = gramian_eltype(k, x, y)
-    LinearFunctionalGradientKernelElement{T}(k, x, y)
-end
-
 function evaluate_block!(Gij, k::GradientKernel, x, y, IT = input_trait(k))
     gradient_kernel!(Gij, k.k, x, y, IT)
 end
@@ -420,7 +396,7 @@ end
 ################################################################################
 # [f, ∂f] ∼ GP([μ, ∂μ], dK) # value + gradient kernel
 # IDEA: For efficiency, maybe create ValueGradientKernelElement like in hessian.jl
-# currently, this is an order of magnitude slower than GradientKernel
+# might not be necessary anymore, benchmark against GradientKernel
 struct ValueGradientKernel{T, K, IT<:InputTrait} <: AbstractDerivativeKernel{T, K}
     k::K
     input_trait::IT

src/gradient_algebra.jl

@@ -1,4 +1,5 @@
 ############################ Gradient Algebra ##################################
+# IDEA: could have specialization for gradient kernels of Power kernels of composite kernels
 ################################### Sum ########################################
 # allocates space for gradient kernel evaluation but does not evaluate
 # the separation from evaluation is useful for ValueGradientKernel