# How to sort tensor along two axis by variance in TensorFlow?

I have a function that visualizes 2D convolutional kernels. The weights come as 4D-tensors [filter_height, filter_width, in_channels, out_channels], as specified for conv2d. Now I try to sort the kernels by their variance before visualizing them, highest variance first.

My current approach would be something like this:

#get variance along 2D kernel dims, returns tensor [in_channels, out_channels]
mean, variance = tf.nn.moments(kernel, [0,1])

#get sorted indices of output kernels by their variance
indices = tf.nn.top_k(variance, k=(variance.shape[1]), sorted=True).indices

#kernel to [in_channels, out_channels,filter_height, filter_width ]
kernel= tf.transpose(kernel, (2, 3, 0, 1))

#get indices to same dimension as kernel
indices=tf.expand_dims(indices,axis=2)
indices=tf.expand_dims(indices,axis=2)

#sort kernel according to indices
kernel=tf.gather_nd(kernel,indices)

#kernel back to [filter_height, filter_width, in_channels, out_channels]
kernel= tf.transpose(kernel, (2, 3, 0, 1))


But after trying endless approaches, I can't even get tf.gather_nd() to output a tensor with the correct amount of dimensions. I think I'm not entirely sure about the functions inner workings though.

Can anyone lead me into the right direction? Any help is highly appreciated!

Edit: to clarify, i would like a tensor [filter_height, filter_width, in_channels, out_channels] where the first of the filter_heightXfilter_width kernels is the one with the highest variance, the second the one with the next highest and so on up to out_channels