Chunked Forwards-Mode

[willtebbutt]

Forwards-mode is taking shape nicely in #389 , but will initially only let you propagate a single tangent forwards at a time. We would ideally like to be able to propagate a collection of tangents forwards at the same time for efficiency reasons -- this is a common optimisation permitted by many systems.

In order to do this, we will require three extensions to the existing system:

1. design + implement a chunked tangent system,
2. extend all frule!!s to handle it, and
3. make some very minor tweaks to build_frule in order to handle this new system.

Points 2 and 3 are self-explanatory, but point 1 requires some discussion.

Chunked-Tangent System

The basic idea here is to pair each primal x with a collection of tangents ts. I'm going to assume the chunk size statically, and bake it into type for the chunks (where appropriate). We might do something like

struct ChunkedPrimitiveTangent{T,N}
    ts::NTuple{N,T}
end

tangent_type(::Type{Float64}, ::Val{N}) where {N} = ChunkedPrimitiveTangent{tangent_type(T),N}

i.e. the chunked tangent type for Float64 with chunk size N is the thing returned above.

For Arrays, we probably want the chunked tangent type to be another Array with an extra dimension to contain the chunks. i.e. the chunked tangent type for a Vector{Float64} would be a Matrix{Float64}, where the second dimension indexes over the chunks. We might want to shove this inside a container to make it semantically meaningful / to ensure that we can enforce statically that all active data has the same chunk size.

For composite types like Tuples, NamedTuples, structs, and mutable structs, we should push the chunks down to the leaves of the objects. i.e. the chunked tangent type for Tuple{Vector{Float64}, Float64} should be a Tuple{Matrix{Float64}, ChunkedPrimitiveTangent{...}} rather than a Vector{Tuple{Vector{Float64},Float64}}.

[yebai]

IIRC, ForwardDiff has some limitations on the argument type it supports. Can you look at ForwardDiff and see the differences between your proposal and ForwardDiff's chunked tangents?

[yebai]

We need to ensure the chunked support remains GPU compatible