I would assume that the majority of mathematically differentiable functions are also differentiable in PyTorch, but to make sure that’s the case, you could check the .grad_fn attribute of the output using inputs, which require gradients e.g. as seen here:
x = torch.randn(10, 10, requires_grad=True)

out = torch.sum(x, dim=1)
print(out.grad_fn)
> <SumBackward1 object at 0x7f79da0b9520> # differentiable

out = torch.argmax(x, dim=1)
print(out.grad_fn)
> None # not differentiable