How to specify the type of a Tracked CuArray?

General Usage Performance
1

Hello,

I have a function that calls neural networks on the gpu and I would like to specify the return type for performances.

using Flux
net = Chain(Dense(3,2,relu)) |> gpu
function foo(net::Chain, input::CuArray{Float32})::TrackedArray
    net(input)
end

out = foo(net, cu([1f0,1f0,1f0]))
typeof(out)
> TrackedArray{...,CuArray{Float32,1}}

I would like to specify that ::TrackedArray{...,CuArray{Float32,1}} instead of simply TrackedArray but with the ..., returns syntax: invalid identifier name “…”. If I use ::TrackedArray{CuArray{Float32,1}} instead it compiles but it throws a convert error:

ERROR: Not implemented: convert TrackedArray{…,CuArray{Float32,1}} to TrackedArray{CuArray{Float32,1},N,A} where A<:AbstractArray{CuArray{Float32,1},N} where N

I think it’s important for performance that the compiler know that the TrackedArray is in the GPU, correct me if I’m wrong, I’m just assuming.