Registry / ai-ml / colt5-attention

colt5-attention

library0.11.1pypypiunverified

Conditionally Routed Attention, an implementation of the CoLT5 architecture for efficient long-context transformers. Current version 0.11.1, rapid release cadence.

ai-ml

pip install colt5-attention

github ↗package ↗

INSTALL

IMPORT

SIG · COLT5-ATTENTION

colt5-attention

ai-mlpythonv0.11.1

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

CoLT5Attention

✓ from colt5_attention import CoLT5Attention

Standard import path

ConditionalRouting

✓ from colt5_attention import ConditionalRouting

Exported from colt5_attention

Initialize CoLT5 attention with routed queries/key-values and run forward pass.

import torch from colt5_attention import CoLT5Attention attn = CoLT5Attention( dim=512, num_routed_queries=64, num_routed_key_values=64, num_heads=8, dropout=0.1 ) x = torch.randn(2, 1024, 512) out = attn(x) print(out.shape) # (2, 1024, 512)

Debug

Known issues

breakingVersion 0.11.0 changed default value of `num_routed_queries` from 128 to 64. Existing code relying on default may see different memory/performance.

fix

Explicitly set `num_routed_queries=128` if you need the old behavior.

affects: >=0.11.0

gotchaAttention masking must use a float mask (e.g., 0.0 for keep, -inf for mask). Boolean masks are not supported and may silently produce wrong outputs.

fix

Convert boolean mask to float: `attn_mask = 0.0 * keep_mask + (-1e9) * (~keep_mask)`

affects: all

gotchaInput tensor must be contiguous. Non-contiguous tensors can cause runtime error or incorrect gradient propagation.

fix

Call `x = x.contiguous()` before passing to attention.

affects: all

Upgrade

Version history

0.11.1latest on PyPI · released Sep 6, 2024

Audit

Dependencies

torchrequiredCore dependency for tensor operations and neural network modules

einopsrequiredUsed for tensor reshaping and attention operations

Agent activity

16 hits · last 30 days

node

Amazon

OpenAI (training)

Resources

githubhttps://github.com/lucidrains/CoLT5-attention ↗

packagehttps://pypi.org/project/colt5-attention/ ↗