permute_ft_output.py

import torch

def permute_ft_output(nnue, permutation):
    l1_size = nnue.layer_stacks.l1.in_features
    assert l1_size == len(permutation)*2

    permutation.extend([x + l1_size // 2 for x in permutation])
    ft_permutation = permutation + list(range(l1_size, nnue.input.num_outputs))

    nnue.input.weight.data = nnue.input.weight.data[:, ft_permutation]
    nnue.input.bias.data = nnue.input.bias.data[ft_permutation]
    nnue.layer_stacks.l1.weight.data = nnue.layer_stacks.l1.weight.data[:, permutation]