OpenViSTAI · Agents · Blockchain · Crypto
Submit a project
Menu
AI Models

Sinkhorn Transformer

Sinkhorn Transformer is a Python and PyTorch library that implements bucket-based sparse attention and provides transformer, language-model, and encoder-decoder model examples.

Explore the practical guide ↓

Independent OpenViST project profile, based on linked sources.

Sinkhorn Transformer — OpenViST project profile in AI Models.
OpenViST project illustration. See the linked sources for product details.
Understand the project

Practical overview

When it fits

Use it when developing or studying transformer models that match query buckets with relevant key buckets, or when experimenting with the documented reversible-network and feed-forward-chunking options.

Who it is for

  • Researchers exploring sparse-attention transformer designs.
  • Developers prototyping language models or encoder-decoder models in PyTorch.

Documented uses

  • Construct a transformer using bucket-based sparse attention.
  • Configure a language-model example with sequence, bucket and architecture parameters.
  • Pass encoder representations and masks as context to a decoder.
  • Use an input-padding helper when sequence length does not match the bucket-size multiple.
Before you choose

Requirements & limits

The repository is public and lists an MIT license. The README documents installation with pip.

  • The README says the implementation has diverged from the paper, using attention in place of the original sorting network and Gumbel Sinkhorn sampling.
  • The simple-sort-network settings are documented as working only for non-causal networks.
  • Model inputs ordinarily need sequence lengths that are multiples of the bucket size; the README provides an Autopadder helper for handling this.
  • The repository examples use PyTorch tensors and expose architecture choices such as bucket size, causal operation, chunking and reversibility.

Check current requirements in the official documentation. Source-based guidance does not mean OpenViST has run or independently tested this project.

Source-based practical guide

Set up a Sinkhorn Transformer language-model configuration

Install the package and define the documented language-model configuration and input shape.

Before you start

  • Python environment with PyTorch available
  • Package installation through pip

Steps

  1. Install the package with the command below.

    Documented command
    pip install sinkhorn_transformer

  2. Import torch and SinkhornTransformerLM from sinkhorn_transformer.

  3. Create a SinkhornTransformerLM with model parameters such as num_tokens, dim, heads, depth, max_seq_len and bucket_size.

  4. Create an integer token tensor with torch.randint and pass it to the model, as shown in the README.

Expected outcome

A configured language-model instance and the documented model-call pattern, which returns a tensor shaped as batch, sequence length and vocabulary size.

Guide limitations

  • The README's example uses a sequence length divisible by the bucket size.
  • The repository documentation is an implementation guide, not a guarantee of successful execution in a particular environment.
Documentation 1 ↗
Source-based practical guide

Handle a non-multiple sequence length with Autopadder

Use the documented helper to pass an odd-length token sequence to a language-model instance.

Before you start

  • Python environment with PyTorch available
  • A SinkhornTransformerLM instance and the Autopadder helper from sinkhorn_transformer

Steps

  1. Create a language model with a defined bucket_size.

  2. Wrap the model with Autopadder(model, pad_left=True).

  3. Create the documented token input with a length that is not a multiple of the bucket size.

  4. Pass the input to the wrapped model.

Expected outcome

The README shows the wrapped model returning output for a 1,117-token input with a bucket size of 128.

Guide limitations

  • The helper pads inputs and can handle an input mask when supplied; the README also says contextual key/value inputs and masks are supported.
  • The README notes that an unpadded input whose length is not a multiple of the bucket size otherwise causes the model to complain.
Documentation 1 ↗
Evidence and transparency

Sources & review

Official repository ↗
Source record details
Source availability
Active
Sources checked
16 October 2026
Original website title
lucidrains/sinkhorn-transformer · Official homepage title · retained in its source language
How we review sourcesSuggest a correction