# Flux unable to differentiate an embedding layer

**URL:** <https://discourse.julialang.org/t/flux-unable-to-differentiate-an-embedding-layer/26067>\
**Category:** Machine Learning\
**Tags:** question, flux\
**Created:** [July 6, 2019, 11:41am UTC](https://discourse.julialang.org/t/flux-unable-to-differentiate-an-embedding-layer/26067 "2019-07-06T11:41:34Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![merckxiaan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/merckxiaan/32/4697_2.png) [@merckxiaan](https://discourse.julialang.org/u/merckxiaan)\
**Post date:** [July 6, 2019, 11:41am UTC](https://discourse.julialang.org/t/flux-unable-to-differentiate-an-embedding-layer/26067/1 "2019-07-06T11:41:34Z")

</div>

I’ve made an embedding layer that should be able to embed a batched input. The forward pass seems to work great, I get back an array with the dimensions I expect.

The backward pass however, doesn’t work. I get an error saying: `DimensionMismatch("tried to assign 10×32×16 array to 320×16 destination")`

```julia
using Flux

struct Embedding
    table
end

Embedding(voc_size, feature_size) = Embedding(param(Flux.glorot_normal(voc_size, feature_size)))

(e::Embedding)(x) = e.table[x, :]

function (e::Embedding)(x::AbstractArray{T, 2}) where {T}
    out = e.table[x, :]
    return(permutedims(out, (1, 3, 2)))
end

@Flux.treelike Embedding

model = Embedding(100, 16) # feature-size=16
input = rand(1:100, (10,32)) # input-size=10, batch-size=32

model(input) # input-size x feature-size x batch-size

loss(x) = sum(model(x)) # not actually meaningful, just as a test
loss(input)

Tracker.gradient(params(model)) do
    loss(input)
end # DimensionMismatch("tried to assign 10×32×16 array to 320×16 destination")

```

Am I using unsupported functionality? I’m thinking the problem lies at the use of permutedims, but how can I solve this?  
I’ve also tried to do this using Zygote (swapping `Tracker.gradient` with `Zygote.gradient` but then I get the same error.)

Thanks for any help!  
Jules

---

<div class="post-metadata">

**Author:** ![simeonschaub](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/simeonschaub/32/216566_2.png) [@simeonschaub](https://discourse.julialang.org/u/simeonschaub)\
**Post date:** [July 6, 2019, 12:14pm UTC](https://discourse.julialang.org/t/flux-unable-to-differentiate-an-embedding-layer/26067/2 "2019-07-06T12:14:48Z")

</div>

Could you provide a self-contained example? I’m not able to run it, since `glorot_normal`, `@treelike` and `test` are not defined.

---

<div class="post-metadata">

**Author:** ![merckxiaan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/merckxiaan/32/4697_2.png) [@merckxiaan](https://discourse.julialang.org/u/merckxiaan)\
**Post date:** [July 6, 2019, 12:22pm UTC](https://discourse.julialang.org/t/flux-unable-to-differentiate-an-embedding-layer/26067/3 "2019-07-06T12:22:32Z")

</div>

Whoops, sorry, now it should run I believe.

---

<div class="post-metadata">

**Author:** ![merckxiaan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/merckxiaan/32/4697_2.png) [@merckxiaan](https://discourse.julialang.org/u/merckxiaan)\
**Post date:** [July 11, 2019, 10:11am UTC](https://discourse.julialang.org/t/flux-unable-to-differentiate-an-embedding-layer/26067/4 "2019-07-11T10:11:01Z")

</div>

@mcabbott created a pr that should fix this: [Non-scalar getindex by mcabbott · Pull Request #256 · FluxML/Zygote.jl · GitHub](https://github.com/FluxML/Zygote.jl/pull/256)  
Thanks!