Is the argsort function differentiable in Tensorflow?

I run the following experiment in colab

import tensorflow as tf

x = tf.constant([[4.0, 2.1, 1.0]])
w = tf.Variable([[1.0, 1.0, 1.0]], name='w')

y_true = tf.constant([[1.0, 2.0, 3.0]])

@tf.function
def loss_fn(y_true, y_pred):
    indices = tf.argsort(y_pred)
    x = tf.gather(y_pred, indices, axis=-1)
    return tf.reduce_sum(tf.square(y_true - x))

with tf.GradientTape() as tape:
  y = x * w
  loss = loss_fn(y_true, y)

tape.gradient(loss, [w])

The computed loss in 1.01 and the gradients for w seem to make sense to me. So I would say the answer is yes, if you are using argsort() for indexing purposes. If you have something else in mind maybe you can tweak the example above and figure out if the gradients behave as you expect.