search

JuliaML / MLUtils.jl

Utilities and abstractions for Machine Learning tasks
MIT License
107 stars 20 forks source link

Type instability in `DataLoader` #168

Open Vilin97 opened 11 months ago

Vilin97 commented 11 months ago

Is DataLoader type-unstable? Seems so. Does it have to be type-unstable?

using Flux
function test_data_loader(data_loader)
    for (x,y) in data_loader
        continue
    end
end
X = rand(Float32, 2, 10)
Y = rand(Float32, 2, 10)
data_loader = Flux.DataLoader((data=X, label=Y))
@code_warntype test_data_loader(data_loader) # lots of red

If it is type-unstable and it's the intended behavior, it would be great to add a sentence This function is type-unstable to the docstring.

ToucheSir commented 11 months ago

I don't think it would be impossible to make iterate(::DataLoader) type stable, but it would be tricky since there are so many conditionally-applied wrappers here. @lorenzoh and @CarloLucibello thoughts? Do you recall why we moved away from creating the wrappers in the DataLoader constructor?