I’d think it might be worth trying whether the layout at the end of the GRU documentation 176 works for GRUCell as well (probably that was added to the RNN/GRU/LSTM-Modules but not to the cells).
Best regards
Thomas