├── .github └── workflows │ ├── pre-commit.yml │ ├── pypi.yml │ ├── static.yml │ ├── testpypi.yml │ └── unit.yml ├── .gitignore ├── .pre-commit-config.yaml ├── CITATION.cff ├── LICENSE ├── README.md ├── docs ├── Makefile ├── README.md ├── make.bat └── source │ ├── conf.py │ ├── e3tools.nn.rst │ ├── e3tools.rst │ ├── index.rst │ └── modules.rst ├── examples ├── models │ ├── __init__.py │ ├── conv.py │ └── transformer.py ├── requirements.txt └── train_qm9.py ├── pyproject.toml ├── src └── e3tools │ ├── __init__.py │ ├── _pack_unpack.py │ ├── _radius.py │ ├── _scatter.py │ ├── nn │ ├── __init__.py │ ├── _axis_to_mul.py │ ├── _conv.py │ ├── _extract_irreps.py │ ├── _gate.py │ ├── _interaction.py │ ├── _layer_norm.py │ ├── _linear.py │ ├── _mlp.py │ ├── _repeat.py │ ├── _scaling.py │ ├── _tensor_product.py │ └── _transformer.py │ └── utils │ ├── __init__.py │ └── _default_dtype_manager.py ├── tests ├── test_attention.py ├── test_equivariance.py ├── test_extract_irreps.py ├── test_fused_conv.py ├── test_layer_norm.py ├── test_pack_unpack.py ├── test_radius.py ├── test_scaling.py └── test_scatter.py └── uv.lock /.github/workflows/pre-commit.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/.github/workflows/pre-commit.yml -------------------------------------------------------------------------------- /.github/workflows/pypi.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/.github/workflows/pypi.yml -------------------------------------------------------------------------------- /.github/workflows/static.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/.github/workflows/static.yml -------------------------------------------------------------------------------- /.github/workflows/testpypi.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/.github/workflows/testpypi.yml -------------------------------------------------------------------------------- /.github/workflows/unit.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/.github/workflows/unit.yml -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | __pycache__ 2 | .venv 3 | docs/build 4 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/.pre-commit-config.yaml -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/CITATION.cff -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/README.md -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/docs/README.md -------------------------------------------------------------------------------- /docs/make.bat: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/docs/make.bat -------------------------------------------------------------------------------- /docs/source/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/docs/source/conf.py -------------------------------------------------------------------------------- /docs/source/e3tools.nn.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/docs/source/e3tools.nn.rst -------------------------------------------------------------------------------- /docs/source/e3tools.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/docs/source/e3tools.rst -------------------------------------------------------------------------------- /docs/source/index.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/docs/source/index.rst -------------------------------------------------------------------------------- /docs/source/modules.rst: -------------------------------------------------------------------------------- 1 | src 2 | === 3 | 4 | .. toctree:: 5 | :maxdepth: 4 6 | 7 | e3tools 8 | -------------------------------------------------------------------------------- /examples/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/examples/models/__init__.py -------------------------------------------------------------------------------- /examples/models/conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/examples/models/conv.py -------------------------------------------------------------------------------- /examples/models/transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/examples/models/transformer.py -------------------------------------------------------------------------------- /examples/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/examples/requirements.txt -------------------------------------------------------------------------------- /examples/train_qm9.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/examples/train_qm9.py -------------------------------------------------------------------------------- /pyproject.toml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/pyproject.toml -------------------------------------------------------------------------------- /src/e3tools/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/__init__.py -------------------------------------------------------------------------------- /src/e3tools/_pack_unpack.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/_pack_unpack.py -------------------------------------------------------------------------------- /src/e3tools/_radius.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/_radius.py -------------------------------------------------------------------------------- /src/e3tools/_scatter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/_scatter.py -------------------------------------------------------------------------------- /src/e3tools/nn/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/__init__.py -------------------------------------------------------------------------------- /src/e3tools/nn/_axis_to_mul.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_axis_to_mul.py -------------------------------------------------------------------------------- /src/e3tools/nn/_conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_conv.py -------------------------------------------------------------------------------- /src/e3tools/nn/_extract_irreps.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_extract_irreps.py -------------------------------------------------------------------------------- /src/e3tools/nn/_gate.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_gate.py -------------------------------------------------------------------------------- /src/e3tools/nn/_interaction.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_interaction.py -------------------------------------------------------------------------------- /src/e3tools/nn/_layer_norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_layer_norm.py -------------------------------------------------------------------------------- /src/e3tools/nn/_linear.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_linear.py -------------------------------------------------------------------------------- /src/e3tools/nn/_mlp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_mlp.py -------------------------------------------------------------------------------- /src/e3tools/nn/_repeat.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_repeat.py -------------------------------------------------------------------------------- /src/e3tools/nn/_scaling.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_scaling.py -------------------------------------------------------------------------------- /src/e3tools/nn/_tensor_product.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_tensor_product.py -------------------------------------------------------------------------------- /src/e3tools/nn/_transformer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/nn/_transformer.py -------------------------------------------------------------------------------- /src/e3tools/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/utils/__init__.py -------------------------------------------------------------------------------- /src/e3tools/utils/_default_dtype_manager.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/src/e3tools/utils/_default_dtype_manager.py -------------------------------------------------------------------------------- /tests/test_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_attention.py -------------------------------------------------------------------------------- /tests/test_equivariance.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_equivariance.py -------------------------------------------------------------------------------- /tests/test_extract_irreps.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_extract_irreps.py -------------------------------------------------------------------------------- /tests/test_fused_conv.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_fused_conv.py -------------------------------------------------------------------------------- /tests/test_layer_norm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_layer_norm.py -------------------------------------------------------------------------------- /tests/test_pack_unpack.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_pack_unpack.py -------------------------------------------------------------------------------- /tests/test_radius.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_radius.py -------------------------------------------------------------------------------- /tests/test_scaling.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_scaling.py -------------------------------------------------------------------------------- /tests/test_scatter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/tests/test_scatter.py -------------------------------------------------------------------------------- /uv.lock: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/prescient-design/e3tools/HEAD/uv.lock --------------------------------------------------------------------------------