Skip to content

rrules do not support chunked mode #566

Open
@oscardssmith

Description

@oscardssmith

This is a general issue, but for a specific incarnation, https://github.com/JuliaDiff/ChainRules.jl/blob/8073c7c4638bdd46f4e822d2ab72423c051c5e4b/src/rulesets/Base/array.jl#L40

function rrule(::typeof(Base.vect), X::Vararg{T, N}) where {T, N}
    vect_pullback(ȳ) = (NoTangent(), NTuple{N}(ȳ)...)
    return Base.vect(X...), vect_pullback
end

This rule implicitly assumes that is a Vector, but if you are taking a jacobian, it will be a Matrix in which case, it should be

function rrule(::typeof(Base.vect), X::Vararg{T, N}) where {T, N}
    vect_pullback(ȳ) = (NoTangent(), ȳ...)
    return Base.vect(X...), vect_pullback
end

Similar problems also exist for the getindex rrules, and I'm sure there are a bunch of other similar cases.
Is there a good general solution to this?

Metadata

Metadata

Assignees

No one assigned

    Labels

    designRequires some desgin before changes are madedocumentationImprovements or additions to documentationhelp wantedExtra attention is needed

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions