{"id":38781308,"url":"https://github.com/atomicarchitects/e3nn_tetris","last_synced_at":"2026-01-17T12:25:06.672Z","repository":{"id":250229759,"uuid":"833386919","full_name":"atomicarchitects/e3nn_tetris","owner":"atomicarchitects","description":"Simple Tetris Network in Torch 2.0 + C++ inference","archived":false,"fork":false,"pushed_at":"2024-11-25T00:43:04.000Z","size":270,"stargazers_count":4,"open_issues_count":1,"forks_count":1,"subscribers_count":2,"default_branch":"main","last_synced_at":"2025-09-09T12:44:44.570Z","etag":null,"topics":[],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":null,"status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/atomicarchitects.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":null,"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,"notice":null,"maintainers":null,"copyright":null,"agents":null,"dco":null,"cla":null}},"created_at":"2024-07-25T00:22:06.000Z","updated_at":"2024-12-15T04:21:05.000Z","dependencies_parsed_at":"2024-08-14T08:41:37.422Z","dependency_job_id":"b1ae54db-c775-44ba-b64d-d2241029d427","html_url":"https://github.com/atomicarchitects/e3nn_tetris","commit_stats":null,"previous_names":["atomicarchitects/e3nn_simple_network","atomicarchitects/e3nn_tetris"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/atomicarchitects/e3nn_tetris","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fe3nn_tetris","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fe3nn_tetris/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fe3nn_tetris/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fe3nn_tetris/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/atomicarchitects","download_url":"https://codeload.github.com/atomicarchitects/e3nn_tetris/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/atomicarchitects%2Fe3nn_tetris/sbom","scorecard":null,"host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":286080680,"owners_count":28508464,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2026-01-17T11:50:55.898Z","status":"ssl_error","status_checked_at":"2026-01-17T11:50:55.569Z","response_time":85,"last_error":"SSL_connect returned=1 errno=0 peeraddr=140.82.121.6:443 state=error: unexpected eof while reading","robots_txt_status":"success","robots_txt_updated_at":"2025-07-24T06:49:26.215Z","robots_txt_url":"https://github.com/robots.txt","online":false,"can_crawl_api":true,"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":"2026-01-17T12:25:06.077Z","updated_at":"2026-01-17T12:25:06.662Z","avatar_url":"https://github.com/atomicarchitects.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"## e3nn_tetris\n\nLets train a simple Nequip Tetris model with e3nn + torch.compile and run inference in C++ !\n\n\n### Train\n\n```python\n(.venv) mkotak@radish:~/atomic_architects/projects/e3nn_tetris$ python train.py \ncompiling...\nW0815 21:09:54.341000 140232380905024 torch/fx/experimental/symbolic_shapes.py:4449] [0/0] xindex is not in var_ranges, defaulting to unknown range.\ninitial accuracy = 12%\ncompilation took 20.7s\ntraining...\n 34%|██████████████████████████████████▋                                                                   | 68/200 [00:00\u003c00:00, 651.29it/s]\nfinal accuracy = 100%\ntraining took 0.1s\nW0815 21:10:06.125000 140242116276224 torch/fx/experimental/symbolic_shapes.py:4449] rindex is not in var_ranges, defaulting to unknown range.\nnode_features tensor([[1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.],\n        [1.]], device='cuda:0')\npos tensor([[0., 0., 0.],\n        [0., 0., 1.],\n        [1., 0., 0.],\n        [1., 1., 0.],\n        [1., 1., 1.],\n        [1., 1., 2.],\n        [2., 1., 1.],\n        [2., 0., 1.],\n        [0., 0., 0.],\n        [1., 0., 0.],\n        [0., 1., 0.],\n        [1., 1., 0.],\n        [0., 0., 0.],\n        [0., 0., 1.],\n        [0., 0., 2.],\n        [0., 0., 3.],\n        [0., 0., 0.],\n        [0., 0., 1.],\n        [0., 1., 0.],\n        [1., 0., 0.],\n        [0., 0., 0.],\n        [0., 0., 1.],\n        [0., 0., 2.],\n        [0., 1., 0.],\n        [0., 0., 0.],\n        [0., 0., 1.],\n        [0., 0., 2.],\n        [0., 1., 1.],\n        [0., 0., 0.],\n        [1., 0., 0.],\n        [1., 1., 0.],\n        [2., 1., 0.]], device='cuda:0')\nedge_index tensor([[ 1,  2,  0,  0,  3,  2,  5,  6,  4,  4,  7,  6,  9, 10,  8, 11,  8, 11,\n          9, 10, 13, 12, 14, 13, 15, 14, 17, 18, 19, 16, 16, 16, 21, 23, 20, 22,\n         21, 20, 25, 24, 26, 27, 25, 25, 29, 28, 30, 29, 31, 30],\n        [ 0,  0,  1,  2,  2,  3,  4,  4,  5,  6,  6,  7,  8,  8,  9,  9, 10, 10,\n         11, 11, 12, 13, 13, 14, 14, 15, 16, 16, 16, 17, 18, 19, 20, 20, 21, 21,\n         22, 23, 24, 25, 25, 25, 26, 27, 28, 29, 29, 30, 30, 31]],\n       device='cuda:0')\nbatch tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 3, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5,\n        6, 6, 6, 6, 7, 7, 7, 7], device='cuda:0')\noutput tensor([[ 0.0344, -0.0344, -0.1076, -0.4316, -0.1765, -0.1971, -0.2294, -0.3128],\n        [-0.0559,  0.0559,  0.0109, -0.1000, -0.3582, -0.6548, -0.4990, -0.4543],\n        [-0.1819,  0.1819,  0.8137, -0.0546, -0.5677, -0.0819, -0.5391,  0.4103],\n        [ 0.0096, -0.0096,  0.3495,  0.4046, -0.3325, -0.0950, -0.0316, -0.3742],\n        [ 0.0054, -0.0054, -0.1696, -0.5817,  0.6609,  0.3416,  0.0284, -0.6731],\n        [ 0.0231, -0.0231,  0.1962, -0.0599, -0.0899,  0.4756,  0.2846, -0.3791],\n        [ 0.0878, -0.0878,  0.2029, -0.0223, -0.0784,  0.3471,  0.5549, -0.4884],\n        [-0.1149,  0.1149,  0.4429, -0.1543, -0.4349, -0.0615, -0.5707,  0.7809]],\n       device='cuda:0')\n```\n\n\n\n### Inference\n\n\n```python\n(.venv) mkotak@radish:~/atomic_architects/projects/e3nn_tetris$ make run\n./build/inference /home/mkotak/atomic_architects/projects/e3nn_tetris/export/model.so\noutput tensor 0.0344 -0.0344 -0.1076 -0.4316 -0.1765 -0.1971 -0.2294 -0.3128\n-0.0559  0.0559  0.0109 -0.1000 -0.3582 -0.6548 -0.4990 -0.4543\n-0.1819  0.1819  0.8137 -0.0546 -0.5677 -0.0819 -0.5391  0.4103\n 0.0096 -0.0096  0.3495  0.4046 -0.3325 -0.0950 -0.0316 -0.3742\n 0.0054 -0.0054 -0.1696 -0.5817  0.6609  0.3416  0.0284 -0.6731\n 0.0231 -0.0231  0.1962 -0.0599 -0.0899  0.4756  0.2846 -0.3791\n 0.0878 -0.0878  0.2029 -0.0223 -0.0784  0.3471  0.5549 -0.4884\n-0.1149  0.1149  0.4429 -0.1543 -0.4349 -0.0615 -0.5707  0.7809\n[ CUDAFloatType{8,8} ]\n```\n\n\n\n\n### (Preliminary) Training time comparision with e3nn-jax on RTX A5500\n\nTODO: Need to test out on bigger datasets/models. The throughput is better in Torch2 currently is better but there might be some initialization differences that make JAX converge faster.\n\n- e3nn + Torch 2\n\n```python\ncompiling...\ninitial accuracy = 12%\ncompilation took 127.5s\ntraining...\n 34%|███████████████████████████████████ | 68/200 [00:00\u003c00:00, 797.43it/s]\nfinal accuracy = 100%\ntraining took 0.1s\n```\n\n- e3nn-jax\n\n```python\ncompiling...\ninitial accuracy = 12%\ncompilation took 5.3s\ntraining...\n 12%|███████████▋ | 24/200 [00:00\u003c00:00, 635.90it/s]\nfinal accuracy = 100%\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fatomicarchitects%2Fe3nn_tetris","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fatomicarchitects%2Fe3nn_tetris","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fatomicarchitects%2Fe3nn_tetris/lists"}