JuliaDiff / JuliaDiff/ChainRulesCore.jl

`ProjectTo(::AbstractArray)` does not infer

Open
#407 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

ProjectTo
Dominant language
Julia
Stars
267
Forks
66
PR merge metrics
No merged PRs in 30d

Description

e.g. see

julia> @inferred ProjectTo(rand(3, 3))(Diagonal(rand(3)))
ERROR: return type Diagonal{Float64, Vector{Float64}} does not match inferred return type Union{Base.ReshapedArray{Float64, 2, Diagonal{Float64, Vector{Float64}}, Tuple{Base.MultiplicativeInverses.SignedMultiplicativeInverse{Int64}}}, Diagonal{Float64, Vector{Float64}}}

the reason for the failure is
https://github.com/JuliaDiff/ChainRulesCore.jl/blob/6efb2d258dc225a50d961fb17fa5ba88d7296ce7/src/projection.jl#L174-L183
which decides whether to dy is dx or a reshape(dx, ...) based on the values of the axes.

We could instead do

    dy = if length(axes(dx)) == length(project.axes)
        axes(dx) == project.axes || throw(_projection_mismatch(project.axes, size(dx)))
        dx
    else
        for d in 1:max(M, length(project.axes))
            if size(dx, d) != length(get(project.axes, d, 1))
                throw(_projection_mismatch(project.axes, size(dx)))
            end
        end
        reshape(dx, project.axes)
    end

which does infer, but throws an error for the arrays that need to be reshaped. In practice, this only means that only

        poffv = ProjectTo(OffsetArray(rand(3), 0:2))
        @test axes(poffv([1, 2, 3])) == (0:2,)

test fails in ChainRulesCore
and all the rrule inference tests (but not the one where inplaceablethunk inference fails) in https://github.com/JuliaDiff/ChainRules.jl/pull/459#issuecomment-884757000 are fixed.

do we want this tradeoff or not?

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with src/projection.jl at the referenced lines and reproduce the @inferred ProjectTo example. Compare the proposed axis-length behavior with the OffsetArray test in ChainRulesCore and the rrule inference tests mentioned in ChainRules#459. Done means deciding the tradeoff and updating the relevant tests and implementation accordingly.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
tooling
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.