In the end I got the desired behaviour by doing the following snippet
i,j,k = x.shape
x = x.view(i*j,k)
x = batch_layer(x)
x = x.view(i,j,k)
This seems to have worked where the usage case was that I have a minibatch of sets / unordered sequences and I wanted the embedding of each token in these sets / sequences to be normalized in the same way with the same parameters.
This isn’t an NLP task but the language is appropiate so I will use it.
I wanted to normalize each word token consistently with the same parameters. In the context of the normalization the batch of sentences where each sentence contained words was abstracted down to just being a large batch of words. I was unsure how batch_norm1d would actually handle things using the permutation solution.
I apologise if I’m unclear as I’m unsure how batch_norm1d behaves with 3d tensors.