torch_ans

torch_ans is a high-performance PyTorch extension for ANS (Asymmetric Numeral Systems) entropy coding, supporting both C++ and CUDA backends across multiple platforms including Linux, macOS, and Windows. It enables fast, parallel rANS (range ANS) and tANS (table-driven ANS) compression and decompression for deep learning and data compression workflows. This extension is designed as an efficient, extensible replacement to torchac and CompressAI range coding.

Features

• High-speed ANS compression/decompression: Efficient encoding and decoding using rANS, suitable for large-scale data and neural network applications.
• Parallel rANS/tANS on CPU and GPU: Leverages PyTorch's CPU parallelism and CUDA acceleration/asynchronous operators for batch processing and high throughput.
• Warp-level decoding on CUDA: rans32_16 with num_interleaves=32 maps each stream onto the 32 lanes of a warp, stages the CDF and lookup tables in shared memory and picks the symbol with a tunable inverse-CDF lookup — achieving up to 25.8 Gsymbol/s decode (~29 GB/s of payload) on an RTX 2080 Ti.
• Multi-platform support: Compatible with Linux/macOS/Windows OS and x86_64/aarch64 based platforms, provided that PyTorch and the necessary build tools (C++ compiler, CUDA toolkit if applicable) are available on that platform.
• Flexible rANS/tANS variants: Supports multiple state sizes, stream sizes, frequency precisions, interleaved coding schemes and decoding acceleration tricks (alias coding, inverse cdf).
• Low-level and high-level APIs: Exposes both granular tensor-based operations and user-friendly interfaces for integration. CompressAI like API are also included, see Cookbook.
• Off-the-shelf for research: Modular design using C++ templates, allowing for rapid prototyping and extension of new rANS variants without compromise on efficiency.

Installation

torch_ans is distributed as a source package and as a prebuilt wheel. The native extension is either compiled on your machine - at install time (--no-build-isolation, which needs a C++ toolchain) or lazily on first use (see Runtime dynamic build) - or taken from the prebuilt stable-ABI wheel, which talks to the libtorch stable C shim instead of to ATen and therefore serves every torch >= 2.10 without compiling anything (see Prebuilt stable-ABI wheel). The former is what you get when no wheel matches your platform, or when your torch is older than 2.10: nothing is pinned to a particular PyTorch build, the extension is compiled against the PyTorch you have installed, so there is no wheel-to-torch version to match.

It is therefore recommended to install a specific version of PyTorch first, then install torch_ans with --no-build-isolation. This is important for PyTorch version compability!

From PyPI

pip install torch_ans --no-build-isolation

On Linux x86_64, macOS arm64 and Windows AMD64 pip picks the prebuilt stable-ABI wheel, so nothing is compiled - provided your torch is >= 2.10, which is what the wheel's .so targets. With an older torch the wheel is still installed (pip cannot express "this wheel needs torch >= 2.10" without being willing to upgrade your torch), and torch_ans tells you so once and compiles for your torch instead:

RuntimeWarning: torch_ans ships a prebuilt extension (_C_cpu...) that needs torch >= 2.10,
but torch 2.9.0 is installed; it is being ignored and torch_ans is compiling for this torch
instead ... Upgrade torch - pip install -U 'torch>=2.10' - to use the prebuilt one, or set
TORCH_ANS_NO_STABLE_ABI=1 to silence this warning.

If no wheel matches (other platforms, or --no-binary torch_ans), the package still installs and the extension is compiled - during installation with --no-build-isolation, otherwise on first use.

To build with CUDA support:

WITH_CUDA=1 pip install torch_ans --no-build-isolation

To build with ROCm/AMDGPU support when using a ROCm-enabled PyTorch installation (not tested yet):

WITH_HIP=1 pip install torch_ans --no-build-isolation

From source

pip install . --no-build-isolation

To build with CUDA support:

WITH_CUDA=1 pip install . --no-build-isolation

To build with ROCm/AMDGPU support when using a ROCm-enabled PyTorch installation (not tested yet):

WITH_HIP=1 pip install . --no-build-isolation

CUDA support

CUDA coding is opt-in at build time: WITH_CUDA=1 is required, a visible GPU is not enough (the same holds for both install paths above). The CUDA toolkit must match the CUDA version your PyTorch build uses (nvcc --version vs torch.version.cuda), and nvcc must be on PATH.

A CPU-only extension (the default) can still be used on a GPU machine, but only with CPU tensors. If it is asked to code CUDA tensors it does not fail silently:

1. torch_ans emits a RuntimeWarning and compiles a CUDA extension for the local torch on the spot (this needs the matching nvcc, and takes a few minutes; the result is cached in ~/.cache/torch_extensions/).
2. If no CUDA extension can be produced, it raises a RuntimeError listing the exact commands to fix it - install the toolkit, or rebuild from source with WITH_CUDA=1 pip install . --no-build-isolation, or set TORCH_ANS_FORCE_RUNTIME_BUILD=1 to ignore a pre-built extension and always compile locally.

TORCH_ANS_FORCE_RUNTIME_BUILD=1 is also useful on its own: it makes torch_ans ignore an existing compiled extension (for example one built by an earlier pip install) and compile for the local torch instead.

Tested PyTorch versions

The CI matrix builds and runs the full test suite against PyTorch 1.10.1 (Python 3.7), 2.1.1 (Python 3.9) and 2.7.1 (Python 3.11) on Linux, macOS and Windows, plus an ARM64 Linux job and a lazy-compile job that installs without a compiler. Newer releases (e.g. 2.11) work for the same reason as any other: the extension is compiled locally against the torch that is installed.

Usage

torch_ans supports a wide range of rANS variants, including different state, stream, and frequency sizes, interleaved coding, and parallel coding on CPU or GPU.

Below we list the currently supported rANS variants. Each rANS family comes with three symbol lookups (divided search, alias sampling and inverse-CDF decoding) and several degrees of interleaving; {iN} in the function names stands for the interleave suffix of the operator (_i2, _i4, _i8, _i32) and is omitted for a single (non-interleaved) state. The table-driven tANS families (tans32 / tans16 / tans8) have their own table in the tANS section: the spread table is their symbol lookup, so they have no alias or inverse-CDF variant:

Variant	State Bits	Stream Bits	Max Freq Bits	Interleaved States	Symbol Lookup	Device Support	init_func	push_func	pop_func
rans64	64	32	16	1, 2, 4, 8	divided search	CPU: all / CUDA: 1	rans64_init_stream	rans64{iN}_push	rans64{iN}_pop
rans64_alias	64	32	16	1, 2, 4, 8	alias	CPU: all / CUDA: 1	-	rans64_alias{iN}_push	rans64_alias{iN}_pop
rans64_invcdf	64	32	16	1, 2, 4, 8	inverse CDF	CPU: all / CUDA: 1	-	rans64{iN}_push	rans64{iN}_invcdf_pop
rans32	32	8	16	1, 2, 4, 8	divided search	CPU: all / CUDA: 1	rans32_init_stream	rans32{iN}_push	rans32{iN}_pop
rans32_alias	32	8	16	1, 2, 4, 8	alias	CPU: all / CUDA: 1	-	rans32_alias{iN}_push	rans32_alias{iN}_pop
rans32_invcdf	32	8	16	1, 2, 4, 8	inverse CDF	CPU: all / CUDA: 1	-	rans32{iN}_push	rans32{iN}_invcdf_pop
rans32_16	32	16	15	1, 2, 4, 8, 32	divided search	CPU: all / CUDA: all	rans32_16_init_stream	rans32_16{iN}_push	rans32_16{iN}_pop
rans32_16_alias	32	16	15	1, 2, 4, 8	alias	CPU: all / CUDA: all	-	rans32_16_alias{iN}_push	rans32_16_alias{iN}_pop
rans32_16_invcdf	32	16	15	1, 2, 4, 8, 32	inverse CDF	CPU: all / CUDA: all	-	rans32_16{iN}_push	rans32_16{iN}_invcdf_pop

For example rans64{iN}_push is rans64_push (a single state), rans64_i2_push, rans64_i4_push or rans64_i8_push, and the inverse-CDF decoder of the 4-way interleaved rans32_16 coder is rans32_16_i4_invcdf_pop.

Legend:

• State Bits: Number of bits in the ANS state. This affects initial stream length, thereby impacting compression ratio when there are less symbols.
• Stream Bits: Number of bits per stream element. This affects the frequency of overflowed state to be written into/read from bitstream, slightly affecting speed.
• Max Freq Bits: Maximum supported frequency precision. This affects the accuracy of entropy estimation, thereby impacting compression ratio (higher is better). However, higher frequency precision also leads to larger memory occupation by CDF tables. In the rANS implementation, State Bits > Stream Bits + Max Freq Bits; the tANS families do not follow that relation (their caps come from the encoder's fixed-point arithmetic and, for tans8, from its 8-bit state) — see tANS.
• Interleaved States: Number of interleaved states for sequential coding in one step (the num_interleaves argument of the family's rans*_init_stream). On CPU this maps to SIMD-friendly sequential interleaving. On CUDA, interleaved coding is implemented for two kernel families: rans32_16 (num_interleaves 2, 4 and 8 map to that many consecutive lanes of a warp — 4-lane sub-warp groups for rans32_16_i4 — and 32 maps the whole stream onto the 32 lanes of a warp using warp-level primitives, technique referenced from Recoil, with one shared bitstream cursor per warp) and the three tANS families (tans_warp_cuda.cuh, one lane group per stream, the bit buffer shared by the lanes). The interleaved rans64/rans32 operators are CPU-only; asking them to code CUDA tensors raises an explicit error instead of returning garbage. Interleaved streams are bit-compatible between the CPU and CUDA implementations (same layout and word order), so streams encoded on one device decode on the other.
• Symbol Lookup: The decode-side symbol search. divided search is the default (folded into the plain pop function); alias is the alias-sampling coder (alias_sampling=True in the high-level API), which has its own push and pop (*_alias*_push/*_alias*_pop); inverse CDF keeps the plain push and replaces only the decoder (*_invcdf_pop, used with inverse_cdf_precision).
• Device Support: Which devices each variant is available on (all = every interleave factor listed for that row). Interleaved CUDA coding exists for rans32_16 and for the tANS families, so the rans64/rans32 interleaved variants are CPU-only there. num_interleaves=32 exists for rans32_16 and for the tANS families (it is the full-warp group of their CUDA kernels), and has no alias variant.
• init_func/push_func/pop_func: Main API functions for this variant. The init function is always the family's rans*_init_stream (- reuses the one shown on the family's first row); alias and inverse-CDF variants differ from the family only in the push/pop shown.

In addition to standard and interleaved rANS, two advanced coding types are supported: alias coding and inverse CDF coding. Both are available at every interleave factor their family supports (alias sampling has no 32-way variant), and both can be selected in the high-level API with alias_sampling=True and inverse_cdf_precision respectively.

Alias Coding: Alias coding modifies both the push (encode) and pop (decode) steps. It accelerates the pop (decode) process by enabling constant-time symbol lookup, but increases memory usage during the push (encode) step due to the need for additional alias tables.

Inverse CDF decoding: Inverse CDF decoding only changes the pop (decode) step. Instead of the data-dependent divided search + binary refine, the decoder loads the symbol from a table keyed by the low bits of the ANS state. The table precision is tunable via inverse_cdf_precision:

• inverse_cdf_precision = freq_precision builds a dense table (one entry per quantized frequency, 2**freq_precision entries per distribution — exact O(1) lookup);
• any smaller value builds a sparse table with only 2**q entries per distribution plus a short bounded linear walk (2**(freq_precision - q) steps at most) — 2**(freq_precision - q)x less memory, usually as fast or faster;
• inverse_cdf_precision = "auto" picks q from the alphabet size at init_params time (about log2(alphabet) - 1, tuned so the row fits the shared-memory staging budget on CUDA), and skips the table entirely for alphabets below 64 symbols where it buys less than ~3%.

With a sparse table the memory footprint is modest and the decode speedup is real: up to +80% on CUDA (large alphabets) and up to 2.8x on CPU over the default lookup (measured, see scripts/bench_sparse_invcdf.py).

Use alias coding for fast decoding when memory usage during encoding is not a concern. Use inverse CDF coding (usually with "auto") for maximum decoding speed when a small per-distribution table is acceptable.

tANS (table-driven ANS)

torch_ans also ships tANS — the table-driven ANS coder (FSE style), ported from the reference in 3rdparty/ans/tans.cpp/tans.hpp — behind the same three call sites as the rANS families, selected with impl="tans32" / "tans16" / "tans8":

Variant	State Bits	Stream Bits	Max Freq Bits	Table slots	Interleaved States	Symbol Lookup	Device Support	init_func	push_func	pop_func
tans32	32	16	15	2**freq_precision	1, 2, 4, 8, 32	tANS spread table	all	tans32_init_stream	tans32_push / tans32_i2_push / ...	tans32_pop / tans32_i2_pop / ...
tans16	16	8	15	2**freq_precision	1, 2, 4, 8, 32	tANS spread table	all	tans16_init_stream	tans16_push / tans16_i2_push / ...	tans16_pop / tans16_i2_pop / ...
tans8	8	8	7	2**freq_precision	1, 2, 4, 8, 32	tANS spread table	all	tans8_init_stream	tans8_push / tans8_i2_push / ...	tans8_pop / tans8_i2_pop / ...

The interleaved variants are one state per lane on top of the bit buffer the lanes share (see the stream header below), so num_interleaves=2/4/8/32 works exactly like it does for the rANS families - including TorchANSInterface(impl="tans32", num_interleaves=4). Measured on an i7-class machine with scripts/bench_tans_vs_rans.py, interleaving pays off more for tANS than for rANS whenever the table is bigger than the cache: up to ~2.6x encode / ~2.2x decode at freq_precision=15 (8 lanes), against ~1.2-1.3x at freq_precision=7 where the table fits in L2.

On CUDA tans_cuda.cu runs the same algorithm in two shapes, byte-compatible with the CPU kernels either way, so TorchANSInterface(impl="tans32", device="cuda") works out of the box - the table is built on the host and moved to the device by init_params:

• one thread per stream row at num_interleaves=1;
• one lane group per stream above it (tans_warp_cuda.cuh): the num_interleaves lanes of a group code one row, each owning its state while the bit buffer and the payload cursor stay shared - the family's fixed-width state cannot absorb the bits it does not emit. The per-lane table lookups run in parallel and the bits are committed in lane order, so the payload is bit-identical to the CPU's interleaved one (32 is a full warp per row).

The device kernels stage the whole lookup table into shared memory when it fits the per-block budget (TORCH_ANS_TANS_SMEM_BUDGET, 16 KiB by default; 0 disables it), and the launch block size is derived from the batch size (TORCH_ANS_TANS_BLOCK_THREADS, and TORCH_ANS_TANS_WARP_BLOCK_THREADS for the lane groups) - one thread (or one group) codes one row. Measured on an RTX 2080 Ti (scripts/bench_tans_vs_rans.py --device cuda, 1M symbols, 1024 rows, 256 distributions, see benchmark_status.md): against single-lane rANS at the same freq_precision, single-lane tANS decodes 1.0-1.5x faster at alphabet >= 64 (up to 1.54x at alphabet 256 / freq_precision=10) and stays within ±5% on encode; the rANS 32-way warp path keeps the encode edge (its lanes own whole words, so nothing is serialised per symbol), at the price of ~0.6-0.8 bit/symbol on 1024-symbol rows - each of its 32 lanes flushes its own state per row. The tANS warp kernels are measured against it below.

The shared-memory staging only reaches configurations whose whole table fits the budget (a handful of distributions at a small freq_precision), and there it is a wash for short rows (1024 symbols: 0.92-0.97x of the global path, since the copy is paid per block and L1 already serves a table that size) and a real gain for long ones (64 streams × 16K symbols with a 6.6 KB table: +4.6-5.4% encode / +3.1-3.6% decode). The 256-distribution sweeps never trigger it. In that same small-table regime single-lane tANS decodes 1.15-1.33x faster than single-lane rANS (2.3x with long streams), so it is the family to pick when the distributions are few.

Interleaving on the device is worth it when the batch is small - the case where a one-thread-per-row kernel cannot fill the machine. On the first device kernels that is where 4 and 8 lanes paid: over a 64-distribution table in global memory (benchmarks/tans_interleaved_cuda.json), tans32 at freq_precision=7 went from 19.1/28.3 MB/s (encode/decode of symbols) at one lane to 29.5/41.5 at 4 lanes (+54%/+47%) and 30.1/40.7 at 8 lanes (+58%/+44%) with 8 stream rows, and to +29%/+29% (4 lanes) with 512 rows, flattening to ~+6%/+5% at 4096 rows.

The bit layer was then reworked (one packed __shfl_sync broadcast per lane, a rolled serial section, a parallel bit read on decode - 229 registers/thread down to 42-44), which improved every factor and made 32 lanes the fastest tANS factor at 8 and 512 stream rows, at 2.59x encode / 11.6x decode over one lane at 8 rows (benchmarks/tans_interleaved_cuda_optimized.json is the after; the before/after table and the head-to-head with the rANS warp path are in benchmark_status.md). What 32 lanes still cost is coded size: its 136-byte header (34 words: one state per lane plus the shared bit buffer) is ~1.1 bits/symbol at 1024-symbol rows (6.96 vs 5.99), which is why its encode is 0.94x of one lane at 4096 rows. Structurally the family still cannot match rANS at a tiny batch on encode - a tANS row shares one bit buffer where every rANS lane owns and renormalizes its own words; on decode the two are within a few percent once the device is saturated.

What is different from rANS:

• The coder is table driven. tans_build_tables(cdfs, cdfs_sizes, table_log) turns the quantized cdf into a spread table of 2**table_log slots per distribution plus the encoder's symbol transformation table, and both directions are then a single table access per symbol — no division on encode, no search on decode. table_log is freq_precision here: torch_ans' quantized cdf is already a normalized counter (every symbol owns at least one unit and the last entry is 2**freq_precision), so its differences are the slot counts directly and FSE's renormalization step is not needed. All three families share one table — it only depends on (cdf, table_log) — and it travels in the cdfs slot of push/pop, so a TorchANSInterface call site does not change.
• Table size is the price. 2**freq_precision slots per distribution (`3 * 2**freq_precision
  • 2 * cdf_rowint32 words: ~0.4 MB per distribution at the maximumfreq_precision=15, ~1.5 kB at tans8's 7), so pick the family by how many distributions you code: tans8for thousands of them,tans32/tans16for a few. A smallerfreq_precision` also quantizes the probabilities coarser, so there is a compression/size trade-off on top.
• The state is fixed width ([2**L, 2**(L+1)) encoding, [0, 2**L) decoding), so unlike a rANS state it cannot hold the bits it does not emit. Those live in a stream header — one state word per interleave lane, then the residual bit buffer and its bit count — which is part of the byte string; that is what keeps the stream layout, the capacity contract, rans_stream_to_byte_strings/rans_byte_strings_to_stream and the whole high-level API identical to rANS. A lane owns its state, but the bit buffer is shared by all of them (one bit sequence, like rANS' shared stream cursor), so the header is num_interleaves + 2 words.
• Bypass coding writes the out-of-range value as raw bits instead of coding it with a uniform tANS table as the reference does; both cost exactly bypass_precision bits per digit, so out-of-range symbols behave like they do in rANS.
• Limits. num_interleaves is 1/2/4/8/32 on both devices (32 is the full-warp group of the CUDA interleaved kernels, and the CPU kernels implement it too), and there is no alias_sampling / inverse_cdf_precision variant — the tANS table is the symbol lookup. Every symbol (plus the bypass sentinel) has to own at least one of the 2**freq_precision table slots, so alphabet + 1 > 2**freq_precision cannot be coded.
• When to use it. See the measured comparison in benchmark_status.md: against the fastest rANS family at the same freq_precision, tANS decodes ~1.3-2.1x faster at freq_precision=7-10 (one table access instead of a search; tans8 at its freq_precision=7 cap is the extreme, and num_interleaves=8 keeps 1.26-1.46x of it) but encodes ~0.5-0.85x as fast, and the gap turns against it as freq_precision grows and the table leaves the cache — so tANS is the right choice for decode-heavy workloads with a small freq_precision. tans8 is the small table tier to pick when there are many distributions (~1.5 kB each, and it gets faster relatively as the distribution count grows: 1.55x encode / 1.74x decode at 8 lanes with 4096 distributions).

Parallel ANS stream

All encoding (push) and decoding (pop) operations in torch_ans are fully parallelized. The batch size of the rANS stream tensor (stream.shape[0]) determines the number of parallel states processed simultaneously. This enables efficient utilization of multi-core CPUs and GPUs for large-scale data compression.

• The stream tensor is typically int32 (defined in DEFAULT_TORCH_TENSOR_TYPE and DEFAULT_TORCH_TENSOR_DTYPE in rans_utils.hpp) and has shape (B, L), where B is the batch size (number of parallel states) and L is the stream length. Note that due to int32 range, each stream length could not exceed 2GB.
• Initial stream length L is computed as 1 + ceil(num_interleaves * sizeof(RANS_STATE_TYPE) / 4). The first element stores the stream length, and the remaining elements store the ANS states. Therefore using 32 state bits could reduce overhead than 64 state bits, especially when a large number of parallel stream is used on small amount of data.
• A tANS stream puts a header there instead: one state word per interleave lane (slots 1..num_interleaves, a full int32 whatever the family's state width is), then the residual bit buffer and its bit count — 4 * (num_interleaves + 2) bytes, i.e. 12 at one lane. That header is part of the serialized byte string, which is what keeps the stream layout, the byte-string helpers and the capacity contract identical to rANS.
• When calling rans*_push or rans*_pop, symbols are distributed along the batch dimension and processed in parallel across all states.
• On CPU you can set the number of threads for parallel coding with torch.set_num_threads().

Speed/Compression Considerations

Increasing the number of parallel states (B) generally improves throughput, but may slightly reduce compression ratio due to larger initial states being stored in the bitstream. The optimal number of parallel states depends on your hardware and data size:

• For small datasets (<100MB), CPUs with fewer parallel states (e.g., 8 for desktop, 32 for server) are usually optimal.
• For large datasets, GPUs become advantageous only with a large number of parallel states (typically >256).
• Example: On an i7-6800k (6C12T) CPU and RTX 2080Ti GPU, rans64 encoding speed is similar for CPU and GPU with 128 parallel states, while with 256 parallel states GPU is 2 times faster than CPU.

On CUDA, throughput of the warp-level interleaved path (rans32_16 + num_interleaves=32) is governed by how many warps the batch provides: each stream is decoded by one warp, so occupancy is rows / (32 * SM count). On a 68-SM GPU, rows >= 2048 saturates the GPU; below that, decode throughput scales almost linearly with rows (measured: 6.6 -> 25.8 Gsymbol/s going from 256 to 2048 rows, inverse_cdf_precision="auto"). See Performance Tuning for the full checklist.

Command-line benchmark tool

After installing the package, you can run a built-in benchmark tool directly with Python:

python -m torch_ans.benchmark

This runs the rANS throughput benchmark across the specified batch sizes and devices. The benchmark supports push, pop, and both modes with --mode.

Example:

python -m torch_ans.benchmark -d cpu cuda --mode push

Example output on i7-6800k (6C12T) CPU and RTX 2080Ti GPU:

torch_ans benchmark
  impl: rans64
  mode: push
  data size: 50.0 MB
  num symbols: 256
  freq precision: 16
 
Benchmarking on cpu:
num threads: 6
  Batch size:    1 | Time: 0.1711 s | Throughput: 292.26 MB/s
  Batch size:    2 | Time: 0.1121 s | Throughput: 446.11 MB/s
  Batch size:    4 | Time: 0.0794 s | Throughput: 629.38 MB/s
  Batch size:    8 | Time: 0.0722 s | Throughput: 692.59 MB/s
  Batch size:   16 | Time: 0.0809 s | Throughput: 618.17 MB/s
  Batch size:   32 | Time: 0.0678 s | Throughput: 737.36 MB/s
  Batch size:   64 | Time: 0.0664 s | Throughput: 753.02 MB/s
  Batch size:  128 | Time: 0.0627 s | Throughput: 797.97 MB/s
  Batch size:  256 | Time: 0.0534 s | Throughput: 937.06 MB/s
  Batch size:  512 | Time: 0.0501 s | Throughput: 998.68 MB/s
  Batch size: 1024 | Time: 0.0473 s | Throughput: 1058.06 MB/s
  Batch size: 2048 | Time: 0.0479 s | Throughput: 1044.59 MB/s
  Batch size: 4096 | Time: 0.0500 s | Throughput: 999.69 MB/s

Benchmarking on cuda:
  Batch size:    1 | Time: 6.5701 s | Throughput: 7.61 MB/s
  Batch size:    2 | Time: 4.8248 s | Throughput: 10.36 MB/s
  Batch size:    4 | Time: 1.4710 s | Throughput: 33.99 MB/s
  Batch size:    8 | Time: 0.7357 s | Throughput: 67.96 MB/s
  Batch size:   16 | Time: 0.3679 s | Throughput: 135.90 MB/s
  Batch size:   32 | Time: 0.1841 s | Throughput: 271.66 MB/s
  Batch size:   64 | Time: 0.0922 s | Throughput: 542.46 MB/s
  Batch size:  128 | Time: 0.0463 s | Throughput: 1081.05 MB/s
  Batch size:  256 | Time: 0.0235 s | Throughput: 2128.31 MB/s
  Batch size:  512 | Time: 0.0128 s | Throughput: 3898.49 MB/s
  Batch size: 1024 | Time: 0.0081 s | Throughput: 6149.65 MB/s
  Batch size: 2048 | Time: 0.0081 s | Throughput: 6161.39 MB/s
  Batch size: 4096 | Time: 0.0080 s | Throughput: 6222.82 MB/s

This will help you determine the best batch size (number of parallel states) for your hardware and data size. For most users, start with CPU and increase batch size until speed plateaus or memory usage becomes excessive.

Low-level API Reference

All low-level functions operate directly on PyTorch tensors for maximum performance and flexibility. Below are the main entry points:

torch_ans.rans_init_stream*

torch_ans.rans*_init_stream(size: int, num_interleaves: int, preallocate_size: int) -> torch.Tensor

Initializes a parallel rANS stream tensor.

• size (int): Number of parallel states (batch size).
• num_interleaves (int): Number of interleaved states per stream.
• preallocate_size (int): Preallocated stream size (bytes). Optional; improves performance for large batches.
• Returns: torch.Tensor of shape (size, L).

torch_ans.rans_push*

torch_ans.rans*_push(
  stream: torch.Tensor,
  symbols: torch.Tensor,
  indexes: torch.Tensor,
  cdfs: torch.Tensor,
  cdfs_sizes: torch.Tensor,
  offsets: torch.Tensor,
  freq_precision: int,
  bypass_coding: bool,
  bypass_precision: int
) -> None

Encodes symbols into the rANS stream in parallel.

• stream (torch.Tensor): rANS stream tensor.

• symbols (torch.Tensor): Symbols to encode (int tensor).

• indexes (torch.Tensor): CDF table index for each symbol.

• cdfs (torch.Tensor): Batched CDF tables.

• cdfs_sizes (torch.Tensor): Size of each CDF table.

• offsets (torch.Tensor): Offset for each distribution.

• freq_precision (int): Frequency precision in bits.

• bypass_coding (bool): Enable bypass coding for out-of-range symbols.

• bypass_precision (int): Precision for bypass coding.

  Range limit: an out-of-range symbol is zig-zag encoded into a raw value stored in the symbol dtype (int32 by default), so its distance from the coded range [offset, offset + cdf_size - 2) must stay below 2^(bits-2) — about 2^30 (~1.07e9) for int32. Symbols further out of range overflow the raw value and are silently mis-coded.

torch_ans.rans_pop*

torch_ans.rans*_pop(
  stream: torch.Tensor,
  indexes: torch.Tensor,
  cdfs: torch.Tensor,
  cdfs_sizes: torch.Tensor,
  offsets: torch.Tensor,
  freq_precision: int,
  bypass_coding: bool,
  bypass_precision: int
) -> torch.Tensor

Decodes symbols from the rANS stream in parallel.

• stream (torch.Tensor): rANS stream tensor.
• indexes (torch.Tensor): CDF table index for each symbol.
• cdfs (torch.Tensor): Batched CDF tables.
• cdfs_sizes (torch.Tensor): Size of each CDF table.
• offsets (torch.Tensor): Offset for each distribution.
• freq_precision (int): Frequency precision in bits.
• bypass_coding (bool): Enable bypass coding for out-of-range symbols.
• bypass_precision (int): Precision for bypass coding.
• Returns: Decoded symbols as torch.Tensor.

torch_ans.rans_stream_to_byte_strings

torch_ans.rans_stream_to_byte_strings(stream: torch.Tensor) -> List[bytes]

Converts a stream tensor to a list of Python bytes objects for serialization or storage.

torch_ans.rans_byte_strings_to_stream

torch_ans.rans_byte_strings_to_stream(byte_strings: List[bytes]) -> torch.Tensor

Reconstructs a stream tensor from a list of Python bytes objects.

See torch_ans_test.py for more advanced usage and validation.

High-level API Reference

The high-level API provides a user-friendly interface for ANS coding similar to CompressAI, abstracting away tensor management and stream details. The main entry point is TorchANSInterface in torch_ans.utils, which supports encoding and decoding with various rANS and tANS implementations.

TorchANSInterface

from torch_ans.utils import TorchANSInterface

Initialization:

coder = TorchANSInterface(impl="rans64", freq_precision=16, device="cpu")

• impl (str): coder variant (e.g., "rans64", "rans32", "rans32_16", "tans32", "tans16", "tans8", etc.)
• freq_precision (int): Frequency precision in bits.
• device (str): "cpu" or "cuda".
• ... (More parameters in torch_ans/utils.py)

Parameter Setup:

coder.init_params(freqs, num_freqs, offsets)

• freqs (torch.Tensor): Batched distribution probability mass functions (PMFs), shape (K, S). K is number of distributions/tables, S is maximum size of symbol alphabet.
• num_freqs (torch.Tensor): Size of symbol alphabet per distribution, shape (K,).
• offsets (torch.Tensor): Offset for each alphabet, shape (K,).

Encoding:

encoded = coder.encode_with_indexes(symbols, indexes)

• symbols (torch.Tensor): Symbols to encode, shape (B, N). B is the same with ANS state batch size (equal to number of parallel states), N is number of symbols.
• indexes (torch.Tensor): Distribution/table index for each symbol, should have same shape with symbols.
• Returns: Encoded stream (typically a list of bytes or tensor).

Decoding:

decoded = coder.decode_with_indexes(encoded, indexes)

• encoded: Encoded stream from encode_with_indexes.
• indexes (torch.Tensor): Distribution/table index for each symbol, shape (B, N). B is the same with ANS state batch size (equal to number of parallel states), N is number of symbols.
• Returns: Decoded symbols as torch.Tensor.

Other Features:

• Supports both CPU and CUDA parallel coding transparently.
• Handles CDF/PMF conversion internally during init_params.
• Provides methods for serialization and deserialization of streams.

See torch_ans_test.py for advanced usage, batch processing, and custom coding schemes.

Full API Reference Table

Function/Class	Description
rans*_init_stream	Initialize rANS stream tensor
rans*_push	Encode symbols into rANS stream
rans*_pop	Decode symbols from rANS stream
tans*_init_stream / tans*_push / tans*_pop	The tANS families' operators (identical signatures to the rANS ones; cdfs carries the spread table)
tans_build_tables	Build the FSE-style spread tables shared by the three tANS families
rans_stream_to_byte_strings	Convert stream tensor to list of bytes
rans_byte_strings_to_stream	Convert list of bytes to stream tensor
TorchANSInterface	High-level API for encoding/decoding with rANS and tANS variants
TorchANSInterface.init_params	Set PMF, alphabet size, and offsets for coder
TorchANSInterface.encode_with_indexes	Encode symbols using high-level API
TorchANSInterface.decode_with_indexes	Decode symbols using high-level API

Cookbook

A high-level usage example with TorchANSInterface:

import torch
from torch_ans.utils import TorchANSInterface

# Prepare PMF and convert to quantized CDF
batch_size = 8
num_dists = 8
num_symbols = 256
pmf = torch.randint(1, 100, (num_dists, num_symbols), dtype=torch.int32)
num_freqs = torch.full((num_dists,), num_symbols, dtype=torch.int32)
offsets = torch.zeros(num_dists, dtype=torch.int32)

# Create coder interface
coder = TorchANSInterface(impl="rans64", freq_precision=16, device="cpu")
coder.init_params(pmf, num_freqs, offsets)

# Prepare symbols and indexes
seq_len = 1024
symbols = torch.randint(0, num_symbols, (batch_size, seq_len), dtype=torch.int32)
indexes = torch.randint(0, num_dists, (batch_size, seq_len), dtype=torch.int32)

# Encode
encoded = coder.encode_with_indexes(symbols, indexes)

# Decode
decoded = coder.decode_with_indexes(encoded, indexes)

assert torch.equal(decoded, symbols)

# Fastest GPU decode path: warp-level 32-way interleaved coding + auto lookup
gpu_coder = TorchANSInterface(impl="rans32_16", freq_precision=15, device="cuda",
                              num_interleaves=32, inverse_cdf_precision="auto")
gpu_coder.init_params(freqs, num_freqs, offsets)
gpu_stream = gpu_coder.encode_with_indexes(
    symbols.to("cuda"), indexes.to("cuda"))
gpu_decoded = gpu_coder.decode_with_indexes(gpu_stream, indexes.to("cuda"))
assert torch.equal(gpu_decoded.cpu(), symbols)

Notes on the GPU example:

• num_interleaves=32 selects the warp-level kernels (one warp per stream).
• freq_precision=15 is the maximum supported by the warp-level path.
• inverse_cdf_precision="auto" builds a symbol-lookup table sized from the alphabet (and skips it for alphabets < 64 symbols).
• Batch at least ~2048 rows (symbols.shape[0]) to saturate the GPU; streams encoded on CPU decode on CUDA and vice versa.

Performance Tuning

In order of impact (measured on an RTX 2080 Ti and a 6-core CPU):

GPU

1. Use impl="rans32_16" with num_interleaves=32 — this selects the warp-level interleaved CUDA kernels (one warp per stream). It requires freq_precision <= 15.
2. Batch enough rows. Each stream is decoded by one warp, so occupancy is rows / (32 * SM count): use rows >= 4 * SM count (e.g. >= 2048 on a 68-SM GPU) to saturate the GPU. Below that, decode throughput scales almost linearly with rows, and no launch-configuration trick can recover it.
3. Enable the inverse-CDF lookup for faster decoding with inverse_cdf_precision="auto": for alphabets with >= 64 symbols it builds a 2**q table (q ~= log2(alphabet) - 1) sized so the whole row fits the shared-memory staging budget. Worth +5% to +80% over the default binary search depending on the alphabet.
4. Keep freq_precision <= 15 — smaller states renormalize less often, and rans32_16 is the only variant with the warp-level path.
5. Avoid bypass coding when symbols are guaranteed in range — the bypass phase is inherently sequential across the warp lanes (its per-symbol skeleton was worth 2.5-3.3x in our ablations).
6. init_params rebuilds the quantized CDFs and the lookup table and costs ~0.2 ms on CUDA. It runs on every call of the dist_freqs API, so for many small tensors prefer passing dist_indexes (precomputed CDF indexes), or reuse the coder.

On CUDA (RTX 2080 Ti, rans32_16, 1024 rows x 4096 symbols, alphabet 255, p=15, Gsymbol/s):

lookup	1-way enc/dec	4-way enc/dec	32-way enc/dec
divided search	1.16 / 0.84	3.69 / 1.64	10.10 / 6.19
inverse-CDF	1.57 / 1.42	3.66 / 2.30	10.49 / 8.85
alias sampling	1.01 / 1.24	1.98 / 2.56	n/a

1. Interleaving matters far more on CUDA than on CPU: 4-way is ~3x and the 32-way warp path is ~7-9x over 1-way (each stream is coded by one warp with a shared cursor, which also coalesces the stream traffic).
2. Decoding: the inverse-CDF table wins at every interleave factor (1.7x over the divided search at 1-way with 256 distributions, 1.4x at 32-way). Alias sampling is also faster than the divided search (1.5x) but slower than the table, and it scales worst (~2x from 1-way to 4-way, against ~3x) because its bucket lookup is a division by a runtime value and it needs an extra dependent table load.
3. Encoding: alias sampling is the slowest by up to 1.9x - the push reads a remap table indexed by the whole 2**freq_precision range, which is 33 MB for 256 distributions. There is also no 32-way alias variant.
4. So the recommendation is the same on both devices: inverse_cdf_precision="auto" (or a dense table) for decode, and no alias sampling except where its 1-way convenience is worth the throughput. Alias needs cdfs_sizes - 1 to be a power of two, e.g. alphabet 255 with bypass coding (a cdf row of 257).

CPU

1. torch.set_num_threads(n) controls the OpenMP parallelism used by push/pop.

2. inverse_cdf_precision="auto" also helps on CPU decoding — up to ~2.5x at num_interleaves=1 and ~1.2-1.6x on top of the interleaved gain (see the next point); the table is built with a single sorted merge, so it is cheap.

3. Interleaved variants boost CPU performance via OoO (Out-of-Order) execution — with num_interleaves=4, the independent rANS states per stream let the out-of-order engine overlap their renormalize/push/pop chains. Measured on an i7-6800K (1M symbols, alphabet 256, 256 distributions, single thread, best-of-5, scripts/bench_interleaves_cpu.py):

   • Encoding: ~1.25-1.4x with num_interleaves=4 (1.5x without bypass coding).
   • Decoding: ~1.8-2.0x with the default (divided-search) symbol lookup, ~1.4-1.8x with the inverse-CDF table — the table's absolute decode throughput is the highest of the two (about 1.2-1.6x better than the divided search at the same num_interleaves), it just has less left to gain from interleaving because a single state is already fast.
   • num_interleaves=8 performs about the same as 4; num_interleaves=32 is not useful on CPU (encode ~0.8x, register pressure) — it exists for the warp-level CUDA path (the same holds for the tANS 32-way kernels, whose warp path is the device one). Use 4 (or 8) on CPU.
   • The inverse-CDF table (inverse_cdf_precision="auto" or an integer) is the fastest CPU decode configuration measured, at every alphabet size and every num_interleaves, provided the auto precision is used (a dense table is far worse: for 256 distributions it is tens of MB and thrashes the cache). This required removing the data-dependent linear walk that used to close the table's gap — see below.
   • The gains come from the fact that the coder is written to be overlap-friendly — see Why interleaving helps below.

   Why interleaving helps (and what it needs). Interleaved coding only pays off if the independent states can actually execute in parallel, and three things used to prevent that:

   • Data-dependent branches. Whether a renormalization is needed is close to a coin flip, so the renormalize if mispredicts on about every second symbol and flushes the out-of-order window — exactly the window the other lanes need. The coders therefore use branch-free renormalization (a speculative store on the encode side, an unconditional load on the decode side), a branch-free monotone search for the symbol lookup, and a branch-free way to close the inverse-CDF table's gap: the table gives a lower bound and the remaining couple of symbols are resolved by loading a small fixed window of cdf entries in parallel and counting how many are <= cum_freq, instead of walking one cdf entry at a time.
   • Non-pipelined 64-bit division. The integer divider is not pipelined, so N interleaved states merely queue on it. Encoding now uses a pipelined floating-point divide plus a rarely-taken exact integer fix-up.
   • Aliasing and table lookups. The states live in registers (locals accessed with a compile-time lane index) instead of in the stream tensor they are written to, and all per-lane table lookups (index -> cdf row, symbol -> (start, freq)) are gathered for the whole interleave group before any state is updated, so the cache misses and the symbol searches of different lanes overlap.

   Note the branch-free renormalization stores one word past the cursor speculatively and the launch configuration keeps the interleaved states inline in the stream buffer, so every buffer keeps one word of slack (see rans_init_stream / rans_push).

4. rans64 remains the most robust default; rans32_16 reduces the initial-state overhead for small payloads. For CPU decode throughput the best combination measured is inverse_cdf_precision="auto" + num_interleaves=4 (or 8); for encode, num_interleaves=4 (the table does not affect encode). See scripts/bench_interleaves_cpu.py to reproduce.

5. Symbol lookup choice. There are three decode lookups; on CPU the inverse-CDF table wins, and alias sampling is the slowest (it is opt-in via alias_sampling=True). Measured on an i7-6800K (1M symbols, 256 distributions, alphabet 255, p=15, single thread, num_interleaves=4, cyc/symbol):

   lookup	encode	decode	tables
   divided search (default)	25.8	63.9	cdf 257 KB
   inverse-CDF ("auto")	26.7	54.4	385 KB
   alias sampling	101.6	75.8	1.3 MB + 32 MB remap

   Alias sampling does benefit from interleaving, just less than the other two (decode 118.2 -> 75.8 cyc/symbol from 1 to 4 interleaved states, i.e. 1.56x, against 1.76x for the inverse-CDF table; encode 212.9 -> 101.6, 2.1x), for two structural reasons: its bucket lookup divides by a runtime value (a non-pipelined integer division - it is a shift here because the builder only accepts a cdf size of 2**k + 1, but the shift has to be recovered per symbol), and its encode side reads a remap table indexed by the whole 2**freq_precision range, which is 32 MB for 256 distributions and misses cache on every symbol. It also requires cdfs_sizes - 1 to be a power of two (e.g. alphabet 255 with bypass coding), unlike the other lookups.

Troubleshooting

• Ensure PyTorch is installed and compatible with your Python version. pybind11 itself is not a dependency: the build takes its headers from torch/include/pybind11 and the stable-ABI path needs none.
• For CUDA support, verify CUDA toolkit is installed and torch.cuda.is_available() returns True.
• If you encounter compilation errors, try cleaning the build directory and reinstalling:

  rm -rf build/ torch_ans.egg-info/
  pip install .
  

FAQ

Algorithmic

Q: What is the difference between parallel and interleaved ANS state?

A: Parallel states corresponds to individual bitstreams, so they require extra space to store bitstream lengths. Despite from this drawback, parallel states are more efficient (especially with a large number of ANS states), easy to implement (device specific code is hardly required) and robust to bit error (corruption in one bitstream will not propagate to another).

In constrast, interleaved states corresponds to a single bitstream, and is sequentially processed. See Fabian's paper for more information. In fact, SIMD ops (such as AVX2 and AVX512) could be used to accelerate interleaved ANS coding, but in our experiments 8 parallel states has better acceleration than the 8 interleaved states with AVX2 ops.

Q: How to choose parameters like State Bits, Stream Bits and Freq Bits? What is their relation to stream size and memory occupation?

A: To be simple, always use rans64 configuration if data size is not extremely small (less than 1KB) and you are not sensitive about memory occupation! If data size is extremely small, you could try rans32 or rans32_16 to reduce initial state.

Q: Why isn't GPU throughput significantly higher than CPU? Some existing implementation like dietgpu achieve over 200GB/s but torch_ans have only 6GB/s throughput!

A: This was true for the old non-interleaved path (one thread per stream, divergent branches, DEFAULT_NUM_THREADS_PER_BLOCK=1), which still exists for rans64/rans32. The rans32_16 variant now has a dedicated warp-level interleaved CUDA path (technique from Recoil/dietgpu): each stream is decoded by one warp, the CDF and lookup tables are staged in shared memory, the symbol lookup uses a tunable inverse-CDF table, and the bypass phase is skipped by a group ballot when no lane needs it.

On an RTX 2080 Ti this reaches ~25.8 Gsymbol/s decode (~29 GB/s of payload at ~9 bits/symbol) with rows >= 2048, which is the same order of magnitude as dietgpu's numbers. Two things to keep in mind when comparing:

• rANS decode throughput in symbols/s depends on the payload entropy: at ~9 bits/symbol, 25.8 Gsymbol/s already moves ~29 GB/s; benchmarks reporting 200 GB/s usually assume fixed 2-byte symbols on data-center GPUs (A100/H100) and no per-symbol distribution routing or bypass coding.
• torch_ans targets the learned-compression use case: multiple distributions per tensor, int32 symbols with per-symbol CDF indexes, and bypass coding for out-of-range values.

Technical

Q: What Python and PyTorch versions are supported?

A: Python >= 3.7 and PyTorch >= 1.10 are recommended. Earlier versions may work but are not tested.

Q: What platforms are supported?

A: torch_ans supports Linux (x86_64/aarch64), macOS (x86_64/aarch64), and Windows (x86_64) platforms. Ensure you have a compatible C++ compiler and optionally CUDA toolkit (for GPU support) installed on your system.

Q: How do I enable CUDA support?

A: Install the CUDA toolkit and ensure PyTorch is built with CUDA. Use WITH_CUDA=1 pip install . to enable CUDA build. Also check with torch.cuda.is_available().

Q: RuntimeError: The detected CUDA version (x.x) mismatches the version that was used to compile PyTorch (y.y). Please make sure to use the same CUDA versions.

A: Ensure that your installed CUDA toolkit version align with PyTorch CUDA version. Or if you do not need CUDA for coding, remove WITH_CUDA environment variable to disable CUDA when building.

Q: Why do I get compilation errors during installation?

A: Make sure you have a C++17 compiler, pybind11, and compatible PyTorch. Try cleaning the build directory: rm -rf build/ torch_ans.egg-info/ and reinstall.

Q: ImportError: ...torch_ans/_C... Undefined Symbol xxxx.

A: In most cases this is caused by version mismatch between PyTorch during runtime and building torch_ans. Reinstalling torch_ans would hopefully solve this.

Q: How do I choose between CPU and GPU?

A: Set the device argument in API calls to "cpu" or "cuda". For large batches, GPU is recommended; for small data, CPU may be faster. Also, as GPU coding process is asynchronous, if some other tasks (such as neural networks in neural compression) are running meanwhile, using GPU coding may increase the overall throughput. For the fastest GPU decode when you have massive data, use impl="rans32_16", num_interleaves=32, freq_precision <= 15, inverse_cdf_precision="auto" and rows >= 2048 (see the Performance Tuning section above).

Q: What is the difference between low-level and high-level APIs?

A: Low-level APIs operate directly on tensors for maximum control, but improper use may lead to segmentation fault. High-level APIs (TorchANSInterface) is designed with a similar API to CompressAI's entropy coder, which simplify usage and enables auto paralellization, ANS stream management and serialization for you.

Q: How do I serialize and deserialize streams in torch.Tensor format?

A: Use torch_ans.rans_stream_to_byte_strings and torch_ans.rans_byte_strings_to_stream for conversion between tensors and Python bytes.

Q: How do I run tests?

A: Run python tests/torch_ans_test.py or use pytest in the project directory.

Q: Where can I find more usage examples?

A: See the README Usage section and torch_ans_test.py for practical examples and test cases. For practical usage of ANS coding in neural compression, see CompressAI.

Q: What can I do if segmentation fault occurs?

This library is developed for research-purpose only, and not as a robust everyday software, so errors may occur occasionally. If you are using the low-level API, try to switch to the high-level API and enable bypass coding, which might be more robust. If error still persists, try checking middle variables such as cdf (should be incremental positive integers and less than 2^freq_precision).

Development

• Requires Python >= 3.7, PyTorch, and a C++17 compatible compiler.
• CUDA toolkit required for GPU support.
• Run tests with:

  pytest -q
  

Known Issues

• Rans64 cuda coding test fails on some newer GPU architectures.
• pip install may compile the extension locally, which needs a C++ toolchain (and nvcc for the CUDA coding paths). That is what happens when no prebuilt wheel matches your platform, or when your torch is older than 2.10: a wheel compiled against ATen cannot be loaded by another torch build, so shipping one would break the import on every other version. The exception is the prebuilt stable-ABI wheel, which is published to PyPI and serves every torch >= 2.10 (see Prebuilt stable-ABI wheel).
• CUDA coding is off unless the extension was built with WITH_CUDA=1: a CPU-only extension on a GPU machine only codes CPU tensors. Asking it for CUDA tensors triggers the automatic local CUDA build described in CUDA support; if that is not possible you get the commands to install/fix it.
• CUDA build fails with std_function.h: parameter packs not expanded with '...': nvcc <= 12.1 cannot parse the std::function headers shipped with libstdc++ from GCC >= 11.4 (e.g. Ubuntu 22.04.3+), so compiling CUDA extensions with the default g++ fails. Fixes (any one of them):
  • Use an older host compiler for nvcc, e.g. g++-10 (install with apt install g++-10). The runtime dynamic build detects this automatically: when a CUDA build fails, it retries with g++-10/g++-9/g++-8 (-ccbin) before falling back to CPU-only, and remembers the working configuration in the torch extensions cache directory (cuda_build_state).
  • Use nvcc >= 12.2, which supports the newer libstdc++ headers.
  • Also make sure the CUDA toolkit version matches the one your PyTorch wheel was built with (e.g. nvcc 11.8 for torch ...+cu118); mismatched toolchains are unsupported.

Testing and coverage

Install the development dependencies and run the Python unit tests:

pip install . --no-build-isolation
pip install pytest pytest-cov gcovr
pytest -q

Run Python coverage for the package:

pytest --cov=torch_ans --cov-report=term-missing --cov-report=html

For release-to-release throughput comparisons see benchmark_status.md: it is measured locally (CI runners are virtualised CPUs and too noisy for this) with

python scripts/bench_version_matrix.py --version-label <ver> --json <ver>.json
python scripts/generate_benchmark_report.py --baseline <old>.json \
    --candidate <new>.json --output benchmark_status.md

To collect native C/C++ coverage for the compiled extension, build with coverage instrumentation and run tests from the repository root:

ENABLE_COVERAGE=1 python setup.py build_ext --inplace
ENABLE_COVERAGE=1 pytest -q

Then generate a native coverage report with gcovr:

gcovr -r . --html-details -o native_coverage.html

This produces Python coverage output in htmlcov/ and native C/C++ coverage output in native_coverage.html.

Runtime dynamic build

If a compiled extension is already present (for example one built by pip install . --no-build-isolation), it is used as is. Otherwise - e.g. when the package was installed without a compiler - torch_ans compiles the native C++/CUDA sources at runtime using PyTorch's cpp_extension, on the first use of an operator.

An extension that cannot be loaded (built against another PyTorch, or with a different ABI) does not break the import: torch_ans warns and falls back to compiling one for the local torch. Set TORCH_ANS_FORCE_RUNTIME_BUILD=1 to always take that path and ignore any pre-built extension.

• Trigger runtime build programmatically:

import torch_ans

• Notes:
  • The build is triggered by the first use of a native operator: creating a TorchANSInterface (or torch_ans._C.<operator> / from torch_ans._C import <operator>, which compile the full extension). Importing torch_ans.utils alone no longer compiles anything. The built module is cached under the torch extensions directory (~/.cache/torch_extensions/) and reused by later processes.
  • Incremental compilation (default): only the operators the interface needs are compiled. TorchANSInterface(impl="rans64") with default settings builds the rans64 push/pop plus the shared helpers rather than all three rANS variants with every interleave factor and symbol lookup (about 4x less compiler work; measured 25s vs 100s for that configuration). Every configuration has its own cached module (torch_ans_C_r64, torch_ans_C_r64_i4_invcdf, ...), and any already loaded module that exports the requested operators is reused, so switching configurations does not duplicate work and a previously built full extension satisfies all of them. Pass incremental_compile=False or set TORCH_ANS_INCREMENTAL_COMPILE=0 to compile every operator up front (this is also what a wheel built at install time always contains).
  • Runtime compilation requires a C/C++ toolchain and ninja (installed automatically as a dependency).
  • The runtime build tries CUDA first and falls back to CPU-only automatically when no usable CUDA runtime or toolchain is available. For CUDA builds you must have a compatible CUDA toolkit and driver installed (see Known Issues for nvcc/GCC compatibility).
  • CPU parallel coding relies on OpenMP: at::parallel_for in the native code is multi-threaded only when the extension is compiled with OpenMP enabled. The runtime build enables it automatically on Linux (GCC -fopenmp) and Windows (MSVC /openmp); on macOS it uses Homebrew libomp (-Xpreprocessor -fopenmp -lomp) when installed and the coding falls back to torch's own thread pool when it is missing - or when a module built against it turns out to abort at use time, which the build checks once in a child process before it accepts the module (the outcome is remembered in an openmp_build_state file in the torch extensions cache directory).
  • If your runtime PyTorch ABI differs from the build-time one, torch_ans will warn by default. Set TORCH_ANS_STRICT_CHECK=1 to re-enable a strict ImportError on mismatch.
  • A compiled torch_ans/_C*.so shadows torch_ans/_C.py completely, so the lazy implementation lives in torch_ans/_lazy_C.py (always importable). That is what the fallbacks switch to: an extension that fails to load, and a CPU-only extension that is asked to code CUDA tensors (see CUDA support).
  • To keep CI/tests stable, the bundled test for dynamic build is guarded; enable it with RUN_DYNAMIC_BUILD_TEST=1 when you want to run the rebuild test locally.

Prebuilt stable-ABI wheel

Every release publishes a prebuilt wheel to PyPI: since v0.4.0 the native code can be compiled once against the libtorch stable ABI (TORCH_TARGET_VERSION = torch 2.10) instead of against each user's torch, so one .so serves every torch >= 2.10 - no compiler, no first-use build, no per-torch-version pinning.

pip install torch_ans            # gets the prebuilt wheel when one matches

Published files of a release (CPU flavour - the CUDA flavour is not published yet):

platform	wheel tag
Linux x86_64	py3-none-manylinux_2_28_x86_64
macOS arm64	py3-none-macosx_11_0_arm64
Windows AMD64	py3-none-win_amd64

plus the sdist, which is what pip falls back to everywhere else - including macOS x86_64: PyTorch publishes no Intel macOS wheel above 2.2, so that platform has no torch >= 2.10 for this wheel to target and gets the sdist.

Building one locally works the same way:

# CPU flavour: builds torch_ans/_C_cpu.so (schemas + CPU and catch-all implementations)
WITH_STABLE_ABI=1 pip install . --no-build-isolation

# CUDA flavour: additionally builds torch_ans/_C_cuda.so (CUDA implementations only)
WITH_STABLE_ABI=1 WITH_CUDA=1 pip install . --no-build-isolation

Why the wheel does not declare torch>=2.10: pip reads the dependencies of the single best-matching wheel of a version, and a platform wheel beats a py3-none-any one, so declaring it would make pip upgrade torch (hundreds of MB) for everyone on an older release. The requirement is therefore enforced at runtime: with torch < 2.10 the prebuilt library is skipped with one RuntimeWarning and the ordinary runtime build runs - the same behaviour a source install has always had. Set TORCH_ANS_NO_STABLE_ABI=1 to silence it, or pip install -U 'torch>=2.10' to use the wheel.

What CI checks (scripts/check_stable_abi_wheel.py, tests/test_stable_abi.py, tests/test_cross_version_load.py):

• the libraries export no CPython and no ATen/C++ symbols - only aoti_torch_* / torch_* from the stable C shim, which is what makes them independent of both the torch release and the Python version;
• the wheel is therefore tagged py3-none-<platform>; the cpXY tag bdist_wheel derives from the mere presence of ext_modules is corrected in CI with python -m wheel tags --python-tag py3 --abi-tag none;
• the CPU and CUDA libraries register under the same operator names (torch.ops.torch_ans.*), so the dispatcher picks the kernel and nothing on the Python side has to know which flavour is installed;
• the same wheel produces a byte-identical stream under every tested torch version (the cross-version-load CI job: 2.10 = the ABI floor, the build version, and the newest release);
• release_wheels.yml (tag push) builds the four wheels with cibuildwheel, retags them to py3-none-<platform>, checks them, installs and uses each one on its own platform, and then uploads them next to the sdist that release_sdist.yml publishes. cibuildwheel.yml stays the ABI regression gate and publishes nothing.

Runtime behaviour and knobs:

• installation is automatic: torch_ans/_stable_C.py looks for _C_cpu* / _C_cuda* next to the package (or in TORCH_ANS_STABLE_LIB_DIR) and loads them with torch.ops.load_library, preferring them over any compilation. A failed load is a warning plus the normal runtime build - never an error;
• TORCH_ANS_NO_STABLE_ABI=1 ignores the prebuilt extension (debugging/CI); TORCH_ANS_STABLE_LIB_DIR=<dir> points the search somewhere else (used by the tests, which build the libraries outside the source tree);
• incremental_compile= has no effect on this path: the prebuilt extension contains every operator, and the profiles exist to save compile time - of which there is none here;
• the API is unchanged (same operator names, same defaults). The one internal difference: the stream capacity is reserved by the Python layer and the native push only compares it (used_bytes), which keeps aten::resize_ and a device->host synchronisation out of every CUDA push (a small push on an RTX 2080 Ti went from 157 to 39 us; scripts/bench_stream_capacity.py reproduces that).

Source distributions are unchanged: without a prebuilt .so the package installs as pure Python and compiles per torch at first use, exactly as before. torch < 2.10 always takes that path, even when the wheel was installed - with the one warning quoted above, and nothing else changes.

A CPU wheel on a GPU machine is not a dead end either: asking for CUDA coding makes utils._ensure_native_for compile a CUDA extension for the local torch (TORCH_ANS_FORCE_RUNTIME_BUILD=1 does the same unconditionally). The runtime build is a pybind11 module, so it cannot collide with the operator schemas the prebuilt library registered in the dispatcher.

TODO

• Docs and examples for high-level API
• Fix interleave ANS state on CUDA, and possibly implement with CUDA Warp-level Primitives Done for rans32_16_i32 / rans32_16_i4: warp-level 32-way (and 4-way sub-warp) interleaved coding on CUDA, bit-compatible with the CPU interleaved streams (see torch_ans/rans_warp_cuda.cuh, technique referenced from Recoil)
• Implement lookup table logic in push/pop steps for possible acceleration Done: inverse-CDF lookup with a tunable (dense/sparse/"auto") 2**q table on CPU and CUDA, warp-level shared-memory staging on CUDA, and a "auto" precision heuristic.
• Implement tANS and its variants with similar high-level API (refer to FSAR) Done: impl="tans32" / "tans16" / "tans8" on the CPU and on CUDA, bit-compatible with the CPU kernels in both shapes (one thread per row at one lane, one lane group per row at 1/2/4/8/32 - see torch_ans/tans_warp_cuda.cuh), with the FSE-style spread table built by tans_build_tables, a bit buffer in the stream header and the same capacity contract / high-level API as rANS (see tANS).
• Test other backends supported in PyTorch (such as ROCm)
• Add more examples such as neural compression

Related Projects

• PyTorch: Deep learning framework used for tensor operations and GPU support.
• pybind11: Python bindings for C++ used in this extension.
• CompressAI: Neural compression library with entropy coding.
• torchac: Fast arithmetic coding library for PyTorch, includes ANS variants.
• ryg_rans: Reference C implementation of rANS.
• fsc: rANS implementation with alias mapping.
• dietgpu: An ultra-fast parallel rANS implementation (over 200GB/s) on NVIDIA GPUs.
• Recoil: A header-only C++20 rANS library including parallel rANS algorithms.
• FiniteStateEntropy: Reference C implementation of tANS/FSE by Yann Collet.
• FSAR: Provides a numpy-based unified rANS/tANS interface implementation.

For more ANS implementations, see Jarek Duda's blog on encode.su

License

MIT License