minGRU-pytorch

minGRU

Implementation of the proposed minGRU in Pytorch, only the log-space numerically stable version.

Appreciation

Adam Sorenti for contributing minLSTM!

Install

$ pip install minGRU-pytorch

Usage

import torch
from minGRU_pytorch import minGRU

min_gru = minGRU(512)

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

out = min_gru(x)

assert x.shape == out.shape

Sanity check

import torch
from minGRU_pytorch import minGRU

min_gru = minGRU(dim = 512, expansion_factor = 1.5)

x = torch.randn(1, 2048, 512)

# parallel

parallel_out = min_gru(x)[:, -1:]

# sequential

prev_hidden = None
for token in x.unbind(dim = 1):
    sequential_out, prev_hidden = min_gru(token[:, None, :], prev_hidden, return_next_prev_hidden = True)

assert torch.allclose(parallel_out, sequential_out, atol = 1e-4)

Test

enwik8

$ python train.py

Citations

@inproceedings{Feng2024WereRA,
    title   = {Were RNNs All We Needed?},
    author  = {Leo Feng and Frederick Tung and Mohamed Osama Ahmed and Yoshua Bengio and Hossein Hajimirsadegh},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:273025630}
}

@inproceedings{anonymous2024hymba,
    title   = {Hymba: A Hybrid-head Architecture for Small Language Models},
    author  = {Anonymous},
    booktitle = {Submitted to The Thirteenth International Conference on Learning Representations},
    year    = {2024},
    url     = {https://openreview.net/forum?id=A1ztozypga},
    note    = {under review}
}

Metadata

Release files for minGRU-pytorch 0.2.1

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

Source distribution (sdist)

Source distribution for minGRU-pytorch 0.2.1
File	Size	Uploaded
mingru_pytorch-0.2.1.tar.gz	36.6 MB	Mar 13, 2025	Details

Built distribution (wheel)

Table of built distributions (wheels) for minGRU-pytorch 0.2.1
File	Interpreter	ABI	Platform	Reset
mingru_pytorch-0.2.1-py3-none-any.whl 8.8 kB Mar 13, 2025	Python 3	none	any	Details

Total release size: 36.6 MB

Release files / mingru_pytorch-0.2.1.tar.gz

Download URL	mingru_pytorch-0.2.1.tar.gz
Size	36.6 MB
Tags	Source
SHA-256 checksum How to use checksums	`fc9b278b4e04db930a5cd4cdab36e2713b1d6e14084269b161d4f133819e758d`
BLAKE2b-256 checksum How to use checksums	`71b5ef07d75b1e7a33175c372332d3885bc4411908e0a9785325f92599005ac0`
Upload date	Mar 13, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.21`

Release files / mingru_pytorch-0.2.1-py3-none-any.whl

Download URL	mingru_pytorch-0.2.1-py3-none-any.whl
Size	8.8 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`22e83b7e9437f4b1a04c21b7e47686fd6ef2fa60bc961cd19e295a592e5e58d7`
BLAKE2b-256 checksum How to use checksums	`60e5d2992e2813314add3a0d368d8b57904007a31917a5455005d1d7b8714c6d`
Upload date	Mar 13, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.21`