Refactor project metadata, add tests, add type-hints, lint/format

.github/workflows/python-publish.yml

@@ -14,23 +14,22 @@ on:
 
 jobs:
   deploy:
-
     runs-on: ubuntu-latest
 
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.x'
-    - name: Install dependencies
-      run: |
-        python -m pip install --upgrade pip
-        pip install build
-    - name: Build package
-      run: python -m build
-    - name: Publish package
-      uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
-      with:
-        user: __token__
-        password: ${{ secrets.PYPI_API_TOKEN }}
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v2
+        with:
+          python-version: "3.x"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install build
+      - name: Build package
+        run: python -m build
+      - name: Publish package
+        uses: pypa/gh-action-pypi-publish@27b31702a0e7fc50959f5ad993c78deac1bdfc29
+        with:
+          user: __token__
+          password: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/tests.yml

@@ -0,0 +1,40 @@
+name: Tests
+
+on:
+  push:
+    branches: [main, master]
+  pull_request:
+    branches: [main, master]
+  workflow_dispatch:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v4
+        with:
+          version: "latest"
+
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          uv sync --dev
+
+      - name: Run tests with pytest
+        run: |
+          uv run pytest --cov=rotary_embedding_torch --cov-report=xml --cov-report=term-missing
+
+      - name: Upload coverage to Codecov
+        uses: codecov/codecov-action@v4
+        with:
+          file: ./coverage.xml
+          fail_ci_if_error: false

.gitignore

@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# IDE stuff
+.idea/

README.md

@@ -8,10 +8,29 @@ My gut also tells me there is something <a href="https://www.nature.com/articles
 
 ## Install
 
+### Using pip
+
 ```bash
 $ pip install rotary-embedding-torch
 ```
 
+### Using uv (recommended for development)
+
+```bash
+$ uv add rotary-embedding-torch
+```
+
+### Development
+
+For development, clone the repository and install with test dependencies:
+
+```bash
+$ git clone https://github.com/lucidrains/rotary-embedding-torch
+$ cd rotary-embedding-torch
+$ uv sync  # create and install env
+$ uv run pytest  # run tests
+```
+
 ## Usage
 
 ```python
@@ -142,7 +161,7 @@ rotary_emb = RotaryEmbedding(
 
 ```bibtex
 @misc{su2021roformer,
-    title   = {RoFormer: Enhanced Transformer with Rotary Position Embedding}, 
+    title   = {RoFormer: Enhanced Transformer with Rotary Position Embedding},
     author  = {Jianlin Su and Yu Lu and Shengfeng Pan and Bo Wen and Yunfeng Liu},
     year    = {2021},
     eprint  = {2104.09864},

pyproject.toml

@@ -0,0 +1,93 @@
+[project]
+name = "rotary-embedding-torch"
+version = "0.8.8"
+description = "Rotary Embedding - Pytorch"
+readme = "README.md"
+license = {text = "MIT"}
+authors = [
+    {name = "Phil Wang", email = "[email protected]"}
+]
+keywords = [
+    "artificial intelligence",
+    "deep learning",
+    "positional embedding"
+]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Developers",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+]
+requires-python = ">=3.9"
+dependencies = [
+    "einops>=0.7",
+    "torch>=2.0",
+]
+
+[project.urls]
+Homepage = "https://github.com/lucidrains/rotary-embedding-torch"
+Repository = "https://github.com/lucidrains/rotary-embedding-torch"
+
+[dependency-groups]
+dev = [
+    "pytest>=8.4.1",
+    "pytest-cov>=6.2.1",
+    "coverage>=7.10.0",
+    "ruff>=0.12.5",
+]
+
+[tool.ruff]
+target-version = "py39"
+line-length = 120
+
+[tool.ruff.lint]
+select = [
+    "E",   # pycodestyle
+    "F",   # pyflakes
+    "I",   # isort
+    "UP",  # pyupgrade
+    "N",   # pep8-naming
+    "ANN", # annotations
+    "S",   # bandit
+    "A",   # builtins
+    "C4",  # comprehensions
+    "EM",  # errmsg
+    "ICN", # import-conventions
+    "PIE", # pie
+    "PT",  # pytest-style
+    "RSE", # raise
+    "RET", # return
+    "TID", # tidy-imports
+    "ARG", # unused-arguments
+    "TD",  # todos
+    "PL",  # pylint
+    "FLY", # flynt
+    "FA",  # flake8-future-annotations
+]
+
+[tool.pydoclint]
+style = "google"
+arg-type-hints-in-docstring = false
+check-return-types = false
+check-yield-types = false
+
+[tool.coverage.run]
+source = ["rotary_embedding_torch"]
+omit = ["tests/*"]
+
+[tool.coverage.report]
+exclude_lines = [
+    "pragma: no cover",
+    "def __repr__",
+    "raise AssertionError",
+    "raise NotImplementedError"
+]
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"

rotary_embedding_torch/__init__.py

@@ -1,6 +1,8 @@
 from rotary_embedding_torch.rotary_embedding_torch import (
-    apply_rotary_emb,
     RotaryEmbedding,
     apply_learned_rotations,
-    broadcat
+    apply_rotary_emb,
+    broadcast,
 )
+
+__all__ = ["RotaryEmbedding", "apply_learned_rotations", "apply_rotary_emb", "broadcast"]