Flux.jl: Zygote: add adjoint for onehotbatch

[Sheemon7]

For online learning or datasets with a lot of classes, where one hot matrix of the whole dataset doesn't fit into memory, it is favourable to compute one hot matrix for each batch again, i.e. using:

loss(x, y) = logitcrossentropy(m(x), onehotbatch(y, 1:C))

The following script fails however, due to non-existence of adjoint for onehotbatch:

using Flux, Zygote
using Statistics: mean
using Base.Iterators: repeated

x_pos = [0, 0]
x_neg = [-1, 1]
X = zeros(2, 10)
Y = rand(1:2, 10)
X = (Y' .- 1) .* x_pos .+ (2 .- Y') .* x_neg

m = Chain(Dense(2, 5), Dense(5, 2))
loss(x, y) = logitcrossentropy(m(x), onehotbatch(y, 1:2))
opt = ADAM()
evalcb = throttle(() -> @show(loss(X, Y)), 1)
Flux.train!(loss, Params(m), repeated((X, Y), 100), opt, cb=evalcb)

ERROR: LoadError: Need an adjoint for constructor Flux.OneHotMatrix{Array{Flux.OneHotVector,1}}. Gradient is of type Array{Float32,2}

The definition of such adjoint shouldn't be hard, but it seems it's not part of the current Flux#zygote. Or am I missing something?

[mcabbott]

Probably it should just be @nograd onehotbatch since y is categorical.

[Sheemon7]

Indeed it does. Thanks for help!

[irinaapopi]

What do you mean by @nograd onehotbatch? I'm having the same issue atm and don't really understand how to fix it