HRM-pytorch

Hierarchical Reasoning Model (wip)

Explorations into the proposed recurrent hierarchical reasoning model by Wang et al. from Sapient Intelligence. Official repository is here

Install

$ pip install HRM-pytorch

Usage

import torch
from HRM import HRM

hrm = HRM(
    networks = [
        dict(
            dim = 32,
            depth = 2,
            attn_dim_head = 8,
            heads = 1,
            use_rmsnorm = True,
            rotary_pos_emb = True,
            pre_norm = False
        ),
        dict(
            dim = 32,
            depth = 4,
            attn_dim_head = 8,
            heads = 1,
            use_rmsnorm = True,
            rotary_pos_emb = True,
            pre_norm = False
        )
    ],
    num_tokens = 256,
    dim = 32,
    reasoning_steps = 10
)

seq = torch.randint(0, 256, (3, 1024))
labels = torch.randint(0, 256, (3, 1024))

loss, hiddens, _ = hrm(seq, labels = labels)
loss.backward()

loss, hiddens, _ = hrm(seq, hiddens = hiddens, labels = labels)
loss.backward()

# after much training

pred = hrm(seq, reasoning_steps = 5)

Citations

@misc{wang2025hierarchicalreasoningmodel,
    title   = {Hierarchical Reasoning Model},
    author  = {Guan Wang and Jin Li and Yuhao Sun and Xing Chen and Changling Liu and Yue Wu and Meng Lu and Sen Song and Yasin Abbasi Yadkori},
    year    = {2025},
    eprint  = {2506.21734},
    archivePrefix = {arXiv},
    primaryClass = {cs.AI},
    url     = {https://arxiv.org/abs/2506.21734},
}

Metadata

Release files for HRM-pytorch 0.1.6

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

Source distribution (sdist)

Source distribution for HRM-pytorch 0.1.6
File	Size	Uploaded
hrm_pytorch-0.1.6.tar.gz	194.5 kB	Aug 4, 2025	Details

Built distribution (wheel)

Table of built distributions (wheels) for HRM-pytorch 0.1.6
File	Interpreter	ABI	Platform	Reset
hrm_pytorch-0.1.6-py3-none-any.whl 9.9 kB Aug 4, 2025	Python 3	none	any	Details

Total release size: 204.4 kB

Release files / hrm_pytorch-0.1.6.tar.gz

Download URL	hrm_pytorch-0.1.6.tar.gz
Size	194.5 kB
Tags	Source
SHA-256 checksum How to use checksums	`545f1954c02521bb7e11f96f79a787213637c4e241445440e24244625ceb6c4a`
BLAKE2b-256 checksum How to use checksums	`b9a53c0471b2e55e82cb3dc3e02c581329fb6e986e5c63c52f55c550676ac6fe`
Upload date	Aug 4, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.23`

Release files / hrm_pytorch-0.1.6-py3-none-any.whl

Download URL	hrm_pytorch-0.1.6-py3-none-any.whl
Size	9.9 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`b9aaf484cb0fe5aad08f62886c08502a9225bb483eaf6e34d06244e0f16b84c1`
BLAKE2b-256 checksum How to use checksums	`d9614d603837284873640eb67bde62705afab1906b53b7d17eaf62adc0894a5c`
Upload date	Aug 4, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.23`