Fix for onecold broadcast bug

[DhairyaLGandhi]

Using mapreduce was returning an Array when used with OneHotMatrix{CuArray}. This causes issues down the chain. map avoids that, and also avoids the reduction giving slightly better speed (~40x).

[DhairyaLGandhi]

bors try

[bors]

try

Build succeeded

• ci/gitlab/trying

[KristofferC]

What was the bug? Could the existing test be tweaked or added to in order to catch it?

[DhairyaLGandhi]

mapreduce is returning an Array with a OneHotMatrix{CuArray}.

I think we already have the test that is supposed to check for this

Flux.jl/test/cuda/cuda.jl

Line 43 in e991228

@test Flux.onecold(y) isa CuArray

Maybe we should add an extra one that checks for the broadcasting just like in the accuracy function?

Edit: I should've written the description properly, and I shall change that to reflect what changes are made and why.

[MikeInnes]

The diff looks odd to me -- if this can be done with a simple map, why was it written with mapreduce before? Is there a behaviour change or was the original code just redundant?

If there's no behaviour change then this is certainly good to go.

[DhairyaLGandhi]

We needed to hit an optimised kernel with the GPU at the time (iirc, because it's been a while since this has been up), but things have moved on massively since then and the regular way works fine now.

No change in behaviour

[darsnack]

I think this will address #864 issues as well. This addresses the comment I made on #864. String labels will still fail to work on the GPU. Any other delays on this other than possibly wanting to be consistent with #1006?

[DhairyaLGandhi]

Bump

[MikeInnes]

I'm not totally clear on the testing situation; if there's a test for this already, why isn't it failing currently? Can we add the test that would fail without this patch?

[lassepe]

A test that would fail right now is the following identity:

using Flux: onehotbatch, onecold
using Test

data = [:b, :a, :c]
labels = [:a, :b, :c]
hot = onehotbatch(data, labels)
cold = onecold(hot, labels)

@test cold == data

This is something that works with the new map version of onecold.

However, I don't know why the the reduce with | was necessary (I guess there is/was a good reason to do it like this). From #1214 one can at least infer that at some point onehotbatch used to produce matrices with eltype Bool.

[lassepe]

Perhaps another relevant data point:

When the mapreduce version of onecold was first introduced, it seems like there was no test for this. The only thing tested is the 1-argument-version. In that case the reduce is simply a 0 | x::Int which should just give x. The 2-argument version where (potentially non-Integer) labels are passed is only tested with Base.Matrix not with Flux.OneHotMatrix.

[darsnack]

I think the original reason was for some CUDA optimization. But this PR really should get some attention.

[CarloLucibello]

@DhairyaLGandhi could you add the test from #764 (comment) and get this merged?

[DhairyaLGandhi]

bors r+

[bors]

Build failed:

• ci/gitlab/gitlab.com

[CarloLucibello]

bors r+

[bors]

Build succeeded:

• ci/gitlab/gitlab.com