FluxML / FluxML/NNlib.jl

Inplace version of batched adjoint/transpose

Open
#502 7 comments 0 reactions 0 assignees View on GitHub
good first issue
Dominant language
Julia
Stars
258
Forks
138
Avg merge
2d 1h
Merged PRs (30d)
3

Description

We are missing an the inplace version of batched adjoint/transpose. They are required to avoid gpu scalar indexing with `Base.copy` like `copy(batched_adjoint(CUDA.randn(3,5,2)))`. They can be implemented as:

```julia
# Inplace
function batched_transpose_f!(f, B::AbstractArray{T, 3}, A::AbstractArray{T, 3}) where T
axes(B,1) == axes(A,2) && axes(B,2) == axes(A,1) && axes(A,3) == axes(B,3) || throw(DimensionMismatch(string(f)))
Threads.@threads for i in axes(A,3)
Bi = @view B[:, :, i]
Ai = @view A[:, :, i]
LinearAlgebra.transpose_f!(f, Bi, Ai)
end
return B
end

using GPUArrays
function batched_transpose_f!(f, B::AnyGPUArray{T, 3}, A::AnyGPUArray{T, 3}) where T
axes(B,1) == axes(A,2) && axes(B,2) == axes(A,1) && axes(A,3) == axes(B,3) || throw(DimensionMismatch(string(f)))
GPUArrays.gpu_call(B, A) do ctx, B, A
idx = GPUArrays.@cartesianidx A
@inbounds B[idx[2], idx[1], idx[3]] = f(A[idx[1], idx[2], idx[3]])
return
end
return B
end

batched_adjoint!(B, A) = batched_transpose_f!(adjoint, B, A)
batched_transpose!(B, A) = batched_transpose_f!(transpose, B, A)

# copy
function Base.copy(x::BatchedAdjoint)
p = parent(x)
a1, a2, a3 = axes(p)
return batched_adjoint!(similar(p, (a2, a1, a3)), p)
end
function Base.copy(x::BatchedTranspose)
p = parent(x)
a1, a2, a3 = axes(p)
return batched_transpose!(similar(p, (a2, a1, a3)), p)
end
```

which require an extra dependency of `GPUArrays`. I have no idea where should we put these code under the `ext`.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by locating the batched adjoint/transpose implementations and the package's ext layout. Review how optional GPU backends and GPUArrays dependencies are declared, then determine where the proposed CPU and GPU methods belong. Done means inplace batched transpose and adjoint work for three-dimensional arrays and Base.copy avoids GPU scalar indexing.

Written by the indexing model from the issue text.

Assessment

Tech stack
julia
Domain
backend-api-design, machine-learning
Issue type
Feature
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.