k-Weisfeiler-Lehman Isomorphism test python implementation

Question

I am trying python implementation of k-dimensional Weisfeiler Lehman. However, I cannot make it work.

This is the code:

def kWL(G, k, verbose=False):
    def n_set(G):
        V = list(G.nodes())
        return list(itertools.combinations(V, k))
    def set_initial_colors(n):
        # return {i: [hashlib.sha224(str(i).encode('utf-8')).hexdigest()] for i in n}
        return {i: [nx.weisfeiler_lehman_graph_hash(G.subgraph(i))] for i in n}
    def find_neighbors(V_k, node):
        return [n for n in V_k if len(set(n) - set(V_k[V_k.index(node)])) == 1]
    
    def base_WL(G, k, verbose, n_set, initial_colors_func, find_neighbors_func):    
        N_set = n_set(G)
        colors = initial_colors_func(N_set)
        old_colors = copy.deepcopy(colors)
        for i in range(len(N_set)):
            for k_tuple in N_set:
                M_neigh_colors = "".join([colors[i][0] for i in find_neighbors_func(N_set, k_tuple)])
                # neigh_colors = "".join([old_colors[i][0] for i in find_neighbors_func(G, n, node)])
                colors[k_tuple].extend([M_neigh_colors])
                colors[k_tuple].sort()
            # Update with the hash
            colors = {k_tuple_color: [hashlib.sha224("".join(colors[k_tuple_color]).encode('utf-8')).hexdigest()] for k_tuple_color in colors}
            
            if list(Counter([item for sublist in colors.values() for item in sublist]).values()) == list(Counter([item for sublist in old_colors.values() for item in sublist]).values()) and i != 0:
                if verbose:
                    print(f'Converged at iteration {i}!')
                break
            
            old_colors = copy.deepcopy(colors)
        canonical_form = sorted(Counter([item for sublist in colors.values() for item in sublist]).items())
        if verbose:
            print(f'Canonical Form Found: \n {canonical_form} \n')
        return canonical_form
    
    
    return base_WL(G, k, verbose, n_set, set_initial_colors, find_neighbors)

I slightly adjusted the not functioning version that I found on: W-L test

For instance, this graph → Example graph when I input it with k=2 should have 4 final colors, the same colors for non-edges {3,2}, {3,0} , then the color for {1,2} and {1,0}, another for {2,0}, and the last for {3,1}.