title: How does pytorch backprop through argmax?
body: <p>I'm building Kmeans in pytorch using gradient descent on centroid locations, instead of expectation-maximisation. Loss is the sum of square distances of each point to its nearest centroid.  To identify which centroid is nearest to each point, I use argmin, which is not differentiable everywhere.  However, pytorch is still able to backprop and update weights (centroid locations), giving similar performance to sklearn kmeans on the data.</p>
  
  <p>Any ideas how this is working, or how I can figure this out within pytorch? Discussion on pytorch github suggests argmax is not differentiable: <a href="https://github.com/pytorch/pytorch/issues/1339" rel="noreferrer">https://github.com/pytorch/pytorch/issues/1339</a>. </p>
  
  <p>Example code below (on random pts):</p>
  
  <pre><code>import numpy as np
  import torch
  
  num_pts, batch_size, n_dims, num_clusters, lr = 1000, 100, 200, 20, 1e-5
  
  # generate random points
  vector = torch.from_numpy(np.random.rand(num_pts, n_dims)).float()
  
  # randomly pick starting centroids
  idx = np.random.choice(num_pts, size=num_clusters)
  kmean_centroids = vector[idx][:,None,:] # [num_clusters,1,n_dims]
  kmean_centroids = torch.tensor(kmean_centroids, requires_grad=True)
  
  for t in range(4001):
      # get batch
      idx = np.random.choice(num_pts, size=batch_size)
      vector_batch = vector[idx]
  
      distances = vector_batch - kmean_centroids # [num_clusters, #pts, #dims]
      distances = torch.sum(distances**2, dim=2) # [num_clusters, #pts]
  
      # argmin
      membership = torch.min(distances, 0)[1] # [#pts]
  
      # cluster distances
      cluster_loss = 0
      for i in range(num_clusters):
          subset = torch.transpose(distances,0,1)[membership==i]
          if len(subset)!=0: # to prevent NaN
              cluster_loss += torch.sum(subset[:,i])
  
      cluster_loss.backward()
      print(cluster_loss.item())
  
      with torch.no_grad():
          kmean_centroids -= lr * kmean_centroids.grad
          kmean_centroids.grad.zero_()
  </code></pre>