# Differentiable argmin? Trying VQ-VAE in Flux.jl
**URL:** https://discourse.julialang.org/t/differentiable-argmin-trying-vq-vae-in-flux-jl/88424
**Category:** Machine Learning
**Tags:** flux, zygote
**Created:** [October 8, 2022, 2:46am UTC](https://discourse.julialang.org/t/differentiable-argmin-trying-vq-vae-in-flux-jl/88424 "2022-10-08T02:46:53Z")
**Posts on this page:** 5
**Page:** 1
### Author:  [@afishy](https://discourse.julialang.org/u/afishy)
#### Post date: [October 8, 2022, 2:46am UTC](https://discourse.julialang.org/t/differentiable-argmin-trying-vq-vae-in-flux-jl/88424/1 "2022-10-08T02:46:53Z")
I’m trying to follow this [VQ-VAE tutorial](https://github.com/zalandoresearch/pytorch-vq-vae/blob/master/vq-vae.ipynb) in Julia, particularly the VectorQuantizer module:
```julia
# Calculate distances
distances = (torch.sum(flat_input**2, dim=1, keepdim=True)
+ torch.sum(self._embedding.weight**2, dim=1)
- 2 * torch.matmul(flat_input, self._embedding.weight.t()))
# Encoding
encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1)
encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device)
encodings.scatter_(1, encoding_indices, 1)
# Quantize and unflatten
quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape)
```
It looks like Pytorch is able to differentiate through `torch.argmin`, while Zygote in my Julia implementation can’t:
```julia
emb = Flux.Embedding(args[:emb_dim], args[:num_embeddings]; init=Flux.glorot_uniform) |> gpu
ps = Flux.params(emb)
loss, grad = withgradient(ps) do
distances = sum(flat_input.^ 2, dims=2)' .+ sum(emb.weight .^ 2, dims=2) + 2.0f0 * emb.weight * flat_input'
encoding_indices = argmin(distances, dims=1)
encodings = NNlib.scatter(+, zeros(size(z_flat, 1), num_embeddings)' |> gpu, enc_inds; dstsize=(num_embeddings, size(flat_input, 1)))
sum(encodings)
end
```
where the grads w.r.t. ps returns `nothing`
I understand that `argmin` isn’t AD-friendly. Is there a way I could implement this vector quantization step differentiably?
---
### Author:  [@avikpal](https://discourse.julialang.org/u/avikpal)
#### Post date: [October 8, 2022, 2:45pm UTC](https://discourse.julialang.org/t/differentiable-argmin-trying-vq-vae-in-flux-jl/88424/2 "2022-10-08T14:45:02Z")
`argmin` is not the problem here. `scatter` is non-differentiable wrt `encoding_indices` [NNlib.jl/scatter.jl at c0b4b8b6e969422ff4af18b473d02192b27c9cf4 · FluxML/NNlib.jl · GitHub](https://github.com/FluxML/NNlib.jl/blob/c0b4b8b6e969422ff4af18b473d02192b27c9cf4/src/scatter.jl#L209)
```julia
julia> using Zygote, NNlib
julia> gradient((src, idx) -> sum(NNlib.scatter(+, src, idx)), [10, 100], [1, 3])
([1.0, 1.0], nothing)
```
---
### Author:  [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)
#### Post date: [October 8, 2022, 3:09pm UTC](https://discourse.julialang.org/t/differentiable-argmin-trying-vq-vae-in-flux-jl/88424/3 "2022-10-08T15:09:02Z")
Isn’t the issue rather that the PyTorch example does a further matmul with `self._embedding.weight` (which will propagate a gradient signal back) whereas the Julia one does not? I wouldn’t think PyTorch’s `scatter` is differentiable wrt indices either.
---
### Author:  [@avikpal](https://discourse.julialang.org/u/avikpal)
#### Post date: [October 8, 2022, 3:11pm UTC](https://discourse.julialang.org/t/differentiable-argmin-trying-vq-vae-in-flux-jl/88424/4 "2022-10-08T15:11:44Z")
Yes, I think so too. I did not mean to say the scatter implementation is incorrect, just that `argmin` was not causing the problem.
---
### Author:  [@afishy](https://discourse.julialang.org/u/afishy)
#### Post date: [October 10, 2022, 2:51am UTC](https://discourse.julialang.org/t/differentiable-argmin-trying-vq-vae-in-flux-jl/88424/5 "2022-10-10T02:51:08Z")
Thank you, I missed that!