Backwards-Assigning Clusters from EdgePooling

[cameronosmith]

I have a few EdgePooling layers, with the clusters produced by each layer's unpooling description in a list clusters.
I want to assign each of the first layer nodes to a cluster in the last layer, essentially mapping the last cluster's nodes back to the original nodes they came from in the first layer. The way I am doing this is as follows:

prev_cluster = clusters[-1]
curr_cluster = None
for cluster_i in reversed(clusters[:-1]):
    curr_cluster = torch.zeros_like(cluster_i)
    curr_cluster[cluster_i] = prev_cluster[cluster_i]
    prev_cluster = curr_cluster

At the end of this loop curr_cluster maps every node in the first layer to a cluster in the last layer.
And then I will assign the first layer's nodes features to the features at the last layer via first_layer_x[curr_cluster] = last_layer_x[curr_cluster].

Is this differentiable so that I can continue in my architecture with the updated first layer nodes?

Thanks for responding to all of these discussions and this is an amazing library.

[cameronosmith]

Sorry for the dumb question, this is what the unpooling method is for.

[rusty1s]

This is differentiable :)