[WIP] Add Zygote repeat rrules

[rayegun]

These are as direct of translations as you can get from Zygote. The last rule, however, is not working because of kwargs shenanigans. Not sure how this should work

Edit: Works towards #383

[rayegun]

Will add missing rrules and frules in another PR

[mzgubic]

Great, thanks! What do you mean by kwarg shenanigans? Could you add the tests as well please?

[mzgubic]

It looks like you are passing optional arguments rather than kwargs actually, kwargs need ;

[rayegun]

Error During Test at /home/will/.julia/packages/ChainRulesTestUtils/bDd51/src/check_result.jl:19 Test threw exception Expression: isapprox(actual, expected; kwargs...) MethodError: no method matching isapprox(::Array{Float64, 3}, ::Array{Float64, 3}; rtol=1.0e-9, atol=1.0e-9, inner=(2, 2, 1), outer=(1, 1, 3))

kwargs are getting passed all the way down to isapprox which doesn't seem right.

[mzgubic]

Oh, I see. Check the test_rrule docstring, you need to do something like fkwargs=(inner=(2, 2, 1), outer=(1, 1, 3)) to pass the kwargs to the function.

[mzgubic]

Thanks, looks good overall! I've added a few style changes and a couple of minor suggestions.

[mzgubic]

Should we add a few more tests? E.g. both matrix and vectors for both 1 and two indices. And similarly a few edge cases for the inner outer one?

Perhaps we could also test for Float64[] or fill(4.0), both zero dimensional arrays

[rayegun]

What do you want the result to be for the empty array? Repeat just returns a zero-dim array for that.

[mzgubic]

Suggested change

	test_rrule(repeat, randn(5), 3)
	test_rrule(repeat, randn(5), 3, 3)
	test_rrule(repeat, randn(3, 3), 2)
	test_rrule(repeat, randn(5, 5), 2,5)
	test_rrule(repeat, randn(5, 4, 3); fkwargs=(inner=(2, 2, 1), outer=(1, 1, 3)))
	test_rrule(repeat, fill(4.0), 3)
	test_frule(repeat, randn(5), 3)
	test_frule(repeat, randn(5), 3,3)
	test_frule(repeat, randn(3, 3), 2)
	test_frule(repeat, randn(3, 3), 2,5)
	test_frule(repeat, randn(5, 4, 3); fkwargs=(inner=(2, 2, 1), outer=(1, 1, 3)))
	test_frule(repeat, fill(4.0), 3)
	test_rrule(repeat, randn(5), 3)
	test_rrule(repeat, randn(5), 3, 3)
	test_rrule(repeat, randn(3, 3), 2)
	test_rrule(repeat, randn(5, 5), 2,5)
	test_rrule(repeat, randn(5, 4, 3); fkwargs=(inner=(2, 2, 1), outer=(1, 1, 3)))
	test_rrule(repeat, fill(4.0), 3)
	@testset "frule" begin
	test_frule(repeat, fill(4.0), 3)
	test_frule(repeat, randn(5), 3)
	test_frule(repeat, randn(5), 3, 3)
	test_frule(repeat, randn(3, 3), 2)
	test_frule(repeat, randn(3, 3), 2, 5)
	test_frule(repeat, randn(5, 4, 3); fkwargs=(inner=(2, 2, 1), outer=(1, 1, 3)))
	end

and similarly for the rrule to keep things tidy

[mcabbott]

Since things like repeat(rand(Int8,2,2,2),1,1,2,2) are allowed, it might be nice to treat the general case? Untested, but my attempt is:

function rrule(::typeof(repeat), x::AbstractArray, scales::Integer...)
    function repeat_pullback_1(dy)
        size2ndims = ntuple(d -> isodd(d) ? size(x,1+d÷2) : get(scales,d÷2,1), 2*ndims(dy))
        sumdy = sum(reshape(dy, size2ndims); dims = ntuple(d -> 2d, ndims(dy)))
        return (NO_FIELDS, reshape(sumdy, size(x)), map(_->DoesNotExist(), scales)...)
    end
    return repeat(x, scales...), repeat_pullback_1
end

What this won't handle is repeat(1:2, 3,2,1,0) but perhaps nobody will write that!

[nickrobinson251]

should this close over length(x) rather than x (to reduce memory requirements)?

Suggested change

	function repeat_pullback(Ȳ)
	return (NO_FIELDS, dropdims(sum(reshape(Ȳ, length(x), :); dims=2); dims=2), DoesNotExist())
	length_x = length(x)
	function repeat_pullback(Ȳ)
	return (NO_FIELDS, dropdims(sum(reshape(Ȳ, length_x, :); dims=2); dims=2), DoesNotExist())

And likewise with the size calls in the rules below.

This is a genuine question, i'm not sure what best practice here is (cc @oxinabox).

I believe this would be manually doing one thing we want an OpaqueClosure to be able to do automatically.
Also maybe this should be mentioned on the "writing good rules docs"?

[oxinabox]

Yeah, I am cautious about doing that since it makes code uglier.
So I wouldn't block a PR over it (it can always be optimized later)
But it does save memory.
We should discuss this with pros and cons in "Writing God Rules"

[mcabbott]

Closed by #460