diff --git a/.flake8 b/.flake8
new file mode 100644
index 0000000..f308054
--- /dev/null
+++ b/.flake8
@@ -0,0 +1,5 @@
+[flake8]
+max-line-length = 88
+select = E, F
+ignore = E402, E721, E731, E741, F722
+per-file-ignores = **/__init__.py:F401
diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml
index 011e33d..b994e4f 100644
--- a/.github/workflows/lint.yml
+++ b/.github/workflows/lint.yml
@@ -18,7 +18,7 @@ jobs:
       max-parallel: 4
       matrix:
         os: [ubuntu-latest]
-        python-version: ["3.10"]
+        python-version: ["3.12"]
     steps:
       - uses: actions/checkout@v3
       - name: Set up python ${{ matrix.python-version }}
@@ -28,6 +28,8 @@ jobs:
       - name: Run styles & lint check
         run: |
           pip install .[testing]
-          ruff format --check jumpax/ tests/ examples/
-          ruff check jumpax/ tests/ examples/
+          pip install "black[jupyter]" flake8 isort
+          black --check jumpax/ tests/ examples/
+          isort --check jumpax/ tests/ examples/
+          flake8 jumpax/ tests/ examples/
           pyright jumpax/
\ No newline at end of file
diff --git a/.gitignore b/.gitignore
index 0b9019d..5b837ee 100644
--- a/.gitignore
+++ b/.gitignore
@@ -206,3 +206,5 @@ cython_debug/
 marimo/_static/
 marimo/_lsp/
 __marimo__/
+
+.vscode
\ No newline at end of file
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..54458cb
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,22 @@
+repos:
+  - repo: https://github.com/PyCQA/isort
+    rev: 8.0.1
+    hooks:
+      - id: isort
+        types_or: [ python, pyi, jupyter ]
+  - repo: https://github.com/psf/black
+    rev: 26.3.1
+    hooks:
+      - id: black
+        types_or: [ python, pyi, jupyter ]
+        additional_dependencies: ["black[jupyter]"]
+  - repo: https://github.com/PyCQA/flake8
+    rev: 7.3.0
+    hooks:
+      - id: flake8
+        types_or: [ python, pyi, jupyter ]
+  - repo: https://github.com/RobertCraigie/pyright-python
+    rev: v1.1.408
+    hooks:
+    - id: pyright
+      additional_dependencies: ["equinox", "pytest", "jax", "jaxtyping"]
diff --git a/jumpax/_aggregators/base.py b/jumpax/_aggregators/base.py
index 9c2c229..955439b 100644
--- a/jumpax/_aggregators/base.py
+++ b/jumpax/_aggregators/base.py
@@ -7,7 +7,6 @@
 from .._custom_meta import AbstractStrictModule
 from .._custom_types import Args, Info, JumpState, RealScalarLike, SolverState, U
 
-
 _SolverState = TypeVar("_SolverState", bound=SolverState)
 
 
diff --git a/jumpax/_custom_meta.py b/jumpax/_custom_meta.py
index 3cc6892..a778f56 100644
--- a/jumpax/_custom_meta.py
+++ b/jumpax/_custom_meta.py
@@ -4,7 +4,6 @@
 import equinox as eqx
 import ihoop
 
-
 try:
 
     class _StrictEqxMeta(ihoop.strict._StrictMeta, eqx._module._module._ModuleMeta):
diff --git a/jumpax/_custom_types.py b/jumpax/_custom_types.py
index fcbf717..b3193ba 100644
--- a/jumpax/_custom_types.py
+++ b/jumpax/_custom_types.py
@@ -9,7 +9,6 @@
     Shaped,
 )
 
-
 if TYPE_CHECKING:
     RealScalarLike = bool | int | float | Array | np.ndarray
 else:
diff --git a/jumpax/_jumps.py b/jumpax/_jumps.py
index f4938b0..a0ae3c1 100644
--- a/jumpax/_jumps.py
+++ b/jumpax/_jumps.py
@@ -10,7 +10,6 @@
 from ._custom_meta import AbstractStrictModule
 from ._custom_types import Args, JumpState, Rate, RealScalarLike, U
 
-
 _Rate = TypeVar("_Rate", bound=Rate)
 _JumpState = TypeVar("_JumpState", bound=JumpState)
 
@@ -197,10 +196,10 @@ def leap_delta(self, t: RealScalarLike, u: U, args: Args) -> Int[Array, "1 S"]:
 
 
 def _massaction_comb(
-    n: Float[Array, "*shape"],
-    k: Int[Array, "*shape"],
-    lgamma_k_plus_1: Float[Array, "*shape"],
-) -> Float[Array, "*shape"]:
+    n: Float[Array, " *shape"],
+    k: Int[Array, " *shape"],
+    lgamma_k_plus_1: Float[Array, " *shape"],
+) -> Float[Array, " *shape"]:
     r"""
     Compute the binomial coefficient $\binom{n}{k}$ using precomputed log-factorials.
 
diff --git a/pyproject.toml b/pyproject.toml
index d4fe126..9d9556b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -39,8 +39,7 @@ examples = [
 testing = [
     "pytest>=7.4.0",
     "nbmake>=1.5.0",
-    "ruff==0.11.0",
-    "pyright==1.1.399",
+    "pyright==1.1.408",
 ]
 docs = [
     "mkdocs>=1.5.0",
@@ -54,31 +53,14 @@ docs = [
 [tool.setuptools]
 packages = { find = { include = ["jumpax", "jumpax.*"] } }
 
-[tool.ruff]
-extend-include = ["*.ipynb"]
-src = []
+[tool.black]
+target-version = ["py312"]
 
-[tool.ruff.lint]
-# I001: isort import sorting
-# F401: Unused imports
-fixable = ["I001", "F401"]
-
-# E402: Module level import not at top of file
-# E731: Do not assign a lambda expression, use a def
-# E741: Ambiguous variable name
-# F722: Syntax error in forward annotation
-ignore = ["E402", "E731", "E741", "F722"]
-
-# Error codes to check for during linting
-# E: pycodestyle errors
-# F: pyflakes errors
-# I001: isort import sorting
-select = ["E", "F", "I001"]
-
-[tool.ruff.lint.isort]
-combine-as-imports = true
-lines-after-imports = 2
-order-by-type = false
+[tool.isort]
+profile = "black"
+combine_as_imports = true
+lines_after_imports = -1
+order_by_type = false
 
 [tool.pyright]
 reportIncompatibleMethodOverride = true
diff --git a/tests/test_cme.py b/tests/test_cme.py
index 3bed617..fabdebc 100644
--- a/tests/test_cme.py
+++ b/tests/test_cme.py
@@ -1,9 +1,10 @@
 import unittest
 
 import jax.numpy as jnp
-import jumpax as jx
 from jax.scipy.special import gammaln
 
+import jumpax as jx
+
 
 def poisson_pmf(k, lam):
     """Poisson probability mass function."""
diff --git a/tests/test_readme.py b/tests/test_readme.py
index c9ca310..e578103 100644
--- a/tests/test_readme.py
+++ b/tests/test_readme.py
@@ -4,6 +4,7 @@
 
 import jax.numpy as jnp
 import jax.random as jr
+
 import jumpax as jpx