KernelFunctions.jl icon indicating copy to clipboard operation
KernelFunctions.jl copied to clipboard

remove zygoterules

Open st-- opened this issue 3 years ago • 9 comments

Curious to see what breaks...

st-- avatar Apr 13 '22 12:04 st--

Codecov Report

Merging #451 (0f20455) into master (8e805ef) will increase coverage by 0.28%. The diff coverage is n/a.

@@            Coverage Diff             @@
##           master     #451      +/-   ##
==========================================
+ Coverage   93.18%   93.46%   +0.28%     
==========================================
  Files          52       51       -1     
  Lines        1261     1255       -6     
==========================================
- Hits         1175     1173       -2     
+ Misses         86       82       -4     
Impacted Files Coverage Δ
src/KernelFunctions.jl 100.00% <ø> (ø)

Continue to review full report at Codecov.

Legend - Click here to learn more Δ = absolute <relative> (impact), ø = not affected, ? = missing data Powered by Codecov. Last update 8e805ef...0f20455. Read the comment docs.

codecov[bot] avatar Apr 13 '22 12:04 codecov[bot]

Huh, this is weird. Is it clear if we're hitting these adjoints at all in our tests?

willtebbutt avatar Apr 13 '22 12:04 willtebbutt

Huh, this is weird. Is it clear if we're hitting these adjoints at all in our tests?

_map is used by the transforms, and most of the transforms have passing AD tests, so I would say "probably"? I've not yet added print() statements to the custom adjoints to actually see...

st-- avatar Apr 13 '22 14:04 st--

_map is used by the transforms, and most of the transforms have passing AD tests, so I would say "probably"?

The rules were defined for map, not _map.

devmotion avatar Apr 13 '22 14:04 devmotion

_map is used by the transforms, and most of the transforms have passing AD tests, so I would say "probably"?

The rules were defined for map, not _map.

The rules were defined for map(t::Transform, x), which is defined in src/transform/transform.jl as _map(t, x). So what are you trying to say by that?

st-- avatar Apr 13 '22 14:04 st--

That https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/blob/8e805ef25d745ab391e4ad424e690507e02cd042/src/zygoterules.jl#L1-L7 are rules for map. Internally, typically we work with _map though (due to AD issues and map being handled to generally by Zygote), e.g., in https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/blob/05fe34077dffe66ed2808c07e21e261496c4f089/src/kernels/transformedkernel.jl#L84-L118, and hence rules for map will not be hit.

devmotion avatar Apr 13 '22 14:04 devmotion

Just verified this locally by adding print statements to the zygote rules -- they don't seem to be hit. Given that we document that one can call map(t, x), I would be in favour of adding a couple of unit tests to ensure that this actually works. We don't need to do that here, but an issue to this effect so that we don't forget about it would be good.

willtebbutt avatar Apr 13 '22 16:04 willtebbutt

@devmotion ah okay, I think I see what you mean now.

Given the map(t::Transform, x) = _map(t, x) definition, should we then just replace all _map calls inside e.g. kernelmatrix methods with just map?

@willtebbutt what would you like to have added unit tests for ?

st-- avatar Apr 13 '22 19:04 st--

Given the map(t::Transform, x) = _map(t, x) definition, should we then just replace all _map calls inside e.g. kernelmatrix methods with just map?

I don't think this will work. The main reason why _map was introduced was that map did not work because Zygote handles map(f, ::AbstractVector). See https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/issues/113, https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/pull/114, and https://github.com/FluxML/Zygote.jl/issues/646.

devmotion avatar Apr 13 '22 20:04 devmotion