Motivation and description

Feature Request: Generalized batched_vec for N-D batches

batched_vec currently only supports the 3D case but not the natural N-D generalization. Please add support when ndims(A) == ndims(B) + 1 and the batch dims match size(A)[3:end] == size(B)[2:end], mirroring batched_mul’s behavior.

Possible Implementation

function batched_vec(A::AbstractArray{T}, B::AbstractArray{T}) where T
    @assert ndims(A) == ndims(B) + 1 && size(A)[3:end] == size(B)[2:end]
    return dropdims(batched_mul(A, reshape(B, size(B,1), 1, size(B)[2:end]...)), dims=2)
end