code: def sparsemax_grad(dout, w_star, dim):
      out = dout.clone()
      supp = w_star > 0
      out[w_star <= 0] = 0
      nnz = supp.to(dtype=dout.dtype).sum(dim=dim, keepdim=True)
      out = out - (out.sum(dim=dim, keepdim=True) / nnz)
      out[w_star <= 0] = 0
      return out
  
file_path: pytorch-struct/torch_struct/semirings/sparse_max.py
imports: import torch
  from .semirings import _BaseLog
label_desc: Use image node_id 2 for calling a global function with example usage: sparsemax_grad(dout, w_star, dim) and returns: out