Added workflows.

Author: patrick-kidger

.github/workflows/release.yml

@@ -0,0 +1,23 @@
+name: Release
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Release
+        uses: patrick-kidger/action_update_python_project@v7
+        with:
+            python-version: "3.11"
+            test-script: |
+                cp -r ${{ github.workspace }}/tests ./tests
+                cp ${{ github.workspace }}/pyproject.toml ./pyproject.toml
+                uv sync --extra dev --no-install-project --inexact
+                uv run --no-sync pytest
+            pypi-token: ${{ secrets.pypi_token }}
+            github-user: patrick-kidger
+            github-token: ${{ github.token }}

.github/workflows/run_tests.yml

@@ -0,0 +1,33 @@
+name: Run tests
+
+on:
+  pull_request:
+
+jobs:
+  run-test:
+    strategy:
+      matrix:
+        python-version: [ "3.10", "3.12" ]
+        os: [ ubuntu-latest ]
+      fail-fast: false
+    runs-on: ${{ matrix.os }}
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v2
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          python -m pip install '.[dev]'
+
+      - name: Checks with pre-commit
+        uses: pre-commit/[email protected]
+
+      - name: Test with pytest
+        run: |
+          python -m pytest

esm2quinox/_esm2.py

@@ -11,7 +11,7 @@
 
 
 @jax.jit
-def _randint(key: PRNGKeyArray, index: int, maxval: Int[ArrayLike, ""]):
+def _randint(key: PRNGKeyArray, index: Int[ArrayLike, ""], maxval: Int[ArrayLike, ""]):
     return jr.randint(jr.fold_in(key, index), (), minval=0, maxval=maxval)
 
 

pyproject.toml

@@ -1,37 +1,45 @@
+[build-system]
+build-backend = "hatchling.build"
+requires = ["hatchling"]
+
 [project]
-name = "esm2quinox"
-version = "0.0.1"
-description = "ESM2 implemented in Equinox+JAX."
-readme = "README.md"
-requires-python =">=3.10"
-license = {file = "LICENSE"}
 authors = [
-  {name = "Patrick Kidger", email = "[email protected]"},
+  {email = "[email protected]", name = "Patrick Kidger"}
 ]
-keywords = ["jax", "neural-networks", "deep-learning", "equinox", "esm2", "protein-language-model"]
 classifiers = [
-    "Development Status :: 3 - Alpha",
-    "Intended Audience :: Developers",
-    "Intended Audience :: Financial and Insurance Industry",
-    "Intended Audience :: Information Technology",
-    "Intended Audience :: Science/Research",
-    "License :: OSI Approved :: Apache Software License",
-    "Natural Language :: English",
-    "Programming Language :: Python :: 3",
-    "Topic :: Scientific/Engineering :: Artificial Intelligence",
-    "Topic :: Scientific/Engineering :: Information Analysis",
-    "Topic :: Scientific/Engineering :: Mathematics",
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Financial and Insurance Industry",
+  "Intended Audience :: Information Technology",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: Apache Software License",
+  "Natural Language :: English",
+  "Programming Language :: Python :: 3",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+  "Topic :: Scientific/Engineering :: Information Analysis",
+  "Topic :: Scientific/Engineering :: Mathematics"
 ]
-urls = {repository = "https://github.com/patrick-kidger/esm2quinox" }
 dependencies = ["equinox>=0.11.11", "jax>=0.5.0", "jaxtyping>=0.2.37"]
+description = "ESM2 implemented in Equinox+JAX."
+keywords = ["jax", "neural-networks", "deep-learning", "equinox", "esm2", "protein-language-model"]
+license = {file = "LICENSE"}
+name = "esm2quinox"
+readme = "README.md"
+requires-python = ">=3.10"
+urls = {repository = "https://github.com/patrick-kidger/esm2quinox"}
+version = "0.1.0"
 
-[build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
+[project.optional-dependencies]
+dev = ["beartype", "fair-esm==2.0.0", "pre-commit", "pytest", "torch"]
 
 [tool.hatch.build]
 include = ["esm2quinox/*"]
 
+[tool.pyright]
+include = ["esm2quinox", "tests"]
+reportIncompatibleMethodOverride = true
+reportIncompatibleVariableOverride = false # Incompatible with eqx.AbstractVar
+
 [tool.pytest.ini_options]
 addopts = "--jaxtyping-packages=esm2quinox,beartype.beartype(conf=beartype.BeartypeConf(strategy=beartype.BeartypeStrategy.On))"
 
@@ -50,8 +58,3 @@ combine-as-imports = true
 extra-standard-library = ["typing_extensions"]
 lines-after-imports = 2
 order-by-type = false
-
-[tool.pyright]
-reportIncompatibleMethodOverride = true
-reportIncompatibleVariableOverride = false  # Incompatible with eqx.AbstractVar
-include = ["esm2quinox", "tests"]