def METHOD_NAME(layer, idx):
    n, a = layer.n_head, layer.d_key
    index = paddle.reshape(
        paddle.index_select(
            paddle.reshape(
                paddle.arange(
                    0, n * a, dtype='int64'), shape=[n, a]),
            idx,
            axis=0),
        shape=[-1])
    def reorder_head_matrix(linearLayer, index, dim=1):
        W = paddle.index_select(linearLayer.weight, index, axis=dim).detach()
        if linearLayer.bias is not None:
            if dim == 0:
                b = paddle.assign(linearLayer.bias).detach()
            else:
                b = paddle.assign(
                    L.index_select(
                        linearLayer.bias, index, dim=0)).detach()
        linearLayer.weight.stop_gradient = True
        linearLayer.weight.set_value(W)
        linearLayer.weight.stop_gradient = False
        if linearLayer.bias is not None:
            linearLayer.bias.stop_gradient = True
            linearLayer.bias.set_value(b)
            linearLayer.bias.stop_gradient = False
    reorder_head_matrix(
        layer.q.fn if hasattr(layer.q, 'fn') else layer.q, index)
    reorder_head_matrix(
        layer.k.fn if hasattr(layer.k, 'fn') else layer.k, index)
    reorder_head_matrix(
        layer.v.fn if hasattr(layer.v, 'fn') else layer.v, index)
    reorder_head_matrix(
        layer.o.fn if hasattr(layer.o, 'fn') else layer.o, index, dim=0)