Given that Zygote does not support mutation, how does Recur gets away with it?

From Flux.jl’s source code here:

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 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 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?

2 Likes