parameters: (op, grad, _)
return_statement: return [array_ops.reshape(
        sparse_ops.sparse_to_dense(ind,
                                   array_ops.reshape(
                                       math_ops.reduce_prod(in_shape), [1]),
                                   array_ops.reshape(grad, [-1]),
                                   validate_indices=False),
        in_shape), array_ops.zeros(
            [], dtype=dtypes.int32)]
docstring: Return the gradients for TopK.
  
    Args:
      op: The TopKOp for which we need to generate gradients.
      grad: Tensor. The gradients passed to the TopKOp.
  
    Returns:
      A list of two tensors, the first being the gradient w.r.t to the input and
      TopK, and the second being the gradient w.r.t. to the indices (all zero).
docstring_summary: Return the gradients for TopK.
docstring_tokens: Return the gradients for TopK .
function: def _TopKGrad(op, grad, _):
    """Return the gradients for TopK.
  
    Args:
      op: The TopKOp for which we need to generate gradients.
      grad: Tensor. The gradients passed to the TopKOp.
  
    Returns:
      A list of two tensors, the first being the gradient w.r.t to the input and
      TopK, and the second being the gradient w.r.t. to the indices (all zero).
    """
    in_shape = array_ops.shape(op.inputs[0])
    ind_shape = array_ops.shape(op.outputs[1])
  
    ind_lastdim = array_ops.gather(ind_shape, array_ops.size(ind_shape) - 1)
    # Flatten indices to 2D.
    ind_2d = array_ops.reshape(op.outputs[1], array_ops.stack([-1, ind_lastdim]))
  
    in_lastdim = array_ops.gather(in_shape, array_ops.size(in_shape) - 1)
    outerdim = array_ops.shape(ind_2d)[0]
    # Compute linear indices (flattened to 1D).
    ind = array_ops.reshape(ind_2d + array_ops.expand_dims(
        math_ops.range(0, outerdim * in_lastdim, in_lastdim), -1), [-1])
  
    # Substitute grad to appropriate locations and fill the rest with zeros,
    # finally reshaping it to the original input shape.
    return [array_ops.reshape(
        sparse_ops.sparse_to_dense(ind,
                                   array_ops.reshape(
                                       math_ops.reduce_prod(in_shape), [1]),
                                   array_ops.reshape(grad, [-1]),
                                   validate_indices=False),
        in_shape), array_ops.zeros(
            [], dtype=dtypes.int32)]
function_tokens: def _TopKGrad ( op , grad , _ ) : in_shape = array_ops . shape ( op . inputs [ 0 ] ) ind_shape = array_ops . shape ( op . outputs [ 1 ] ) ind_lastdim = array_ops . gather ( ind_shape , array_ops . size ( ind_shape ) - 1 ) # Flatten indices to 2D. ind_2d = array_ops . reshape ( op . outputs [ 1 ] , array_ops . stack ( [ - 1 , ind_lastdim ] ) ) in_lastdim = array_ops . gather ( in_shape , array_ops . size ( in_shape ) - 1 ) outerdim = array_ops . shape ( ind_2d ) [ 0 ] # Compute linear indices (flattened to 1D). ind = array_ops . reshape ( ind_2d + array_ops . expand_dims ( math_ops . range ( 0 , outerdim * in_lastdim , in_lastdim ) , - 1 ) , [ - 1 ] ) # Substitute grad to appropriate locations and fill the rest with zeros, # finally reshaping it to the original input shape. return [ array_ops . reshape ( sparse_ops . sparse_to_dense ( ind , array_ops . reshape ( math_ops . reduce_prod ( in_shape ) , [ 1 ] ) , array_ops . reshape ( grad , [ - 1 ] ) , validate_indices = False ) , in_shape ) , array_ops . zeros ( [ ] , dtype = dtypes . int32 ) ]