Aoo_Lmay:

How do i determine if a function is differentiable at a glance?


You can pass an input tensor with requires_grad=True to a particular operation and check, if the output tensor has a valid .grad_fn as seen here:

x = torch.randn(10, 10, requires_grad=True)
out = torch.argmax(x, dim=1)
print(out.grad_fn)
> None

val, idx = torch.max(x, dim=1)
print(val.grad_fn)
> <MaxBackward0 object at ...>
print(idx.grad_fn)
> None

Also, it’s often helpful to “draw” the function and think how the gradient could look like.



 Aoo_Lmay:

Is there an argmax analog that could be employed?


Yes, softmax (sometimes also called softargmax) is the smooth version of argmax, which provides valid gradients.