def METHOD_NAME(z, linear: Linear, a=True):
    d_hid = linear.bias.shape[0] // 2
    index = 0 if a else d_hid
    p = (
        nn.functional.linear(z, linear.weight[index:index + d_hid])
        + linear.bias[index:index + d_hid])
    return p