permutation optimization in tensorflow

Problem setup

I have a batch of 4x4 matrices, containing real valued entries.

L = tf.placeholder('float32', shape=[None, 4, 4], name='pairwise-loss')

I'd like to find a batch of 4-permutations such that each 4-permutation minimizes the sum of applying the permutation as one-hot mask.

def get_best_permutation(L_):
    '''
    Returns a batch of permutations `P` such that `tf.reduce_sum(tf.one_hot(P) * L_, axis=(1,2))`
      is minimized for each batch element.

    Args:
        L_: real valued tensor of shape [None, 4, 4]

    Returns:
        P: tf.int64 tensor of shape [None, 4]

    '''
    raise NotImpelmentedError()

The permutation size is small and constant, typically 3-4. However batch size is expected to be very large. Ideally, everything should be done inside a graph, and on GPU so less data transfers.

EDIT it's safe to assume all matrix entries are positive.

Background

This is to implement something similar to Permutation Invariant Training.

Quick and dirty solution

One can just pre-compute all possible permutations since they are small, then apply all of them in parallel. Finally apply tf.argmin to find the best one. However I'd like a more efficient solution.