flax_gated_linear_rnn
