Fixing implicit gradients

src/kernels/transformedkernel.jl

@@ -63,11 +63,11 @@ end
 # Kernel matrix operations
 
 function kerneldiagmatrix!(K::AbstractVector, κ::TransformedKernel, x::AbstractVector)
-    return kerneldiagmatrix!(K, κ.kernel, map(κ.transform, x))
+    return kerneldiagmatrix!(K, κ.kernel, _map(κ.transform, x))
 end
 
 function kernelmatrix!(K::AbstractMatrix, κ::TransformedKernel, x::AbstractVector)
-    return kernelmatrix!(K, kernel(κ), map(κ.transform, x))
+    return kernelmatrix!(K, kernel(κ), _map(κ.transform, x))
 end
 
 function kernelmatrix!(
@@ -76,17 +76,17 @@ function kernelmatrix!(
     x::AbstractVector,
     y::AbstractVector,
 )
-    return kernelmatrix!(K, kernel(κ), map(κ.transform, x), map(κ.transform, y))
+    return kernelmatrix!(K, kernel(κ), _map(κ.transform, x), _map(κ.transform, y))
 end
 
 function kerneldiagmatrix(κ::TransformedKernel, x::AbstractVector)
-    return kerneldiagmatrix(κ.kernel, map(κ.transform, x))
+    return kerneldiagmatrix(κ.kernel, _map(κ.transform, x))
 end
 
 function kernelmatrix(κ::TransformedKernel, x::AbstractVector)
-    return kernelmatrix(kernel(κ), map(κ.transform, x))
+    return kernelmatrix(kernel(κ), _map(κ.transform, x))
 end
 
 function kernelmatrix(κ::TransformedKernel, x::AbstractVector, y::AbstractVector)
-    return kernelmatrix(kernel(κ), map(κ.transform, x), map(κ.transform, y))
+    return kernelmatrix(kernel(κ), _map(κ.transform, x), _map(κ.transform, y))
 end

src/transform/scaletransform.jl

@@ -17,7 +17,7 @@ end
 
 set!(t::ScaleTransform,ρ::Real) = t.s .= [ρ]
 
-(t::ScaleTransform)(x) = first(t.s) .* x
+(t::ScaleTransform)(x) = first(t.s) * x
 
 _map(t::ScaleTransform, x::AbstractVector{<:Real}) = first(t.s) .* x
 _map(t::ScaleTransform, x::ColVecs) = ColVecs(first(t.s) .* x.X)

test/kernels/transformedkernel.jl

@@ -52,4 +52,27 @@
         end
     end
     test_ADs(x->transform(SqExponentialKernel(), x[1]), rand(1))# ADs = [:ForwardDiff, :ReverseDiff])
+    # Test implicit gradients
+    @testset "Implicit gradients" begin
+        k = transform(SqExponentialKernel(), 2.0)
+        ps = Flux.params(k)
+        X = rand(10, 1); x = vec(X)
+        A = rand(10, 10)
+        # Implicit
+        g1 = Flux.gradient(ps) do
+            tr(kernelmatrix(k, X, obsdim = 1) * A)
+        end
+        # Explicit
+        g2 = Flux.gradient(k) do k
+          tr(kernelmatrix(k, X, obsdim = 1) * A)
+        end
+
+        # Implicit for a vector
+        g3 = Flux.gradient(ps) do
+          tr(kernelmatrix(k, x) * A)
+        end
+        @test g1[first(ps)] ≈ first(g2).transform.s
+        @test g1[first(ps)] ≈ g3[first(ps)]
+    end
+
 end