The checkpoint for Slot Attention object discovery model: google-research/slot_attention at master · google-research/google-research · GitHub (https://console.cloud.google.com/storage/browser/gresearch/slot-attention) has a tensorflow.keras.layers.GRUCell (from TFv2).
Corresponding keys (with shapes) in the checkpoint are:
{
  'network/layer_with_weights-0/slot_attention/gru/recurrent_kernel/.ATTRIBUTES/VARIABLE_VALUE' : torch.Size([64, 192]), 
  'network/layer_with_weights-0/slot_attention/gru/kernel/.ATTRIBUTES/VARIABLE_VALUE' : torch.Size([64, 192]),
  'network/layer_with_weights-0/slot_attention/gru/bias/.ATTRIBUTES/VARIABLE_VALUE' : torch.Size([2, 192])
}

Parameters for PyTorch one seem to be
{
  'slot_attention.gru.weight_ih': torch.Size([192, 64]),
  'slot_attention.gru.weight_hh': torch.Size([192, 64]),
  'slot_attention.gru.bias_ih': torch.Size([192]), 
  'slot_attention.gru.bias_hh': torch.Size([192])
}

How does one convert them to PyTorch GRUCell (especially in the view of
GRUcell is different in pytorch and tensorflow 4)?
Would the following be correct (in pseudo-code)?:
weight_hh = recurrent_kernel.t()
weight_ih = kernel.t()
bias_ih, bias_hh = bias.unbind()

What should be the bias format? (maybe related 4)
Thank you!