# Given that Zygote does not support mutation, how does Recur gets away with it?
**URL:** https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290
**Category:** Machine Learning
**Tags:** question, flux, zygote
**Created:** [March 21, 2020, 5:41am UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290 "2020-03-21T05:41:38Z")
**Posts on this page:** 8
**Page:** 1
### Author:  [@Azamat](https://discourse.julialang.org/u/Azamat)
#### Post date: [March 21, 2020, 5:41am UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290/1 "2020-03-21T05:41:38Z")
From Flux.jl’s source code [here](https://github.com/FluxML/Flux.jl/blob/1605a010398ea2dae423ca8e346d19e213594de5/src/layers/recurrent.jl#L37):
```nohighlight
mutable struct Recur{T}
cell::T
init
state
end
Recur(m, h = hidden(m)) = Recur(m, h, h)
function (m::Recur)(xs...)
h, y = m.cell(m.state, xs...)
m.state = h
return y
end
```
it looks like `Recur` struct, which is used for every recurrent layer, does mutation of its `state` field in the forward pass. But Zygote.jl does not support mutation, so why is this not throwing something like `ERROR: Mutation is not supported!` as it usually does in such cases?
**For a context of where this came up** : I was implementing my custom stateful layer. Initially, I just defined `MyCustomRecurrentCell` [similarly](https://github.com/FluxML/Flux.jl/blob/1605a010398ea2dae423ca8e346d19e213594de5/src/layers/recurrent.jl#L61) to `RNNCell` and relied on `Recur` to handle the `MyCustomRecurrentCell`’s state mutation, just like `RNNCell` does. But then I discovered that because `Recur`’s fields are not type annotated I was getting `Any` type outputs for my state which was propagating to all the other layer and I was getting `Any` everywhere because of this. Then I filed [this issue](https://github.com/FluxML/Flux.jl/issues/1092) and decided to reimplement `Recur` with type annotated fields. When I did so, I’m now getting something like `Mutation is not supported` error. So I was wondering how is `Recur` able to get away with mutating state?
---
### Author:  [@jeremiedb](https://discourse.julialang.org/u/jeremiedb)
#### Post date: [June 23, 2020, 7:24am UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290/2 "2020-06-23T07:24:50Z")
@Azamat I was wondering if you figured out what was going on with the state mutation?
I’m also still puzzled about figuring how the `Recur` struct handles its state to be mutated.
More specifically, I’m wondering if the reason could come from the usage of a broadcast for applying the forward pass over the sequence, such as in: `rnn.(x)`. Do you know if there’s any guarantee when applying the broadcast that it is actually applied in a sequence, which is needed for the state to be updated in right order?
---
### Author:  [@Azamat](https://discourse.julialang.org/u/Azamat)
#### Post date: [June 23, 2020, 2:44pm UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290/3 "2020-06-23T14:44:58Z")
@jeremiedb No, I haven’t. Maybe @MikeInnes can shed some light on this.
Besides being type-unstable, there are some other issues with the current implementation of recurrent layers in Flux (see e.g. [Flux.jl#1089](https://github.com/FluxML/Flux.jl/issues/1089)), which is why I think they need to be redesigned completely with the attention to performance
---
### Author:  [@MikeInnes](https://discourse.julialang.org/u/MikeInnes)
#### Post date: [June 23, 2020, 3:01pm UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290/4 "2020-06-23T15:01:50Z")
Zygote doesn’t support mutation _of arrays_, but you can mutate other objects just fine (eg try code with dictionaries in).
What might be happening when you add a type restraint is that Julia has to convert the array before storing it, which might call a mutating kernel. Whatever the exact reason, the error must be coming from mutation of an array.
Feel free to try it, but I’d be very surprised is adding type restrictions to `Recur` gave any meaningful performance improvement. I know the Julia manual talks about type inference and globals etc., but ML really has very different performance constraints; there’s plenty of active work here and it’s definitely not the case that these things are written without attention to performance.
---
### Author:  [@jeremiedb](https://discourse.julialang.org/u/jeremiedb)
#### Post date: [June 27, 2020, 4:53am UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290/5 "2020-06-27T04:53:19Z")
Although I remain unclear on how Zygote gets away with the apparent array mutation that happens within the `Recur`, from the test below, it effectively looks like implementing an immutable struct for the `RNNCell` and `Recur` along with type annotation doesn’t bring any improvement (both for time and memory allocations):
```julia
using Flux
using Flux: Recur, @functor, glorot_uniform, hidden
using BenchmarkTools
# immutable alternative - no hidden state since redundant with one defined in Recur
struct MyRNNCell{F,T}
σ::F
Wi::Matrix{T}
Wh::Matrix{T}
b::Vector{T}
end
# initializer
MyRNNCell(in::Integer, out::Integer, σ = tanh; init = glorot_uniform) = MyRNNCell(σ, init(out, in), init(out, out), init(out))
# overload
function (m::MyRNNCell{F,T})(h::Matrix{T}, x::Matrix{T}) where {F,T}
σ, Wi, Wh, b = m.σ, m.Wi, m.Wh, m.b
h = σ.(Wi*x .+ Wh*h .+ b)
return h, h
end
@functor MyRNNCell Wi, Wh, b
# immutable alternative implementation of Recur
struct MyRecur{F,T}
cell::MyRNNCell{F,T}
init::Vector{T}
state::Matrix{T}
end
# overload
function (m::MyRecur)(xs...)
h, y = m.cell(m.state, xs...)
m.state .= h
return y
end
@functor MyRecur cell, init
# define model based on alternative immutable struct
m = MyRecur(MyRNNCell(128,256), zeros(Float32,256), zeros(Float32,256, 512))
# original / benchmark model
bm = RNN(128,256)
# generate data
x = rand(Float32, 128,512)
xx = [x = rand(Float32, 128,512) for i in 1:100]
# all clean inference
@code_warntype m(x)
# inference issues raised
@code_warntype bm(x)
@btime m(x)
2.562 ms (11 allocations: 1.50 MiB)
@btime bm(x)
2.553 ms (11 allocations: 1.50 MiB)
```
Any pointers to understand how both approach have same performance despite the expected gain from the type stability (at least from the @code\_warntype) would be helpful, as it effectively breaks the rationale I had from the manual performance tips.
---
### Author:  [@MikeInnes](https://discourse.julialang.org/u/MikeInnes)
#### Post date: [June 29, 2020, 12:19pm UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290/6 "2020-06-29T12:19:07Z")
> [@jeremiedb](#):
>
> I remain unclear on how Zygote gets away with the apparent array mutation
There is no array mutation. You can see that the `Recur` struct modifies its own mutable `state` field [here](https://github.com/FluxML/Flux.jl/blob/318ef9d90640cc7effd29bfe8c6b11e924920d29/src/layers/recurrent.jl#L37). But it’s a struct that gets modified, not an array.
> [@jeremiedb](#):
>
> it effectively breaks the rationale I had from the manual performance tips.
Performance is always context-dependent. If code is not type inferred, Julia inserts dynamic dispatches to figure out what method should be called at run time; this costs about 100ns - 1μs. That’s a disaster if you have a scalar loop (since scalar operations can take about a nanosecond, so the overhead is significant), but it’s a non-issue if you’re working with big arrays (as in ML) since individual array operations can easily take milliseconds. Adding a microsecond of dispatch to each operation isn’t even noticeable.
Type inference might matter more in future once the compiler does more array optimisations, but for now it’s largely not worth worrying about.
---
### Author:  [@Azamat](https://discourse.julialang.org/u/Azamat)
#### Post date: [June 29, 2020, 7:18pm UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290/7 "2020-06-29T19:18:53Z")
> Type inference might matter more in future once the compiler does more array optimisations, but for now it’s largely not worth worrying about.
In my use-case, I was using `LSTM`s with [OMEinsum.jl](https://github.com/under-Peter/OMEinsum.jl) and because of their current type-unstable implementation the `LSTM`’s output was getting inferred as `Any` or `Array{Any}` and because of that it was not getting dispatching to the right GPU–optimized kernel in OMEinsum, but to a generic slow one, which is when I’ve filed the issue above. All of this was happening in the backward pass, so really hard to debug and fix. So type inference was a deal-breaker in my case.
---
### Author:  [@kristoffer.carlsson](https://discourse.julialang.org/u/kristoffer.carlsson)
#### Post date: [June 29, 2020, 7:59pm UTC](https://discourse.julialang.org/t/given-that-zygote-does-not-support-mutation-how-does-recur-gets-away-with-it/36290/8 "2020-06-29T19:59:38Z")
> [@Azamat](#):
>
> the `LSTM` 's output was getting inferred as `Any` or `Array{Any}` and because of that it was not getting dispatching to the right GPU–optimized kernel in OMEinsum, but to a generic slow one,
Hm, inference shouldn’t influence dispatch, the actual runtime types are used to determine what method is called. Unless you somehow manually call into inference and do logic based on that, that is.