using Zygote
using BenchmarkTools
using KernelFunctions
x = [3.0]
macro old_check_args(K, param, cond, desc=string(cond))
 quote
 if !($(esc(cond)))
 throw(
 ArgumentError(
 string(
 $(string(K)),
 ": ",
 $(string(param)),
 " = ",
 $(esc(param)),
 " does not ",
 "satisfy the constraint ",
 $(string(desc)),
 ".",
 ),
 ),
 )
 end
 end
end
struct OldLinearKernel{Tc<:Real} <: KernelFunctions.SimpleKernel
 c::Vector{Tc}
 function OldLinearKernel(c::Real)
 @old_check_args(LinearKernel, c, c >= zero(c), "c ≥ 0")
 return new{typeof(c)}([c])
 end
end
function f(x)
 k = LinearKernel(;c = x[1])
 sum(k.c)
end
function g(x)
 k = LinearKernel(;c = x[1], check_args=false)
 sum(k.c)
end
function h(x)
 k = OldLinearKernel(x[1])
 sum(k.c)
end
@btime Zygote.gradient($f, $x) # 15.980 μs (150 allocations: 5.89 KiB)
@btime Zygote.gradient($g, $x) # 13.853 μs (142 allocations: 5.72 KiB)
@btime Zygote.gradient($h, $x) # 4.700 μs (51 allocations: 1.89 KiB)

 function OldLinearKernel(c::Real; check_args=true)
 check_args && @old_check_args(LinearKernel, c, c >= zero(c), "c ≥ 0")
 return new{typeof(c)}([c])
 end

function h(x)
 k = OldLinearKernel(x[1])
 sum(k.c)
end
function i(x)
 k = OldLinearKernel(x[1]; check_args=false)
 sum(k.c)
end
@btime Zygote.gradient($h, $x) # 6.086 μs (65 allocations: 2.58 KiB)
@btime Zygote.gradient($i, $x) # 10.404 μs (91 allocations: 3.38 KiB)