Skip to main content

Sparsemax pytorch

Project description

sparsemax Documentation Status Updates coverage.svg

A PyTorch implementation of SparseMax ( with gradients checked and tested

Sparsemax is an alternative to softmax when one wants to generate hard probability distributions. It has been used to great effect in recent papers like ProtoAttend (


pip install -U sparsemax


Use as if it was nn.Softmax()! Nice and simple.

from sparsemax import Sparsemax
import torch
import torch.nn as nn

sparsemax = Sparsemax(dim=-1)
softmax = torch.nn.Softmax(dim=-1)

logits = torch.randn(2, 3, 5)
logits.requires_grad = True

softmax_probs = softmax(logits)
print("\nSoftmax probabilities")

sparsemax_probs = sparsemax(logits)
print("\nSparsemax probabilities")

Advantages over existing implementations

This repo borrows heavily from:

However, there are a few key advantages:

  1. Backward pass equations implemented natively as a torch.autograd.Function, resulting in 30% speedup, compared to the above repository.

  2. The package is easily pip-installable (no need to copy the code).

  3. The package works for multi-dimensional tensors, operating over any axis.

  4. The operator forward and backward passes are tested (backward-pass check due to torch.autograd.gradcheck

Check that gradients are computed correctly

from torch.autograd import gradcheck
from sparsemax import Sparsemax

input = (torch.randn(6, 3, 20,dtype=torch.double,requires_grad=True))
test = gradcheck(sparsemax, input, eps=1e-6, atol=1e-4)


This package was created with Cookiecutter and the audreyr/cookiecutter-pypackage project template.


0.1.0 (2020-05-25)

  • First release on PyPI.

Project details

Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

sparsemax-0.1.9.tar.gz (12.1 kB view hashes)

Uploaded Source

Built Distribution

sparsemax-0.1.9-py2.py3-none-any.whl (5.5 kB view hashes)

Uploaded Python 2 Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page