Backward through topk operation on Variable

Sourabh_Daptardar · October 29, 2017, 9:35am

I believe ‘topk’ is differentiable and ‘indices’ is not (in the autograd graph sense).

Digging in autograd code I found the following line related to class Topk(_MultiSelectionFunction):

pytorch/pytorch/blob/4f33b136d8ba547725022840ee8adf77e1ed245e/torch/autograd/_functions/tensor.py#L537


      
          
          @staticmethod
          def forward(ctx, input, dim, return_indices, args):
              fn = getattr(input, ctx._forward_cls.__name__.lower())
              ctx.return_indices = return_indices
              ctx.input_size = input.size()
              ctx.dim = dim
              output, indices = fn(*args)
              if return_indices:
                  ctx.save_for_backward(indices)
                  ctx.mark_non_differentiable(indices)
                  return output, indices
              else:
                  ctx.indices = indices
                  return output
          
          @staticmethod
          def backward(ctx, grad_output, grad_indices=None):
              grad_input = Variable(grad_output.data.new(ctx.input_size).zero_())
              if ctx.return_indices:
                  indices, = ctx.saved_variables

Perhaps, @apaszke or any one familiar with autograd code, could clarify and confirm. Is there some documentation on this ?