parameters: (op, grad)
return_statement: return [None, grad_x, None]
docstring: Gradients for SparseSoftmax.
  
    The calculation is the same as SoftmaxGrad:
  
      grad_x = grad_softmax * softmax - sum(grad_softmax * softmax) * softmax
  
    where we now only operate on the non-zero values present in the SparseTensors.
  
    Args:
      op: the SparseSoftmax op.
      grad: the upstream gradient w.r.t. the non-zero SparseSoftmax output values.
  
    Returns:
      Gradients w.r.t. the input (sp_indices, sp_values, sp_shape).
docstring_summary: Gradients for SparseSoftmax.
docstring_tokens: Gradients for SparseSoftmax .
function: def _SparseSoftmaxGrad(op, grad):
    """Gradients for SparseSoftmax.
  
    The calculation is the same as SoftmaxGrad:
  
      grad_x = grad_softmax * softmax - sum(grad_softmax * softmax) * softmax
  
    where we now only operate on the non-zero values present in the SparseTensors.
  
    Args:
      op: the SparseSoftmax op.
      grad: the upstream gradient w.r.t. the non-zero SparseSoftmax output values.
  
    Returns:
      Gradients w.r.t. the input (sp_indices, sp_values, sp_shape).
    """
    indices, shape = op.inputs[0], op.inputs[2]
    out_vals = op.outputs[0]
    sp_output = ops.SparseTensor(indices, out_vals, shape)
    sp_grad = ops.SparseTensor(indices, grad, shape)
    sp_product = ops.SparseTensor(
        indices, sp_output.values * sp_grad.values, shape)
  
    # [..., B, 1], dense.
    sum_reduced = -sparse_ops.sparse_reduce_sum(sp_product, [-1], keep_dims=True)
    # sparse [..., B, C] + dense [..., B, 1] with broadcast; outputs sparse.
    sp_sum = sparse_ops.sparse_dense_cwise_add(sp_grad, sum_reduced)
  
    grad_x = sp_sum.values * sp_output.values
    return [None, grad_x, None]
function_tokens: def _SparseSoftmaxGrad ( op , grad ) : indices , shape = op . inputs [ 0 ] , op . inputs [ 2 ] out_vals = op . outputs [ 0 ] sp_output = ops . SparseTensor ( indices , out_vals , shape ) sp_grad = ops . SparseTensor ( indices , grad , shape ) sp_product = ops . SparseTensor ( indices , sp_output . values * sp_grad . values , shape ) # [..., B, 1], dense. sum_reduced = - sparse_ops . sparse_reduce_sum ( sp_product , [ - 1 ] , keep_dims = True ) # sparse [..., B, C] + dense [..., B, 1] with broadcast; outputs sparse. sp_sum = sparse_ops . sparse_dense_cwise_add ( sp_grad , sum_reduced ) grad_x = sp_sum . values * sp_output . values return [ None , grad_x , None ]