return lazy Fill arrays for mean functions from #197

JuliaGaussianProcesses · Oct 28, 2021 · becacaa · becacaa
1 parent a521ae5
commit becacaa
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/mean_function.jl b/src/mean_function.jl
@@ -10,7 +10,7 @@ struct ZeroMean{T<:Real} <: MeanFunction end
 """
 This is an AbstractGPs-internal workaround for AD issues; ideally we would just extend Base.map
 """
-_map_meanfunction(::ZeroMean{T}, x::AbstractVector) where {T} = zeros(T, length(x))
+_map_meanfunction(::ZeroMean{T}, x::AbstractVector) where {T} = Zeros{T}(length(x))
 
 function ChainRulesCore.rrule(::typeof(_map_meanfunction), m::ZeroMean, x::AbstractVector)
     map_ZeroMean_pullback(Δ) = (NoTangent(), NoTangent(), ZeroTangent())
@@ -28,7 +28,7 @@ struct ConstMean{T<:Real} <: MeanFunction
     c::T
 end
 
-_map_meanfunction(m::ConstMean, x::AbstractVector) = fill(m.c, length(x))
+_map_meanfunction(m::ConstMean, x::AbstractVector) = Fill(m.c, length(x))
 
 """
     CustomMean{Tf} <: MeanFunction