Skip to content

Support testing chunked forward mode #37

Open
@oxinabox

Description

@oxinabox

Consider this (From #36)

simo(x) = (x, 2x)
function ChainRulesCore.frule((_, ẋ), simo, x)
    y = simo(x)
    return y, Composite{typeof(y)}(ẋ, 2ẋ)
end

I believe that that the following should work:

frule_test(simo, (randn(), randn(4)))  # chunked mode, scalar/vector
frule_test(simo, (randn(3), randn(2, 3)))  # chunked mode, vector/matrix

as I believe the following is the correct chunked mode behavour.
@YingboMa am i right?

at least for the scalar primal and vector differential

julia> frule((Zero(), [1, 2, 3]), simo, π)
((π, 6.283185307179586), ([1, 2, 3], [2, 4, 6]))

julia> frule((Zero(), [1 1; 0 1; 1 0]), simo, [1, 1, 1])
(([1, 1, 1], [2, 2, 2]), Composite{Tuple{Array{Int64,1},Array{Int64,1}}}([1 1; 0 1; 1 0], [2 2; 0 2; 2 0]))

Metadata

Metadata

Assignees

No one assigned

    Labels

    enhancementNew feature or requesthelp wantedExtra attention is needed

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions