{"id":29023266,"url":"https://github.com/rish-16/torch-efa","last_synced_at":"2025-08-06T13:20:42.113Z","repository":{"id":300355002,"uuid":"1001186901","full_name":"rish-16/torch-efa","owner":"rish-16","description":"PyTorch implementation of Euclidean Fast Attention and Euclidean RoPE","archived":false,"fork":false,"pushed_at":"2025-06-21T07:09:07.000Z","size":761,"stargazers_count":1,"open_issues_count":0,"forks_count":0,"subscribers_count":0,"default_branch":"main","last_synced_at":"2025-06-21T08:21:03.519Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":null,"language":"Jupyter Notebook","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/rish-16.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null,"zenodo":null}},"created_at":"2025-06-13T01:03:23.000Z","updated_at":"2025-06-21T07:09:10.000Z","dependencies_parsed_at":"2025-06-21T08:21:10.507Z","dependency_job_id":"63d3e240-e045-4cfd-a336-e02cbd2e456c","html_url":"https://github.com/rish-16/torch-efa","commit_stats":null,"previous_names":["rish-16/torch-efa"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/rish-16/torch-efa","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rish-16%2Ftorch-efa","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rish-16%2Ftorch-efa/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rish-16%2Ftorch-efa/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rish-16%2Ftorch-efa/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/rish-16","download_url":"https://codeload.github.com/rish-16/torch-efa/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/rish-16%2Ftorch-efa/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":261990349,"owners_count":23241189,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":[],"created_at":"2025-06-26T03:05:56.714Z","updated_at":"2025-08-06T13:20:42.099Z","avatar_url":"https://github.com/rish-16.png","language":"Jupyter Notebook","funding_links":[],"categories":[],"sub_categories":[],"readme":"# torch-efa\nPyTorch implementation of Euclidean Fast Attention (EFA) and Euclidean RoPE (ERoPE)\n\n## Euclidean Fast Attention\n\n[EFA](https://arxiv.org/abs/2412.08541) is an $O(3)$-equivariant, linear self-attention mechanism for 3D geometric data like molecules. Using kernelized attention, EFA can achieve all-to-all message passing. Euclidean Rotary Positional Encodings (ERoPE) is used to encode orientational information between pairs of nodes directly into the node features. EFA uses higher-order equivariant features from spherical harmonics, making it highly expressive. \n\n![EFA Architecture](./assets/efa_banner.png)\n\n`torch-efa` is a fast and scalable PyTorch implementation of EFA and ERoPE for ML-based interatomic potentials or molecular systems.\n\n## Installation\n\nYou can install `torch-efa` via `uv` and `pip`:\n\n```bash\n$ uv pip install torch-efa\n```\n\nYou can also build the package from source:\n\n```\nbash\n$ git clone https://github.com/rish-16/torch-efa.git\n$ cd torch-efa\n$ uv sync\n```\n\n## Usage\n\nYou can use the `EuclideanFastAttention` module like any other `torch.nn.Module` in place of existing equivariant modules:\n\n```python\nimport torch\nfrom torch_efa import EuclideanFastAttention\n\nnode_ftrs = torch.randn(B, N_atoms, parity, max_deg, dim)\npos = torch.randn(B, N_atoms, 3)\n\nefa = EuclideanFastAttention(\n    in_dim=o3.Irreps(\"...\"), # input irreps\n    out_dim=o3.Irreps(\"...\"), # output irrep\n    max_degree=3, # highest tensor degree\n    num_lebedev=11, # number of grid points on sphere for Lebedev Quadrature \n)\n\ny = efa(node_ftrs, pos) # [B, N_atoms, out_dim]\n```\n\n## Contributing\n\nAppreciate any contributions in the form of PRs or Issues! \n\n## Citation\n\nIf `torch-efa` has been helpful in your work, please consider citing the repo!\n\n```\n@misc{anand2025torchefa,\n  author       = {Rishabh Anand},\n  title        = {torch-efa},\n  year         = {2025},\n  howpublished = {\\url{https://github.com/rish-16/torch-efa}},\n  publisher    = {GitHub}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Frish-16%2Ftorch-efa","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Frish-16%2Ftorch-efa","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Frish-16%2Ftorch-efa/lists"}