remove zygoterules

[st--]

Curious to see what breaks...

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 93.46%. Comparing base (8e805ef) to head (0f20455).
Report is 58 commits behind head on master.

Additional details and impacted files

@@            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     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[willtebbutt]

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

[st--]

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...

[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.

[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.

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?

[devmotion]

That

KernelFunctions.jl/src/zygoterules.jl

Lines 1 to 7 in 8e805ef

	ZygoteRules.@adjoint function Base.map(t::Transform, X::ColVecs)
	return ZygoteRules.pullback(_map, t, X)
	end
	ZygoteRules.@adjoint function Base.map(t::Transform, X::RowVecs)
	return ZygoteRules.pullback(_map, t, X)
	end

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

KernelFunctions.jl/src/kernels/transformedkernel.jl

Lines 84 to 118 in 05fe340

	function kernelmatrix_diag!(K::AbstractVector, κ::TransformedKernel, x::AbstractVector)
	return kernelmatrix_diag!(K, κ.kernel, _map(κ.transform, x))
	end
	function kernelmatrix_diag!(
	K::AbstractVector, κ::TransformedKernel, x::AbstractVector, y::AbstractVector
	)
	return kernelmatrix_diag!(K, κ.kernel, _map(κ.transform, x), _map(κ.transform, y))
	end
	function kernelmatrix!(K::AbstractMatrix, κ::TransformedKernel, x::AbstractVector)
	return kernelmatrix!(K, κ.kernel, _map(κ.transform, x))
	end
	function kernelmatrix!(
	K::AbstractMatrix, κ::TransformedKernel, x::AbstractVector, y::AbstractVector
	)
	return kernelmatrix!(K, κ.kernel, _map(κ.transform, x), _map(κ.transform, y))
	end
	function kernelmatrix_diag(κ::TransformedKernel, x::AbstractVector)
	return kernelmatrix_diag(κ.kernel, _map(κ.transform, x))
	end
	function kernelmatrix_diag(κ::TransformedKernel, x::AbstractVector, y::AbstractVector)
	return kernelmatrix_diag(κ.kernel, _map(κ.transform, x), _map(κ.transform, y))
	end
	function kernelmatrix(κ::TransformedKernel, x::AbstractVector)
	return kernelmatrix(κ.kernel, _map(κ.transform, x))
	end
	function kernelmatrix(κ::TransformedKernel, x::AbstractVector, y::AbstractVector)
	return kernelmatrix(κ.kernel, _map(κ.transform, x), _map(κ.transform, y))
	end

, and hence rules for map will not be hit.

[willtebbutt]

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.

[st--]

@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 ?

[devmotion]

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 #113, #114, and FluxML/Zygote.jl#646.