gotennet-pytorch

GotenNet - Pytorch

Implementation of GotenNet, new SOTA 3d equivariant transformer, in Pytorch

I know a lot of researchers have moved on from geometric learning after Alphafold3. However, I just cannot help but wonder. Hedging my bets

The official repository has been released here

Install

$ pip install gotennet-pytorch

Usage

import torch
torch.set_default_dtype(torch.float64) # recommended for equivariant network training

from gotennet_pytorch import GotenNet

model = GotenNet(
    dim = 256,
    max_degree = 2,
    depth = 1,
    heads = 2,
    dim_head = 32,
    dim_edge_refinement = 256,
    return_coors = False
)

atom_ids = torch.randint(0, 14, (1, 12)) # negative atom indices will be assumed to be padding - length of molecule is thus `(atom_ids >= 0).sum(dim = -1)`
coors = torch.randn(1, 12, 3)
adj_mat = torch.randint(0, 2, (1, 12, 12)).bool()

invariant, coors_out = model(atom_ids, adj_mat = adj_mat, coors = coors)

Citations

@inproceedings{anonymous2024rethinking,
    title   = {Rethinking Efficient 3D Equivariant Graph Neural Networks},
    author  = {Anonymous},
    booktitle = {Submitted to The Thirteenth International Conference on Learning Representations},
    year    = {2024},
    url     = {https://openreview.net/forum?id=5wxCQDtbMo},
    note    = {under review}
}

@inproceedings{Zhou2024ValueRL,
    title   = {Value Residual Learning For Alleviating Attention Concentration In Transformers},
    author  = {Zhanchao Zhou and Tianyi Wu and Zhiyun Jiang and Zhenzhong Lan},
    year    = {2024},
    url     = {https://api.semanticscholar.org/CorpusID:273532030}
}

@article{Zhu2024HyperConnections,
    title   = {Hyper-Connections},
    author  = {Defa Zhu and Hongzhi Huang and Zihao Huang and Yutao Zeng and Yunyao Mao and Banggu Wu and Qiyang Min and Xun Zhou},
    journal = {ArXiv},
    year    = {2024},
    volume  = {abs/2409.19606},
    url     = {https://api.semanticscholar.org/CorpusID:272987528}
}

Metadata

Release files for gotennet-pytorch 0.3.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 gotennet-pytorch 0.3.1
File	Size	Uploaded
gotennet_pytorch-0.3.1.tar.gz	144.6 kB	Apr 7, 2025	Details

Built distribution (wheel)

Table of built distributions (wheels) for gotennet-pytorch 0.3.1
File	Interpreter	ABI	Platform	Reset
gotennet_pytorch-0.3.1-py3-none-any.whl 11.4 kB Apr 7, 2025	Python 3	none	any	Details

Total release size: 156.0 kB

Release files / gotennet_pytorch-0.3.1.tar.gz

Download URL	gotennet_pytorch-0.3.1.tar.gz
Size	144.6 kB
Tags	Source
SHA-256 checksum How to use checksums	`e1ed40b1acd751beebcdcfe793aae21dae09e325b381a6e72799178eee2ea456`
BLAKE2b-256 checksum How to use checksums	`dfb49e01fbe97d0a01c87e9ba491cdea7e9214d9f9448406aa0dbc59b25a10de`
Upload date	Apr 7, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.21`

Release files / gotennet_pytorch-0.3.1-py3-none-any.whl

Download URL	gotennet_pytorch-0.3.1-py3-none-any.whl
Size	11.4 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`bb449eeb2ee4f2360423278ca874f110612570ef97cc60768ed8e9b589a73edd`
BLAKE2b-256 checksum How to use checksums	`0c372ebf8509250eeec4f7e7d2ad8490464f852e46599cb62f6db03537343625`
Upload date	Apr 7, 2025
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/6.1.0 CPython/3.9.21`