diff --git a/.coveragerc b/.coveragerc
deleted file mode 100644
index 796666d..0000000
--- a/.coveragerc
+++ /dev/null
@@ -1,8 +0,0 @@
-[run]
-relative_files = true
-
-[report]
-exclude_lines =
-  raise NotImplementedError
-  except ImportError
-  def backward
diff --git a/.github/workflows/cpp.yml b/.github/workflows/cpp.yml
index a277d9e..90500a4 100644
--- a/.github/workflows/cpp.yml
+++ b/.github/workflows/cpp.yml
@@ -1,10 +1,17 @@
 name: C++
-
 on:
   push:
-    branches: [main]
+    branches:
+      - main
+    tags:
+      - "*"
+    paths-ignore:
+      - "docs/**"
   pull_request:
-    branches: [main]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref }}
+  cancel-in-progress: true
 
 jobs:
   build:
diff --git a/.github/workflows/python.yml b/.github/workflows/python.yml
index c4c4f91..74bcfe4 100644
--- a/.github/workflows/python.yml
+++ b/.github/workflows/python.yml
@@ -1,114 +1,74 @@
 name: Python
 on:
   push:
-    branches: [main]
+    branches:
+      - main
+    tags:
+      - "*"
+    paths-ignore:
+      - "docs/**"
   pull_request:
-    branches: [main]
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref }}
+  cancel-in-progress: true
 
 jobs:
-  build:
-    runs-on: ${{ matrix.os }}
+  tests:
+    runs-on: ubuntu-latest
     strategy:
+      fail-fast: false
       matrix:
-        python-version: ["3.8", "3.9"]
-        os: [ubuntu-latest, windows-latest, macos-latest]
+        python-version: ["3.9"]
+        extras: ["core", "jax", "pymc3"]
+        include:
+          - python-version: "3.10"
+            extras: "core"
+
     steps:
-      - uses: actions/checkout@v2
+      - name: Checkout
+        uses: actions/checkout@v2
         with:
-          submodules: true
           fetch-depth: 0
-      - name: Set up Python ${{ matrix.python-version }}
+          submodules: true
+
+      - name: Setup Python
         uses: actions/setup-python@v2
         with:
           python-version: ${{ matrix.python-version }}
-      - uses: ilammy/msvc-dev-cmd@v1
-      - name: Install dependencies
-        run: |
-          python -m pip install -U pip
-          python -m pip install -e ".[test]"
-        env:
-          DISTUTILS_USE_SDK: 1
-          MSSdk: 1
-      - name: Run the unit tests
-        run: python -m pytest --cov celerite2 python/test
-      - name: Coveralls
-        if: startsWith(matrix.os, 'ubuntu')
-        uses: AndreMiras/coveralls-python-action@v20201129
-        with:
-          parallel: true
-          flag-name: Unit Tests
 
-  theano:
-    runs-on: "ubuntu-latest"
-    steps:
-      - uses: actions/checkout@v2
-        with:
-          submodules: true
-          fetch-depth: 0
-      - name: Set up Python
-        uses: conda-incubator/setup-miniconda@v2
-        with:
-          python-version: 3.8
-          auto-update-conda: true
       - name: Install dependencies
-        shell: bash -l {0}
         run: |
-          conda install -q numpy scipy theano mkl-service
           python -m pip install -U pip
-          python -m pip install --use-feature=2020-resolver -e ".[test,theano]"
-      - name: Get theano compiledir
-        id: compiledir
-        shell: bash -l {0}
-        run: |
-          python -c "import theano; print('::set-output name=compiledir::' + theano.config.compiledir.split('/')[-1])"
-      - name: "Cache ~/.theano"
-        uses: actions/cache@v2
-        with:
-          path: ~/.theano
-          key: theano-${{ steps.compiledir.outputs.compiledir }}-${{ hashFiles('python/test/theano/*.py') }}
-          restore-keys: |
-            theano-${{ steps.compiledir.outputs.compiledir }}-
-            theano-
-      - name: Run the unit tests
-        shell: bash -l {0}
-        run: python -m pytest --cov celerite2 python/test/theano
-      - name: Coveralls
-        uses: AndreMiras/coveralls-python-action@v20201129
-        with:
-          parallel: true
-          flag-name: Unit Tests
+          python -m pip install -U coveralls coverage[toml] tox tox-gh-actions pybind11
 
-  jax:
-    runs-on: "ubuntu-latest"
-    steps:
-      - uses: actions/checkout@v2
-        with:
-          submodules: true
-          fetch-depth: 0
-      - name: Set up Python
-        uses: conda-incubator/setup-miniconda@v2
-        with:
-          python-version: 3.8
-          auto-update-conda: true
-      - name: Install dependencies
-        shell: bash -l {0}
+      - name: Run tests
+        run: python -m tox
+        env:
+          EXTRAS: ${{ matrix.extras }}
+
+      - name: Combine and upload coverage
         run: |
-          python -m pip install -U pip
-          python -m pip install --use-feature=2020-resolver -e ".[test,jax]"
-      - name: Run the unit tests
-        shell: bash -l {0}
-        run: python -m pytest --cov celerite2 python/test/jax
-      - name: Coveralls
-        uses: AndreMiras/coveralls-python-action@v20201129
-        with:
-          parallel: true
-          flag-name: Unit Tests
+          python -m coverage combine
+          python -m coverage xml -i
+          python -m coveralls --service=github
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          COVERALLS_PARALLEL: true
+          COVERALLS_FLAG_NAME: py${{ matrix.python-version }}-${{ matrix.extras }}
 
   coverage:
-    needs: [build, theano, jax]
+    needs: tests
     runs-on: ubuntu-latest
     steps:
-      - name: Coveralls Finished
-        uses: AndreMiras/coveralls-python-action@v20201129
+      - name: Setup Python
+        uses: actions/setup-python@v2
         with:
-          parallel-finished: true
+          python-version: "3.9"
+      - name: Finish coverage collection
+        run: |
+          python -m pip install -U pip
+          python -m pip install -U coveralls
+          python -m coveralls --finish
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml
index 6342b08..bb25379 100644
--- a/.github/workflows/tutorials.yml
+++ b/.github/workflows/tutorials.yml
@@ -9,6 +9,10 @@ on:
     types:
       - published
 
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref }}
+  cancel-in-progress: true
+
 jobs:
   notebooks:
     name: "Build the notebooks for the docs"
@@ -27,6 +31,7 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install -U pip
+          python -m pip install "numpy<1.22"
           python -m pip install ".[tutorials]"
 
       - name: Get theano compiledir
diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
index 99a7992..75f38f3 100644
--- a/.github/workflows/wheels.yml
+++ b/.github/workflows/wheels.yml
@@ -14,14 +14,11 @@ jobs:
         with:
           submodules: true
           fetch-depth: 0
-      - uses: ilammy/msvc-dev-cmd@v1
-      - uses: joerick/cibuildwheel@v1.9.0
+      - uses: pypa/cibuildwheel@v2.3.1
         env:
-          CIBW_BUILD: "cp3?-*"
-          CIBW_SKIP: "cp35-* *-win32 *-manylinux_i686"
+          CIBW_SKIP: "pp* *-win32 *-manylinux_i686 *-musllinux*"
           CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
-          DISTUTILS_USE_SDK: 1
-          MSSdk: 1
+          CIBW_ARCHS_MACOS: "x86_64 universal2 arm64"
       - uses: actions/upload-artifact@v2
         with:
           path: ./wheelhouse/*.whl
diff --git a/.gitignore b/.gitignore
index 8ebc6ee..9a88f65 100644
--- a/.gitignore
+++ b/.gitignore
@@ -35,3 +35,6 @@ dist
 .envrc
 .coverage
 *.ipynb
+.coverage*
+.eggs
+.tox
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ca414c6..e279e01 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.0.1
+    rev: v4.1.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
diff --git a/pyproject.toml b/pyproject.toml
index eb74094..58a45a6 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,5 +1,5 @@
 [build-system]
-requires = ["setuptools>=40.6.0", "wheel", "setuptools_scm", "numpy>=1.13.0", "pybind11>=2.4"]
+requires = ["setuptools>=40.6.0", "wheel", "setuptools_scm", "oldest-supported-numpy", "pybind11>=2.4"]
 build-backend = "setuptools.build_meta"
 
 [tool.black]
@@ -28,5 +28,27 @@ use_parentheses = true
 known_third_party = ["numpy"]
 
 [tool.pytest.ini_options]
-norecursedirs = "python/test/*"
 addopts = "-v"
+
+[tool.coverage.run]
+parallel = true
+branch = true
+source = ["celerite2"]
+omit = [
+  "*_test.py",
+  "*__init__*",
+  "*/celerite2/celerite2_version.py",
+]
+
+[tool.coverage.paths]
+source = ["python", "*/site-packages"]
+
+[tool.coverage.report]
+show_missing = true
+exclude_lines = [
+  "raise NotImplementedError",
+  "except ImportError",
+  "pragma: no cover",
+  "def R_op",
+  "if verbose"
+]
diff --git a/python/celerite2/backprop.cpp b/python/celerite2/backprop.cpp
index 827b411..e5c813a 100644
--- a/python/celerite2/backprop.cpp
+++ b/python/celerite2/backprop.cpp
@@ -24,9 +24,9 @@ auto factor_fwd(py::array_t<double, py::array::c_style> t, py::array_t<double, p
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -80,9 +80,9 @@ auto factor_rev(py::array_t<double, py::array::c_style> t, py::array_t<double, p
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -140,11 +140,11 @@ auto solve_lower_fwd(py::array_t<double, py::array::c_style> t, py::array_t<doub
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -203,11 +203,11 @@ auto solve_lower_rev(py::array_t<double, py::array::c_style> t, py::array_t<doub
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -270,11 +270,11 @@ auto solve_upper_fwd(py::array_t<double, py::array::c_style> t, py::array_t<doub
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -333,11 +333,11 @@ auto solve_upper_rev(py::array_t<double, py::array::c_style> t, py::array_t<doub
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -400,11 +400,11 @@ auto matmul_lower_fwd(py::array_t<double, py::array::c_style> t, py::array_t<dou
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -463,11 +463,11 @@ auto matmul_lower_rev(py::array_t<double, py::array::c_style> t, py::array_t<dou
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -530,11 +530,11 @@ auto matmul_upper_fwd(py::array_t<double, py::array::c_style> t, py::array_t<dou
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -593,11 +593,11 @@ auto matmul_upper_rev(py::array_t<double, py::array::c_style> t, py::array_t<dou
 
   // Parse dimensions
   if (tbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t");
-  ssize_t N = tbuf.shape[0];
+  py::ssize_t N = tbuf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -662,13 +662,13 @@ auto general_matmul_lower_fwd(py::array_t<double, py::array::c_style> t1, py::ar
 
   // Parse dimensions
   if (t1buf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t1");
-  ssize_t N = t1buf.shape[0];
+  py::ssize_t N = t1buf.shape[0];
   if (t2buf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t2");
-  ssize_t M = t2buf.shape[0];
+  py::ssize_t M = t2buf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (t1buf.ndim != 1 || t1buf.shape[0] != N) throw std::invalid_argument("Invalid shape: t1");
@@ -723,13 +723,13 @@ auto general_matmul_upper_fwd(py::array_t<double, py::array::c_style> t1, py::ar
 
   // Parse dimensions
   if (t1buf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t1");
-  ssize_t N = t1buf.shape[0];
+  py::ssize_t N = t1buf.shape[0];
   if (t2buf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: t2");
-  ssize_t M = t2buf.shape[0];
+  py::ssize_t M = t2buf.shape[0];
   if (cbuf.ndim <= 0) throw std::invalid_argument("Invalid number of dimensions: c");
-  ssize_t J = cbuf.shape[0];
+  py::ssize_t J = cbuf.shape[0];
   if (Ybuf.ndim <= 1) throw std::invalid_argument("Invalid number of dimensions: Y");
-  ssize_t nrhs = Ybuf.shape[1];
+  py::ssize_t nrhs = Ybuf.shape[1];
 
   // Check shapes
   if (t1buf.ndim != 1 || t1buf.shape[0] != N) throw std::invalid_argument("Invalid shape: t1");
diff --git a/python/celerite2/driver.cpp b/python/celerite2/driver.cpp
index 15e6747..21dff0b 100644
--- a/python/celerite2/driver.cpp
+++ b/python/celerite2/driver.cpp
@@ -31,10 +31,10 @@ auto factor (
     // Parse dimensions
     if (tbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t");
-    ssize_t N = tbuf.shape[0];
+    py::ssize_t N = tbuf.shape[0];
     if (cbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: c");
-    ssize_t J = cbuf.shape[0];
+    py::ssize_t J = cbuf.shape[0];
 
     // Check shapes
     if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -82,13 +82,13 @@ auto solve_lower (
     // Parse dimensions
     if (tbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t");
-    ssize_t N = tbuf.shape[0];
+    py::ssize_t N = tbuf.shape[0];
     if (cbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: c");
-    ssize_t J = cbuf.shape[0];
+    py::ssize_t J = cbuf.shape[0];
     if (Ybuf.ndim <= 1)
         throw std::invalid_argument("Invalid number of dimensions: Y");
-    ssize_t nrhs = Ybuf.shape[1];
+    py::ssize_t nrhs = Ybuf.shape[1];
 
     // Check shapes
     if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -139,13 +139,13 @@ auto solve_upper (
     // Parse dimensions
     if (tbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t");
-    ssize_t N = tbuf.shape[0];
+    py::ssize_t N = tbuf.shape[0];
     if (cbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: c");
-    ssize_t J = cbuf.shape[0];
+    py::ssize_t J = cbuf.shape[0];
     if (Ybuf.ndim <= 1)
         throw std::invalid_argument("Invalid number of dimensions: Y");
-    ssize_t nrhs = Ybuf.shape[1];
+    py::ssize_t nrhs = Ybuf.shape[1];
 
     // Check shapes
     if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -196,13 +196,13 @@ auto matmul_lower (
     // Parse dimensions
     if (tbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t");
-    ssize_t N = tbuf.shape[0];
+    py::ssize_t N = tbuf.shape[0];
     if (cbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: c");
-    ssize_t J = cbuf.shape[0];
+    py::ssize_t J = cbuf.shape[0];
     if (Ybuf.ndim <= 1)
         throw std::invalid_argument("Invalid number of dimensions: Y");
-    ssize_t nrhs = Ybuf.shape[1];
+    py::ssize_t nrhs = Ybuf.shape[1];
 
     // Check shapes
     if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -253,13 +253,13 @@ auto matmul_upper (
     // Parse dimensions
     if (tbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t");
-    ssize_t N = tbuf.shape[0];
+    py::ssize_t N = tbuf.shape[0];
     if (cbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: c");
-    ssize_t J = cbuf.shape[0];
+    py::ssize_t J = cbuf.shape[0];
     if (Ybuf.ndim <= 1)
         throw std::invalid_argument("Invalid number of dimensions: Y");
-    ssize_t nrhs = Ybuf.shape[1];
+    py::ssize_t nrhs = Ybuf.shape[1];
 
     // Check shapes
     if (tbuf.ndim != 1 || tbuf.shape[0] != N) throw std::invalid_argument("Invalid shape: t");
@@ -312,16 +312,16 @@ auto general_matmul_lower (
     // Parse dimensions
     if (t1buf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t1");
-    ssize_t N = t1buf.shape[0];
+    py::ssize_t N = t1buf.shape[0];
     if (t2buf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t2");
-    ssize_t M = t2buf.shape[0];
+    py::ssize_t M = t2buf.shape[0];
     if (cbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: c");
-    ssize_t J = cbuf.shape[0];
+    py::ssize_t J = cbuf.shape[0];
     if (Ybuf.ndim <= 1)
         throw std::invalid_argument("Invalid number of dimensions: Y");
-    ssize_t nrhs = Ybuf.shape[1];
+    py::ssize_t nrhs = Ybuf.shape[1];
 
     // Check shapes
     if (t1buf.ndim != 1 || t1buf.shape[0] != N) throw std::invalid_argument("Invalid shape: t1");
@@ -376,16 +376,16 @@ auto general_matmul_upper (
     // Parse dimensions
     if (t1buf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t1");
-    ssize_t N = t1buf.shape[0];
+    py::ssize_t N = t1buf.shape[0];
     if (t2buf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: t2");
-    ssize_t M = t2buf.shape[0];
+    py::ssize_t M = t2buf.shape[0];
     if (cbuf.ndim <= 0)
         throw std::invalid_argument("Invalid number of dimensions: c");
-    ssize_t J = cbuf.shape[0];
+    py::ssize_t J = cbuf.shape[0];
     if (Ybuf.ndim <= 1)
         throw std::invalid_argument("Invalid number of dimensions: Y");
-    ssize_t nrhs = Ybuf.shape[1];
+    py::ssize_t nrhs = Ybuf.shape[1];
 
     // Check shapes
     if (t1buf.ndim != 1 || t1buf.shape[0] != N) throw std::invalid_argument("Invalid shape: t1");
@@ -442,7 +442,7 @@ auto get_celerite_matrices(
     auto U = U_out.mutable_unchecked<2>();
     auto V = V_out.mutable_unchecked<2>();
 
-    ssize_t N = x.shape(0), Jr = ar.shape(0), Jc = ac.shape(0), J = Jr + 2 * Jc;
+    py::ssize_t N = x.shape(0), Jr = ar.shape(0), Jc = ac.shape(0), J = Jr + 2 * Jc;
 
     if (bc.shape(0) != Jc) throw std::invalid_argument("dimension mismatch: bc");
     if (dc.shape(0) != Jc) throw std::invalid_argument("dimension mismatch: dc");
@@ -454,16 +454,16 @@ auto get_celerite_matrices(
     if (V.shape(0) != N || V.shape(1) != J) throw std::invalid_argument("dimension mismatch: V");
 
     double sum = 0.0;
-    for (ssize_t j = 0; j < Jr; ++j) sum += ar(j);
-    for (ssize_t j = 0; j < Jc; ++j) sum += ac(j);
+    for (py::ssize_t j = 0; j < Jr; ++j) sum += ar(j);
+    for (py::ssize_t j = 0; j < Jc; ++j) sum += ac(j);
 
-    for (ssize_t n = 0; n < N; ++n) {
+    for (py::ssize_t n = 0; n < N; ++n) {
         a(n) = diag(n) + sum;
-        for (ssize_t j = 0; j < Jr; ++j) {
+        for (py::ssize_t j = 0; j < Jr; ++j) {
             V(n, j) = 1.0;
             U(n, j) = ar(j);
         }
-        for (ssize_t j = 0, ind = Jr; j < Jc; ++j, ind += 2) {
+        for (py::ssize_t j = 0, ind = Jr; j < Jc; ++j, ind += 2) {
             double arg = dc(j) * x(n);
             double cos = V(n, ind) = std::cos(arg);
             double sin = V(n, ind + 1) = std::sin(arg);
diff --git a/python/spec/templates/backprop.cpp b/python/spec/templates/backprop.cpp
index 9077067..c444132 100644
--- a/python/spec/templates/backprop.cpp
+++ b/python/spec/templates/backprop.cpp
@@ -21,7 +21,7 @@ auto {{mod.name}}_fwd (
     {% for dim in mod.dimensions -%}
     if ({{mod.inputs[dim.coords[0]].name}}buf.ndim <= {{dim.coords[1]}})
         throw std::invalid_argument("Invalid number of dimensions: {{mod.inputs[dim.coords[0]].name}}");
-    ssize_t {{dim.name}} = {{mod.inputs[dim.coords[0]].name}}buf.shape[{{dim.coords[1]}}];
+    py::ssize_t {{dim.name}} = {{mod.inputs[dim.coords[0]].name}}buf.shape[{{dim.coords[1]}}];
     {% endfor %}
     // Check shapes
     {% for arg in mod.inputs + mod.outputs + mod.extra_outputs -%}
@@ -97,7 +97,7 @@ auto {{mod.name}}_rev (
     {% for dim in mod.dimensions -%}
     if ({{mod.inputs[dim.coords[0]].name}}buf.ndim <= {{dim.coords[1]}})
         throw std::invalid_argument("Invalid number of dimensions: {{mod.inputs[dim.coords[0]].name}}");
-    ssize_t {{dim.name}} = {{mod.inputs[dim.coords[0]].name}}buf.shape[{{dim.coords[1]}}];
+    py::ssize_t {{dim.name}} = {{mod.inputs[dim.coords[0]].name}}buf.shape[{{dim.coords[1]}}];
     {% endfor %}
     // Check shapes
     {% for arg in mod.rev_inputs + mod.rev_outputs -%}
diff --git a/python/spec/templates/driver.cpp b/python/spec/templates/driver.cpp
index 739b5d1..b26e329 100644
--- a/python/spec/templates/driver.cpp
+++ b/python/spec/templates/driver.cpp
@@ -22,7 +22,7 @@ auto {{mod.name}} (
     {% for dim in mod.dimensions -%}
     if ({{mod.inputs[dim.coords[0]].name}}buf.ndim <= {{dim.coords[1]}})
         throw std::invalid_argument("Invalid number of dimensions: {{mod.inputs[dim.coords[0]].name}}");
-    ssize_t {{dim.name}} = {{mod.inputs[dim.coords[0]].name}}buf.shape[{{dim.coords[1]}}];
+    py::ssize_t {{dim.name}} = {{mod.inputs[dim.coords[0]].name}}buf.shape[{{dim.coords[1]}}];
     {% endfor %}
     // Check shapes
     {% for arg in mod.inputs + mod.outputs -%}
@@ -98,7 +98,7 @@ auto get_celerite_matrices(
     auto U = U_out.mutable_unchecked<2>();
     auto V = V_out.mutable_unchecked<2>();
 
-    ssize_t N = x.shape(0), Jr = ar.shape(0), Jc = ac.shape(0), J = Jr + 2 * Jc;
+    py::ssize_t N = x.shape(0), Jr = ar.shape(0), Jc = ac.shape(0), J = Jr + 2 * Jc;
 
     if (bc.shape(0) != Jc) throw std::invalid_argument("dimension mismatch: bc");
     if (dc.shape(0) != Jc) throw std::invalid_argument("dimension mismatch: dc");
@@ -110,16 +110,16 @@ auto get_celerite_matrices(
     if (V.shape(0) != N || V.shape(1) != J) throw std::invalid_argument("dimension mismatch: V");
 
     double sum = 0.0;
-    for (ssize_t j = 0; j < Jr; ++j) sum += ar(j);
-    for (ssize_t j = 0; j < Jc; ++j) sum += ac(j);
+    for (py::ssize_t j = 0; j < Jr; ++j) sum += ar(j);
+    for (py::ssize_t j = 0; j < Jc; ++j) sum += ac(j);
 
-    for (ssize_t n = 0; n < N; ++n) {
+    for (py::ssize_t n = 0; n < N; ++n) {
         a(n) = diag(n) + sum;
-        for (ssize_t j = 0; j < Jr; ++j) {
+        for (py::ssize_t j = 0; j < Jr; ++j) {
             V(n, j) = 1.0;
             U(n, j) = ar(j);
         }
-        for (ssize_t j = 0, ind = Jr; j < Jc; ++j, ind += 2) {
+        for (py::ssize_t j = 0, ind = Jr; j < Jc; ++j, ind += 2) {
             double arg = dc(j) * x(n);
             double cos = V(n, ind) = std::cos(arg);
             double sin = V(n, ind + 1) = std::sin(arg);
diff --git a/python/test/jax/test_celerite2.py b/python/test/jax/test_jax_celerite2.py
similarity index 90%
rename from python/test/jax/test_celerite2.py
rename to python/test/jax/test_jax_celerite2.py
index ebf781e..e5fe75e 100644
--- a/python/test/jax/test_celerite2.py
+++ b/python/test/jax/test_jax_celerite2.py
@@ -2,9 +2,13 @@
 import celerite2
 import numpy as np
 import pytest
-from celerite2 import terms as pyterms
-from celerite2.jax import GaussianProcess, terms
-from celerite2.testing import check_gp_models
+
+try:
+    from celerite2 import terms as pyterms
+    from celerite2.jax import GaussianProcess, terms
+    from celerite2.testing import check_gp_models
+except (ImportError, ModuleNotFoundError):
+    pytestmark = pytest.mark.skip("jax not installed")
 
 term_mark = pytest.mark.parametrize(
     "name,args",
diff --git a/python/test/jax/test_ops.py b/python/test/jax/test_jax_ops.py
similarity index 89%
rename from python/test/jax/test_ops.py
rename to python/test/jax/test_jax_ops.py
index 2c2629e..c302642 100644
--- a/python/test/jax/test_ops.py
+++ b/python/test/jax/test_jax_ops.py
@@ -1,10 +1,15 @@
 # -*- coding: utf-8 -*-
 import numpy as np
-from celerite2 import driver
-from celerite2.jax import ops
-from celerite2.testing import get_matrices
-from jax import jit
-from jax.test_util import check_grads
+import pytest
+
+try:
+    from celerite2 import driver
+    from celerite2.jax import ops
+    from celerite2.testing import get_matrices
+    from jax import jit
+    from jax.test_util import check_grads
+except (ImportError, ModuleNotFoundError):
+    pytestmark = pytest.mark.skip("jax not installed")
 
 
 def check_op(op, input_arrays, expected_outputs, grad=True):
diff --git a/python/test/jax/test_terms.py b/python/test/jax/test_jax_terms.py
similarity index 80%
rename from python/test/jax/test_terms.py
rename to python/test/jax/test_jax_terms.py
index 56aa734..1bf4d15 100644
--- a/python/test/jax/test_terms.py
+++ b/python/test/jax/test_jax_terms.py
@@ -3,29 +3,22 @@
 
 import numpy as np
 import pytest
-from celerite2 import terms as pyterms
-from celerite2.testing import check_tensor_term
 
 try:
+    from celerite2 import terms as pyterms
+    from celerite2.jax import terms
+    from celerite2.testing import check_tensor_term
     from jax.config import config
-except ImportError:
-    HAS_JAX = False
+except (ImportError, ModuleNotFoundError):
+    pytestmark = pytest.mark.skip("jax not installed")
 else:
-    HAS_JAX = True
-
     config.update("jax_enable_x64", True)
 
-    from celerite2.jax import terms
-
-pytestmark = pytest.mark.skipif(not HAS_JAX, reason="jax is not installed")
-
-
-def evaluate(x):
-    assert x.dtype == np.float64 or x.dtype == np.int64
-    return np.asarray(x)
-
+    def evaluate(x):
+        assert x.dtype == np.float64 or x.dtype == np.int64
+        return np.asarray(x)
 
-compare_terms = partial(check_tensor_term, evaluate)
+    compare_terms = partial(check_tensor_term, evaluate)
 
 
 @pytest.mark.parametrize(
diff --git a/python/test/theano/conftest.py b/python/test/pymc3/conftest.py
similarity index 100%
rename from python/test/theano/conftest.py
rename to python/test/pymc3/conftest.py
diff --git a/python/test/theano/test_celerite2.py b/python/test/pymc3/test_pymc3_celerite2.py
similarity index 88%
rename from python/test/theano/test_celerite2.py
rename to python/test/pymc3/test_pymc3_celerite2.py
index 25dd3d2..b769826 100644
--- a/python/test/theano/test_celerite2.py
+++ b/python/test/pymc3/test_pymc3_celerite2.py
@@ -1,11 +1,16 @@
 # -*- coding: utf-8 -*-
+
 import celerite2
 import numpy as np
 import pytest
-from celerite2 import terms as pyterms
-from celerite2.testing import check_gp_models
-from celerite2.theano import GaussianProcess, terms
-from celerite2.theano.celerite2 import CITATIONS
+
+try:
+    from celerite2 import terms as pyterms
+    from celerite2.testing import check_gp_models
+    from celerite2.theano import GaussianProcess, terms
+    from celerite2.theano.celerite2 import CITATIONS
+except (ImportError, ModuleNotFoundError):
+    pytestmark = pytest.mark.skip("aesara_theano_fallback not installed")
 
 term_mark = pytest.mark.parametrize(
     "name,args",
@@ -97,7 +102,7 @@ def test_errors(data):
 
 
 def test_marginal(data):
-    import pymc3 as pm
+    pm = pytest.importorskip("pymc3")
 
     x, diag, y, t = data
 
@@ -113,7 +118,7 @@ def test_marginal(data):
 
 
 def test_citations(data):
-    import pymc3 as pm
+    pm = pytest.importorskip("pymc3")
 
     x, diag, y, t = data
 
diff --git a/python/test/theano/test_ops.py b/python/test/pymc3/test_pymc3_ops.py
similarity index 92%
rename from python/test/theano/test_ops.py
rename to python/test/pymc3/test_pymc3_ops.py
index 8dab51f..7802af7 100644
--- a/python/test/theano/test_ops.py
+++ b/python/test/pymc3/test_pymc3_ops.py
@@ -1,10 +1,16 @@
 # -*- coding: utf-8 -*-
-import aesara_theano_fallback.tensor as tt
+
 import numpy as np
-from aesara_theano_fallback import aesara as theano
-from celerite2 import backprop, driver
-from celerite2.testing import get_matrices
-from celerite2.theano import ops
+import pytest
+
+try:
+    import aesara_theano_fallback.tensor as tt
+    from aesara_theano_fallback import aesara as theano
+    from celerite2 import backprop, driver
+    from celerite2.testing import get_matrices
+    from celerite2.theano import ops
+except (ImportError, ModuleNotFoundError):
+    pytestmark = pytest.mark.skip("aesara_theano_fallback not installed")
 
 
 def convert_values_to_types(values):
diff --git a/python/test/theano/test_terms.py b/python/test/pymc3/test_pymc3_terms.py
similarity index 84%
rename from python/test/theano/test_terms.py
rename to python/test/pymc3/test_pymc3_terms.py
index 4848ed6..ce6211b 100644
--- a/python/test/theano/test_terms.py
+++ b/python/test/pymc3/test_pymc3_terms.py
@@ -1,13 +1,18 @@
 # -*- coding: utf-8 -*-
+
 from functools import partial
 
 import numpy as np
 import pytest
-from celerite2 import terms as pyterms
-from celerite2.testing import check_tensor_term
-from celerite2.theano import terms
 
-compare_terms = partial(check_tensor_term, lambda x: x.eval())
+try:
+    from celerite2 import terms as pyterms
+    from celerite2.testing import check_tensor_term
+    from celerite2.theano import terms
+except (ImportError, ModuleNotFoundError):
+    pytestmark = pytest.mark.skip("aesara_theano_fallback not installed")
+else:
+    compare_terms = partial(check_tensor_term, lambda x: x.eval())
 
 
 def test_complete_implementation():
diff --git a/setup.py b/setup.py
index 38268ec..f728989 100644
--- a/setup.py
+++ b/setup.py
@@ -35,12 +35,12 @@
     "style": ["isort", "black", "black_nbconvert"],
     "test": [
         "coverage[toml]",
-        "pytest==6.0.0rc1",
+        "pytest",
         "pytest-cov",
         "scipy",
         "celerite>=0.3.1",
     ],
-    "theano": [
+    "pymc3": [
         "pymc3>=3.9, <3.12",
         "aesara-theano-fallback>=0.0.2",
     ],
@@ -68,6 +68,7 @@
         "numpyro",
     ],
 }
+EXTRA_REQUIRE["theano"] = EXTRA_REQUIRE["pymc3"]
 EXTRA_REQUIRE["dev"] = (
     EXTRA_REQUIRE["style"]
     + EXTRA_REQUIRE["test"]
diff --git a/tox.ini b/tox.ini
new file mode 100644
index 0000000..548b77a
--- /dev/null
+++ b/tox.ini
@@ -0,0 +1,32 @@
+[tox]
+envlist = py{38,39,310}-{core,jax,pymc3},lint
+
+[gh-actions]
+python =
+    3.8: py38
+    3.9: py39
+    3.10: py310
+
+[gh-actions:env]
+EXTRAS =
+    core: core
+    jax: jax
+    pymc3: pymc3
+
+[testenv]
+passenv = GITHUB_*
+deps =
+    pymc3: numpy<1.22
+extras =
+    test
+    jax: jax
+    pymc3: pymc3
+commands =
+    pip freeze
+    python -m coverage run -m pytest -v {posargs}
+
+[testenv:lint]
+skip_install = true
+deps = pre-commit
+commands =
+    pre-commit run --all-files