ProjectTo access to undefined reference

[maartenvd]

struct test
       a
end
 ProjectTo(Vector{test}(undef,4))

fails, as ProjectTo tries to iterate the elements of the vector, which are not defined

[mcabbott]

Can you give more context? I don't think it makes much sense to track gradients with respect to undefined quantities.

[maartenvd]

I have quite a bit of code that I would like to differentiate through, but it seems fairly cumbersome to avoid mutation. As a workaround I tried this quick'n dirty hack:

function _setindex(a::AbstractArray,v,args...)
    b = copy(a);
    b[args...] = v
    b
end
function ChainRulesCore.rrule(::typeof(_setindex),a::AbstractArray,tv,args...) 
    t = _setindex(a,tv,args...);
    pr_a = ProjectTo(a);
    function toret(v)
        lol = copy(pr_a(v));
        lol[args...] = zero.(lol[args...]);
        (NoTangent(),lol,pr_a(v)[args...],fill(ZeroTangent(),length(args))...)
    end
    t,toret
end

the plan is to then write a macro that automatically translate setindex! calls to _setindex

This then runs into the mentioned error, as in my code I currently initialize a Vector{T}(undef,length), and then fill it up with _setindex. Zygote wants to calculate the adjoint of _setindex, which calls ProjectTo. I'm not sure at what level I should tackle this - I can either try to fix _setindex, or maybe this can be solved at the level of chainrules?

[maartenvd]

Alternatively, is there a more canonical workaround? It seems more people require something like this ( https://gist.github.com/sdewaele/9a9c705eb4e8da9d798056d18d04c383 , https://discourse.julialang.org/t/zygote-and-setindex/44554 ).

---

I tried to fix it like :

function ProjectTo(xs::AbstractArray)
    elements = map(eachindex(xs)) do ind
        if isassigned(xs,ind)
            ProjectTo(xs[ind])
        else
            NoTangent()
        end
    end
    if elements isa AbstractArray{<:ProjectTo{<:AbstractZero}}
        return ProjectTo{NoTangent}()  # short-circuit if all elements project to zero
    else
        # Arrays of arrays come here, and will apply projectors individually:
        return ProjectTo{AbstractArray}(; elements=elements, axes=axes(xs))
    end
end

but this function also doesn't work (aside from https://github.com/JuliaLang/julia/issues/45633 )

I give it something with type T subtyping AbstractArray, and I get a projector to AbstractArrays (no attempt is made at restoring T). Then, the tangent type I give it in my function is something constructed by Zygote (a Tangent - namedtuple), which in turn simply passes through the tangent. The Tangent type itself is also problematic - it's a Tangent{Any,...}... - so it somehow lost its primal type.

How do I properly hook into chainrulescore to prevent this? Is there documentation that I am missing? I have a very simple type that pretty much acts like an abstractarray, and I want the tangent types to be of the same type (or at least to be indexeable)

[oxinabox]

I feel handling things containing undef isn't a priority in general, I am not 100% sure what the right behavour is. I would just remove the ProjectTo from your rule, tbh. I don't think it is helping you. But if you do want it, i think you should definite it only for your custom type.

but this function also doesn't work

Except it is monkeypatching the basic behavour what is wrong is that it is defining only the constructor for the ProjectTo but really you want to overload the function (::ProjectTo{MyType}(d::AbstractArray) NamedDims.jl has a decent example: https://github.com/invenia/NamedDims.jl/blob/10a2c7e2fa398527f2e261982c593029470ef5c2/src/chainrules.jl#L12-L21

The Tangent type itself is also problematic - it's a Tangent{Any,...}... - so it somehow lost its primal type.

Zygote always does this. it isn't on you.

How do I properly hook into chainrulescore to prevent this? Is there documentation that I am missing? I

https://juliadiff.org/ChainRulesCore.jl/dev/rule_author/superpowers/projectto.html is the docs, but it only described how to use ProjectTo a bit more is described in https://juliadiff.org/ChainRulesCore.jl/dev/api.html#ChainRulesCore.ProjectTo

We should get documentation https://github.com/JuliaDiff/ChainRulesCore.jl/issues/599

[mcabbott]

I don't know what to make of undef in the primal. But I do wonder whether we ought to simplify ProjectTo in a way which would probably solve this:

Maybe we should just let it ignore the contents of a container. Arrays of non-numbers to start with, possibly tuples, etc too. Whatever operations act on their contents should already impose projection, and perhaps there is no need to do it again?

This would mean removing the methods which store an array of projections. That always seems like quite a lot of overhead (although I have not benchmarked anything). It would also mean that containers like this could be left alone, rather than being mapped over:

julia> ProjectTo(eachcol(rand(2,3)))(eachcol(ones(2,3)))  # 1.9, destroys Slices
3-element Vector{SubArray{Float64, 1, Matrix{Float64}, Tuple{Base.Slice{Base.OneTo{Int64}}, Int64}, true}}:
 [1.0, 1.0]
 [1.0, 1.0]
 [1.0, 1.0]

This would, I think, incidentally remove the undef error above.