def extract_roles(x, roles):
    """
    x is [N, B, R, *shape]
    roles is [N, B]
    output is [N, B, *shape]
    """
    N, B, R, *shape = x.shape
    assert roles.shape == (N, B)
    return (x.reshape(N*B, R, *shape)[range(N*B), roles.reshape(N*B)]).reshape(N, B, *shape)