# Flux: \`nothing\` gradients with custom layer

**URL:** <https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565>\
**Category:** General Usage\
**Created:** [May 29, 2023, 1:38pm UTC](https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565 "2023-05-29T13:38:37Z")\
**Posts on this page:** 7\
**Page:** 1

<div class="post-metadata">

**Author:** ![FedeClaudi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/fedeclaudi/32/33490_2.png) [@FedeClaudi](https://discourse.julialang.org/u/FedeClaudi)\
**Post date:** [May 29, 2023, 1:38pm UTC](https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565/1 "2023-05-29T13:38:37Z")

</div>

Hi all,

I’m trying to define a custom Layer for Flux:

```julia
struct CANCell
    topology::Symbol
    n::NTuple{D,Int} where {D} # number of neurons in each dimension
    d::Int # number of dimensions
    I::Vector{Tuple} # index (i,j...) of each neuron in the lattice
    Z # gpu(Matrix)
    W # gpu(Matrix)
    ∂ᵢW # gpu.(::Vector{Matrix})
    b₀::Float32
    τ::Float32
    state0
    β # gpu(::Vector{Float32})
end

```

where `β` is the only trainable parameter.

It’s a recurrent model so I define a forward pass:

```julia
function (can::CANCell)(h, v::AbstractVecOrMat)
    W, b₀, τ, ∂ᵢW = can.W, can.b₀, can.τ, can.∂ᵢW
    σ = Flux.NNlib.fast_act(relu, v)

    if size(v, 2) != size(h, 2)
        # reshape h
        rs = gpu(zeros(Float32, size(h, 1), size(v, 2)))
        h = h .+ rs
    end
    dₜh = muladd(W, h, b₀)

    
    idx = argmax.(eachcol(h))
    for d in 1:can.d
        ∂W = can.β .* (∂ᵢW[d][:, idx] .* v[d, :]')
        dₜh = dₜh + ∂W
    end

    h = h + (σ.(dₜh) .- h)/τ

    idx = argmax.(eachcol(h))
    y = can.Z[:, idx]
       
    return h, reshape_cell_output(y, v)
end

```

and then use `@functor` and `Recur`:

```julia

@functor CANCell
Flux.trainable(c::CANCell) = (; β = c.β) # parameters to train

CAN(args...; kwargs...) = Recur(CANCell(args...))

Recur(m::CANCell) = Recur(m, m.state0)

```

Finally, the whole thing is a layer of a larger model:

```julia
Chain(
        input = Chain(
            Dense(N, nh),
            Dense(nh, nh, tanh),  
        ),
        recur = Chain(
            RNN(nh, d, tanh),
            Dense(d, nh),
            Dense(nh, nh),
            Dense(nh, size(can.cell.W, 1)),
            can,
        ),
    )

```

This doesn’t train correctly though. If I print the gradients during training I get:

```julia
(layers = (input = nothing, recur = (layers = (nothing, nothing, nothing, nothing, 
(cell = (topology = nothing, n = nothing, d = nothing, I = nothing, Z = Float32[0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 2.6447595f-5 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0], W = nothing, ∂ᵢW = nothing, b₀ = nothing, τ = nothing, state0 = nothing, β = nothing), state = nothing)),)),)

```

so `nothing` for all layers/params except for `Z` which should not even be trainable. Also the `input` layer is a single `nothing` instead of a tuple with the gradient of each layer. I’m very confused.

I know that the problem is with the layer definition and not the rest of my training code because if I use Flux default layers it all trains correctly. I also used other custom layers where all numeric parameters were trainable and that was also fine.

It seems like defining `trainable` breaks gradient propagation but I can’t figure out why. Using `@functor CANCell (β,)` has the same effect. Any idea what could be the problem?

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [May 29, 2023, 2:19pm UTC](https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565/2 "2023-05-29T14:19:38Z")

</div>

The definition of `trainable` is redundant if `@functor` is already used, but without a MWE with the loss and gradient taking it’s difficult to say what if anything could be at issue.

---

<div class="post-metadata">

**Author:** ![FedeClaudi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/fedeclaudi/32/33490_2.png) [@FedeClaudi](https://discourse.julialang.org/u/FedeClaudi)\
**Post date:** [May 29, 2023, 9:00pm UTC](https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565/3 "2023-05-29T21:00:33Z")

</div>

Thanks for getting back to me.

Removing `Flux.trainable` and using only `@functor CANCell (β,)` doesn’t change.  
The loss function is just `Flux.mse(m(x), y)` and the training loop is standard too. The same code and the same data works perfectly if I don’t use my custom layer (e.g. if I just replace it with an `RNN`). So the problem should be with the layer definition.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [May 29, 2023, 9:17pm UTC](https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565/4 "2023-05-29T21:17:08Z")

</div>

> [@FedeClaudi](#):
>
> The loss function is just `Flux.mse(m(x), y)` and the training loop is standard too.

Neither of those should work out of the box for RNNs because the input handling is different. That’s why a runnable MWE is important, because you’ve surely done other changes to make them work 🙂

---

<div class="post-metadata">

**Author:** ![FedeClaudi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/fedeclaudi/32/33490_2.png) [@FedeClaudi](https://discourse.julialang.org/u/FedeClaudi)\
**Post date:** [May 29, 2023, 9:52pm UTC](https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565/5 "2023-05-29T21:52:17Z")

</div>

No I get it but the whole thing is part of a bigger code base so it’s hard to get a MWE out of it.

But this model works:

```julia
Chain(
        input = Dense(N, nh),
        recur = Chain(
            Dense(nh, nh, tanh),
            Dense(nh, 64,),
            RNN(64, 1)
        ),
    )

```

with all the gradients etc:

```julia
(layers = (input = (weight = Float32[-0.010323494 0.009741096 -0.004862191 -0.0014745682 -0.0041073635 0.00091178465 -0.006024482 0.0027308932 -0.002870893 -0.0006800644 -0.00555486 0.021792755 -0.0030689763 -0.006078959 -0.002313804 0.0011246824 0.0019230729 -0.00074779923 -0.017342443 -0.008566279 0.0034860787 0.00018054058....

```

But this:

```julia
Chain(
        input = Dense(N, nh),
        recur = Chain(
            Dense(nh, nh, tanh),
            Dense(nh, 2,),
            can
        ),
    )

```

does not give the right gradient:

```julia
(layers = (input = nothing, recur = (layers = (nothing, nothing, (cell = (topology = nothing, n = nothing, d = nothing, I = nothing, Z = Float32[-0.009424061 0.0 0.

```

What I find particularly confusing is that the `input` component’s gradient is `nothing`.

Notice that the only difference is the last layer in `recur`. They are both essentially RNNs just that my custom layer has some non-trainable parameters.

The data is passed as batch where each `x` is a vector of matrices as described in the [docs](https://fluxml.ai/Flux.jl/stable/models/recurrence/#Working-with-sequences).

```julia
    function dobatch(b)
        l, grads = Flux.withgradient(model) do m
            length(b) == 1 ? loss(m, b) : loss(m, b...)
        end
        any(isnan.(l)) && error("NaN loss during training")
        grads = grads isa Tuple ? grads[1] : grads

        println(grads)

        update!(opt_state, model, grads) # model params update
        l
    end

```

The total loss for a batch is the sum of `Flux.mse` over each sample.

---

<div class="post-metadata">

**Author:** ![ToucheSir](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/touchesir/32/14411_2.png) [@ToucheSir](https://discourse.julialang.org/u/ToucheSir)\
**Post date:** [May 29, 2023, 10:41pm UTC](https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565/6 "2023-05-29T22:41:34Z")

</div>

> [@FedeClaudi](#):
>
> No I get it but the whole thing is part of a bigger code base so it’s hard to get a MWE out of it.

Placeholder data is always an option. Usually code for a single batch’s worth of training is enough and you’ve posted maybe 70% of that, so filling in the gaps (e.g. what is `can` and how is it defined?) would be enough.

> [@FedeClaudi](#):
>
> just that my custom layer has some non-trainable parameters.

Note that `@functor` and `trainable` don’t actually affect what gradients you get back, so those can be ruled out immediately.

Some more things to try:

- Reduce the example from a `Chain` with a bunch of `Dense` layers to a single `RNN` vs a single `Recur{CanCell}`.
- Drop the `Recur` part and see if you get gradients when manually differentiating through a single timestep with a `CanCell` as your model. If that works, reintroduce the `Recur`.

---

<div class="post-metadata">

**Author:** ![FedeClaudi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/fedeclaudi/32/33490_2.png) [@FedeClaudi](https://discourse.julialang.org/u/FedeClaudi)\
**Post date:** [May 30, 2023, 1:03pm UTC](https://discourse.julialang.org/t/flux-nothing-gradients-with-custom-layer/99565/7 "2023-05-30T13:03:02Z")

</div>

Hello.

After a bit more investigation, it looks like the problem was with using`argmax` which is not differentiable.  
Changing the forward pass of the layer to use `softmax` instead gives correct gradients.  
Thanks for helping out!