Is max operation differentiable in Pytorch?
I am using Pytorch to training some neural networks. The part I am confused about is:
prediction = myNetwork(img_batch)
max_act = prediction.max(1)[0].sum()
loss = softcrossentropy_loss - alpha * max_act

In the above codes, &quot;prediction&quot; is the output tensor of &quot;myNetwork&quot;.
I hope to maximize the larget output of &quot;prediction&quot; over a batch.
For example:
[[-1.2, 2.0, 5.0, 0.1, -1.5] [9.6, -1.1, 0.7, 4,3, 3.3]]
For the first prediction vector, the 3rd element is the larget, while for the second vector, the 1st element is the largets. And I want to maximize &quot;5.0+9.6&quot;, although we cannot know what index is the larget output for a new input data.
In fact, my training seems to be successful, because the &quot;max_act&quot; part was really increased, which is the desired behavior to me. However, I heard some discussion about whether max() operation is differentiable or not:
Some says, mathmatically, max() is not differentiable.
Some says, max() is just an identity function to select the largest element, and this largest element is differentiable.

So I got confused now, and I am worried if my idea of maximizing &quot;max_act&quot; is wrong from the beginning.
Could someone provide some guidance if max() operation is differentiable in Pytorch?

max is differentiable with respect to the values, not the indices. It is perfectly valid in your application.
From the gradient point of view, d(max_value)/d(v) is 1 if max_value==v and 0 otherwise. You can consider it as a selector.
d(max_index)/d(v) is not really meaningful as it is a discontinuous function, with only 0 and undefined as possible gradients.
