seems this is indeed the layout. by experimenting a bit:
import torch
from torch import nn, optim


def run_findz():
    """
    I 3 O 4
    gru.weight_ih.size() torch.Size([12, 3])
    gru.bias_hh.size() torch.Size([12])
    output tensor([[ 0.5000,  0.5000,  0.5000,  0.5000],
            [ 1.0000,  1.0000,  1.0000,  1.0000],
            [ 1.5000,  1.5000,  1.5000,  1.5000],
            [ 2.0000,  2.0000,  2.0000,  2.0000]])
    output tensor([[ 1.,  1.,  1.,  1.],
            [ 2.,  2.,  2.,  2.],
            [ 3.,  3.,  3.,  3.],
            [ 4.,  4.,  4.,  4.]])
    output tensor([[ 1.0000,  1.0000,  1.0000,  1.0000],
            [ 1.5000,  1.5000,  1.5000,  1.5000],
            [ 2.0000,  2.0000,  2.0000,  2.0000],
            [ 2.5000,  2.5000,  2.5000,  2.5000]])    

    (so z is block 1)
    """
    print('')
    print('find z')

    N = 4
    I = 3
    O = 4
    print('I', I, 'O', O)
    gru = nn.GRUCell(I, O)
    print('gru.weight_ih.size()', gru.weight_ih.size())
    print('gru.bias_hh.size()', gru.bias_hh.size())
    gru.weight_ih.data.fill_(0)
    gru.weight_hh.data.fill_(0)
    gru.bias_ih.data.fill_(0)
    gru.bias_hh.data.fill_(0)

    for i in range(3):
        gru.bias_hh.data.fill_(0)
        gru.bias_hh.data[i * 4:(i + 1) * 4].fill_(20)

        input = torch.zeros(N, I)
        state = torch.zeros(N, O)
        for i in range(N):
            state[i].fill_(i + 1)
        output = gru(input, state)
        print('output', output)


def run_findn():
    """
    find N
    I 3 O 4
    gru.weight_ih.size() torch.Size([12, 3])
    gru.bias_hh.size() torch.Size([12])
    i 0 output tensor([[ 0.,  0.,  0.,  0.]])
    i 1 output tensor([[ 0.,  0.,  0.,  0.]])
    i 2 output tensor([[ 1.,  1.,  1.,  1.]])    

    (so n is block 2)
    """
    print('')
    print('find n')
    N = 1
    I = 3
    O = 4
    print('I', I, 'O', O)
    gru = nn.GRUCell(I, O)
    print('gru.weight_ih.size()', gru.weight_ih.size())
    print('gru.bias_hh.size()', gru.bias_hh.size())
    gru.weight_ih.data.fill_(0)
    gru.weight_hh.data.fill_(0)
    gru.bias_ih.data.fill_(0)
    gru.bias_hh.data.fill_(0)

    gru.bias_hh.data[4:8].fill_(-20)

    # gru.bias_hh.data[0:4].fill_(20)
    for i in range(3):
        gru.bias_ih.data.fill_(0)
        gru.bias_ih.data[i * 4:(i + 1) * 4].fill_(20)

        input = torch.zeros(N, I)
        state = torch.zeros(N, O)
        output = gru(input, state)
        print('i', i, 'output', output)


def run_findr():
    """
    output:

    find r
    I 3 O 4
    gru.weight_ih.size() torch.Size([12, 3])
    gru.bias_hh.size() torch.Size([12])
    output tensor([[ 0.,  0.,  0.,  0.]])
    output tensor([[ 1.,  1.,  1.,  1.]])

    (so r is block 0)
    """
    print('')
    print('find r')
    N = 1
    I = 3
    O = 4
    print('I', I, 'O', O)
    gru = nn.GRUCell(I, O)
    print('gru.weight_ih.size()', gru.weight_ih.size())
    print('gru.bias_hh.size()', gru.bias_hh.size())
    gru.weight_ih.data.fill_(0)
    gru.weight_hh.data.fill_(0)
    gru.bias_ih.data.fill_(0)
    gru.bias_hh.data.fill_(0)

    gru.bias_hh.data[4:8].fill_(-20)
    gru.bias_hh.data[8:12].fill_(20)

    gru.bias_hh.data[0:4].fill_(-100)

    input = torch.zeros(N, I)
    state = torch.zeros(N, O)
    output = gru(input, state)
    print('output', output)

    gru.bias_hh.data[0:4].fill_(20)

    input = torch.zeros(N, I)
    state = torch.zeros(N, O)
    output = gru(input, state)
    print('output', output)


if __name__ == '__main__':
    # run_findz()
    # run_findn()
    run_findr()