Ordenar una matriz en TensorFlow

Supongamos que tengo una matriz en TensorFlow:

[ 0.12300211, 0.51767069, 0.13886075, 0.55363625], [ 0.47279349, 0.50432992, 0.48080254, 0.51576483], [ 0.84347934, 0.44505221, 0.88839239, 0.48857492], [ 0.93650454, 0.43652734, 0.96464157, 0.47236174], .. 

Me gustaría ordenar esta matriz por la tercera columna. ¿Cómo hago esto? Puedo ordenar cada columna individualmente usando tf.nn.top_k() , lo que me da los valores ordenados y los índices respectivos. Podría usar los índices de esta tercera columna para reordenar los otros, pero no puedo encontrar una Op. De reordenación.

Suponiendo que quiero mantener las cosas en la gráfica (no hay travesuras de Python):

  • ¿Cómo ordeno (la matriz anterior) en TensorFlow?
  • ¿Cómo hago un nuevo pedido en TensorFlow cuando tengo índices para volver a ordenar?

Los siguientes trabajos:

 a = tf.constant(...) # the array reordered = tf.gather(a, tf.nn.top_k(a[:, 2], k=4).indices)