light-recurrent-unit-pytorch

Light Recurrent Unit - Pytorch

Implementation of the Light Recurrent Unit in Pytorch

Install

$ pip install light-recurrent-unit-pytorch

Usage

import torch
from light_recurrent_unit_pytorch import LightRecurrentUnitCell

cell = LightRecurrentUnitCell(256)

x = torch.randn(2, 256)
hidden = torch.randn(2, 256)

next_hidden = cell(x, hidden) # (2, 256)

Single layer

import torch
from light_recurrent_unit_pytorch import LightRecurrentUnitLayer

layer = LightRecurrentUnitLayer(256)

x = torch.randn(2, 1024, 256)

out = layer(x) # (2, 1024, 256)

assert out.shape == x.shape

Stacked

import torch
from light_recurrent_unit_pytorch import LightRecurrentUnit

lru = LightRecurrentUnit(256, depth = 4)

x = torch.randn(2, 1024, 256)

out, layer_hiddens = lru(x) # (2, 1024, 256), List[(2, 256)]

assert out.shape == x.shape

Multi head

With heads > 1, the K parallel candidates are softmax selected before the gate, as in this paper

import torch
from light_recurrent_unit_pytorch import LightRecurrentUnit

lru = LightRecurrentUnit(256, depth = 4, heads = 4)

x = torch.randn(2, 1024, 256)

out, layer_hiddens = lru(x) # (2, 1024, 256), List[(2, 256)]

assert out.shape == x.shape

Test

Enwik8

$ python train.py

Citations

@Article{electronics13163204,
    AUTHOR  = {Ye, Hong and Zhang, Yibing and Liu, Huizhou and Li, Xuannong and Chang, Jiaming and Zheng, Hui},
    TITLE   = {Light Recurrent Unit: Towards an Interpretable Recurrent Neural Network for Modeling Long-Range Dependency},
    JOURNAL = {Electronics},
    VOLUME  = {13},
    YEAR    = {2024},
    NUMBER  = {16},
    ARTICLE-NUMBER = {3204},
    URL     = {https://www.mdpi.com/2079-9292/13/16/3204},
    ISSN    = {2079-9292},
    DOI     = {10.3390/electronics13163204}
}

@article{Merrill2024TheIO,
    title   = {The Illusion of State in State-Space Models},
    author  = {William Merrill and Jackson Petty and Ashish Sabharwal},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2404.08819},
    url     = {https://api.semanticscholar.org/CorpusID:269149086}
}

Metadata

Release files for light-recurrent-unit-pytorch 0.1.4

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for light-recurrent-unit-pytorch 0.1.4
File	Size	Uploaded
light_recurrent_unit_pytorch-0.1.4.tar.gz	6.1 kB	Oct 9, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for light-recurrent-unit-pytorch 0.1.4
File	Interpreter	ABI	Platform	Reset
light_recurrent_unit_pytorch-0.1.4-py3-none-any.whl 6.1 kB Oct 9, 2026	Python 3	none	any	Details

Total release size: 12.3 kB

Release files / light_recurrent_unit_pytorch-0.1.4.tar.gz

Download URL	light_recurrent_unit_pytorch-0.1.4.tar.gz
Size	6.1 kB
Tags	Source
SHA-256 checksum How to use checksums	`0971399815f79aea3c707fa87ea5812ab4cdc508f4ac93cdd9dc76975fa521ba`
BLAKE2b-256 checksum How to use checksums	`04c38058ef369044c9dcbe07e04e8e20d6869c202f2c952adb43deb7217c76c2`
Upload date	Oct 9, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`

Release files / light_recurrent_unit_pytorch-0.1.4-py3-none-any.whl

Download URL	light_recurrent_unit_pytorch-0.1.4-py3-none-any.whl
Size	6.1 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`0d6eb1ed930c60f798820f2e656c4267a42d252f9f09400bff42bb6c34cdbe0b`
BLAKE2b-256 checksum How to use checksums	`c0eae20c889cea5b53d5e8f5b1e0c41dded35319356048f5c00e3ef323c7fc10`
Upload date	Oct 9, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`uv/0.8.17`