diff --git a/.github/workflows/actions.yml b/.github/workflows/actions.yml
deleted file mode 100644
index 3183a4f..0000000
--- a/.github/workflows/actions.yml
+++ /dev/null
@@ -1,74 +0,0 @@
-name: build
-
-env:
-  PYTHON_MAIN_VERSION: 3.9
-  
-on:
-  pull_request:
-    branches:
-      - '*'
-  push:
-    branches:
-      - '*'
-    tags:
-        - '*'
-
-jobs:
-  build-linux:
-    
-    strategy:
-      max-parallel: 5
-      matrix:
-        python-version: ['3.9', '3.10', '3.11']
-        os: [ubuntu-latest]
-
-    runs-on: ${{ matrix.os }}
-
-    steps:
-    - 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: |
-        sudo apt-get update -qq
-        python -m pip install --upgrade pip
-        python -m pip install flake8 pytest
-        python -m pip install jaxlib
-        python -m pip install jax
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-            
-    - name: install package
-      run: |
-        pip install .
-        pip list
-    
-    - name: Lint with flake8
-      run: |
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
-    
-    - name: Generate coverage report
-      run: |
-        pip install pytest
-        pip install pytest-cov
-        pytest --cov=./ --cov-report=xml
-        
-    - name: Upload coverage to Codecov
-      if: ${{ matrix.python-version == env.PYTHON_MAIN_VERSION }}
-      uses: codecov/codecov-action@v1
-      with:
-        token: ${{ secrets.CODECOV_TOKEN }}
-        files: ./coverage.xml
-        # directory: ./coverage/reports/
-        flags: unittests
-        env_vars: OS,PYTHON
-        name: codecov-umbrella
-        fail_ci_if_error: true
-        # path_to_write_report: ./coverage/codecov_report.txt
-        verbose: true
diff --git a/.github/workflows/ci-deploy.yml b/.github/workflows/ci-deploy.yml
new file mode 100644
index 0000000..14f09d0
--- /dev/null
+++ b/.github/workflows/ci-deploy.yml
@@ -0,0 +1,48 @@
+name: CI/Deploy
+
+on:
+  push:
+    tags: ["v*"]
+
+jobs:
+
+  unit:
+    uses: ./.github/workflows/unit.yml
+
+  notebooks:
+    uses: ./.github/workflows/notebook_smoke.yml
+
+  build_and_publish:
+    name: Upload release to PyPI
+    if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags')
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/gpax
+    permissions:
+      id-token: write
+
+    needs:
+      - unit
+      - notebooks
+
+    steps:
+
+    - name: Checkout
+      uses: actions/checkout@v3
+      with:
+        fetch-depth: 0
+
+    - name: Set up Python 3.9
+      uses: actions/setup-python@v2
+      with:
+        python-version: 3.9
+
+    - name: Build and apply version
+      run: bash scripts/build.sh
+
+    - name: Publish package distributions to PyPI
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        # CURRENTLY USING TEST SERVER FOR NOW!!!!
+        repository-url: https://test.pypi.org/legacy/
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..b2df5c1
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,17 @@
+name: CI
+
+on:
+
+  pull_request:
+    branches: ["master", "dev/*", mc/issue-59]
+
+  push:
+    branches: ["master", "dev/*", mc/issue-59]
+
+jobs: 
+
+  unit:
+    uses: ./.github/workflows/unit.yml
+
+  notebooks:
+    uses: ./.github/workflows/notebook_smoke.yml
diff --git a/.github/workflows/notebook_smoke.yml b/.github/workflows/notebook_smoke.yml
index 312a5df..a6d8ab7 100644
--- a/.github/workflows/notebook_smoke.yml
+++ b/.github/workflows/notebook_smoke.yml
@@ -4,14 +4,7 @@ env:
   CI_SMOKE: True
 
 on:
-  pull_request:
-    branches:
-      - '*'
-  push:
-    branches:
-      - '*'
-    tags:
-      - '*'
+  workflow_call:
 
 jobs:
   build-linux:
@@ -34,22 +27,9 @@ jobs:
 
     - name: Install dependencies
       run: |
-        sudo apt-get update -qq
-        python -m pip install --upgrade pip
-        python -m pip install flake8 pytest
-        python -m pip install jaxlib
-        python -m pip install jax
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-            
-    - name: install package
-      run: |
-        pip install .
-        pip list
+        bash scripts/install.sh
 
     - name: Notebook smoke tests
       run: |
-        pip install ipython
-        pip install nbformat
-        pip install seaborn
         bash scripts/test_notebooks.sh
 
diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml
deleted file mode 100644
index 8ba346d..0000000
--- a/.github/workflows/publish-to-pypi.yml
+++ /dev/null
@@ -1,50 +0,0 @@
-name: Publish Python 🐍 distribution 📦 to PyPI and TestPyPI
-
-on: push
-
-jobs:
-  build:
-    name: Build distribution 📦
-    runs-on: ubuntu-latest
-
-    steps:
-    - uses: actions/checkout@v4
-    - name: Set up Python
-      uses: actions/setup-python@v4
-      with:
-        python-version: "3.x"
-    - name: Install pypa/build
-      run: >-
-        python3 -m
-        pip install
-        build
-        --user
-    - name: Build a binary wheel and a source tarball
-      run: python3 -m build
-    - name: Store the distribution packages
-      uses: actions/upload-artifact@v3
-      with:
-        name: python-package-distributions
-        path: dist/
-
-  publish-to-pypi:
-    name: >-
-      Publish Python 🐍 distribution 📦 to PyPI
-    if: startsWith(github.ref, 'refs/tags/')  # only publish to PyPI on tag pushes
-    needs:
-    - build
-    runs-on: ubuntu-latest
-    environment:
-      name: pypi
-      url: https://pypi.org/p/gpax
-    permissions:
-      id-token: write  # IMPORTANT: mandatory for trusted publishing
-
-    steps:
-    - name: Download all the dists
-      uses: actions/download-artifact@v3
-      with:
-        name: python-package-distributions
-        path: dist/
-    - name: Publish distribution 📦 to PyPI
-      uses: pypa/gh-action-pypi-publish@release/v1
diff --git a/.github/workflows/unit.yml b/.github/workflows/unit.yml
new file mode 100644
index 0000000..d7b69d0
--- /dev/null
+++ b/.github/workflows/unit.yml
@@ -0,0 +1,54 @@
+name: Unit
+
+env:
+  PYTHON_MAIN_VERSION: 3.9
+
+on:
+  workflow_call:
+
+jobs:
+  build:
+    runs-on: ${{ matrix.os }}
+
+    strategy:
+      matrix:
+        os: [ubuntu-latest]
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+      - uses: actions/checkout@v3
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v2
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          bash scripts/install.sh
+          bash scripts/install.sh test
+
+      - name: Lint with flake8
+        run: |
+          # stop the build if there are Python syntax errors or undefined names
+          flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
+          # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
+          flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
+
+      - name: Run tests
+        run: |
+          pytest -v --cov --cov-report xml tests
+
+      - name: Upload coverage to Codecov
+        if: ${{ matrix.python-version == env.PYTHON_MAIN_VERSION }}
+        uses: codecov/codecov-action@v1
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          files: ./coverage.xml
+          # directory: ./coverage/reports/
+          flags: unittests
+          env_vars: OS,PYTHON
+          name: codecov-umbrella
+          fail_ci_if_error: true
+          # path_to_write_report: ./coverage/codecov_report.txt
+          verbose: true
diff --git a/.gitignore b/.gitignore
index b6e4761..c0fb41c 100644
--- a/.gitignore
+++ b/.gitignore
@@ -127,3 +127,6 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
+
+# Pyright
+pyrightconfig.json
diff --git a/gpax/__init__.py b/gpax/__init__.py
index 84de056..17e4f82 100644
--- a/gpax/__init__.py
+++ b/gpax/__init__.py
@@ -1,13 +1,47 @@
-from .__version__ import version as __version__
-from . import priors
-from . import utils
-from . import kernels
-from . import acquisition
-from .hypo import sample_next
-from .models import (DKL, CoregGP, ExactGP, MultiTaskGP, iBNN, vExactGP,
-                     vi_iBNN, viDKL, viGP, sPM, viMTDKL, VarNoiseGP, UIGP,
-                     MeasuredNoiseGP, viSparseGP, BNN)
+from gpax import acquisition, kernels, utils
+from gpax._version import __version__
+from gpax.hypo import sample_next
+from gpax.models import (
+    BNN,
+    DKL,
+    UIGP,
+    CoregGP,
+    ExactGP,
+    MeasuredNoiseGP,
+    MultiTaskGP,
+    VarNoiseGP,
+    iBNN,
+    sPM,
+    vExactGP,
+    vi_iBNN,
+    viDKL,
+    viGP,
+    viMTDKL,
+    viSparseGP,
+)
 
-__all__ = ["priors", "utils", "kernels", "mtkernels", "acquisition", "ExactGP", "vExactGP", "DKL",
-           "viDKL", "iBNN", "vi_iBNN", "MultiTaskGP", "viMTDKL", "viGP", "sPM", "VarNoiseGP",
-           "UIGP", "MeasuredNoiseGP", "viSparseGP", "CoregGP", "BNN", "sample_next", "__version__"]
+__all__ = [
+    "priors",
+    "utils",
+    "kernels",
+    "mtkernels",
+    "acquisition",
+    "ExactGP",
+    "vExactGP",
+    "DKL",
+    "viDKL",
+    "iBNN",
+    "vi_iBNN",
+    "MultiTaskGP",
+    "viMTDKL",
+    "viGP",
+    "sPM",
+    "VarNoiseGP",
+    "UIGP",
+    "MeasuredNoiseGP",
+    "viSparseGP",
+    "CoregGP",
+    "BNN",
+    "sample_next",
+    "__version__",
+]
diff --git a/gpax/__version__.py b/gpax/__version__.py
deleted file mode 100644
index 7fce6b3..0000000
--- a/gpax/__version__.py
+++ /dev/null
@@ -1 +0,0 @@
-version = '0.1.8'
diff --git a/gpax/_version.py b/gpax/_version.py
new file mode 100644
index 0000000..9807bb2
--- /dev/null
+++ b/gpax/_version.py
@@ -0,0 +1,10 @@
+"""Version file. This is overwritten during build and will contain a static
+__version__ variable."""
+
+try:
+    import dunamai as _dunamai
+
+    __version__ = _dunamai.Version.from_any_vcs().serialize()
+    del _dunamai
+except ImportError:
+    __version__ = "dev"
diff --git a/pyproject.toml b/pyproject.toml
index 1eeed4c..b225a6f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,3 +1,45 @@
+[build-system]
+requires = ["flit_core >=3.2,<4"]
+build-backend = "flit_core.buildapi"
+
+[project]
+name = "gpax"
+authors = [
+    {"name" = "Maxim Ziatdinov", "email" = "maxim.ziatdinov@ai4microcopy.com"}
+]
+readme = "README.md"
+requires-python = ">=3.9"
+license = {"file" = "LICENSE"}
+description='Gaussian processes in NumPyro and JAX'
+classifiers=[
+    'Programming Language :: Python',
+    'Development Status :: 3 - Alpha',
+    'Intended Audience :: Science/Research',
+    'Operating System :: POSIX :: Linux',
+    'Operating System :: MacOS :: MacOS X',
+    'Topic :: Scientific/Engineering'
+]
+
+# Core dependencies
+dependencies = [
+    "matplotlib>=3.1.1",
+    "jax>=0.4.8",
+    "numpyro>=0.11.0",
+    "dm-haiku>=0.0.5",
+    "dunamai==1.19.2",
+    "jaxopt", # TODO: do we want to require this?
+]
+
+dynamic = ["version"]
+
+[project.optional-dependencies]
+test = [
+    "flake8",
+    "jaxlib",
+    "pytest",
+    "pytest-cov",
+]
+
 [tool.black]
 line-length = 127
 include = '\.pyi?$'
diff --git a/scripts/build.sh b/scripts/build.sh
new file mode 100644
index 0000000..e746e62
--- /dev/null
+++ b/scripts/build.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+
+pip install flit~=3.7
+pip install dunamai==1.19.2
+echo "__version__ = '$(dunamai from any --style=pep440 --no-metadata)'" >gpax/_version.py
+flit build
diff --git a/scripts/install.sh b/scripts/install.sh
new file mode 100644
index 0000000..6ca7894
--- /dev/null
+++ b/scripts/install.sh
@@ -0,0 +1,10 @@
+#!/bin/bash
+
+pip install toml
+
+if [ "$1" = "test" ]; then
+	python3 -c 'import toml; c = toml.load("pyproject.toml"); print("\n".join(c["project"]["optional-dependencies"]["test"]))' | pip install -r /dev/stdin
+
+else
+	python3 -c 'import toml; c = toml.load("pyproject.toml"); print("\n".join(c["project"]["dependencies"]))' | pip install -r /dev/stdin
+fi
diff --git a/scripts/test_notebooks.sh b/scripts/test_notebooks.sh
index f1f1771..e63cef5 100644
--- a/scripts/test_notebooks.sh
+++ b/scripts/test_notebooks.sh
@@ -1,6 +1,8 @@
 #!/bin/bash
 
-#!/bin/bash
+pip install ipython
+pip install nbformat
+pip install seaborn
 for nb in examples/*.ipynb; do
     echo "Running notebook smoke test on $nb"
     ipython -c "%run $nb"
diff --git a/setup.py b/setup.py
deleted file mode 100644
index 537b639..0000000
--- a/setup.py
+++ /dev/null
@@ -1,40 +0,0 @@
-__author__ = "Maxim Ziatdinov"
-__copyright__ = "Copyright Maxim Ziatdinov (2021)"
-__maintainer__ = "Maxim Ziatdinov"
-__email__ = "maxim.ziatdinov@ai4microcopy.com"
-
-from setuptools import setup, find_packages
-import os
-
-module_dir = os.path.dirname(os.path.abspath(__file__))
-
-with open(os.path.join(module_dir, 'gpax/__version__.py')) as f:
-    __version__ = f.read().split("'")[1]
-
-if __name__ == "__main__":
-    setup(
-        name='gpax',
-        python_requires='>=3.9',
-        version=__version__,
-        description='Gaussian processes in NumPyro and JAX',
-        long_description=open(os.path.join(module_dir, 'README.md')).read(),
-        long_description_content_type='text/markdown',
-        url='https://github.com/ziatdinovmax/gpax/',
-        author='Maxim Ziatdinov',
-        author_email='maxim.ziatdinov@ai4microcopy.com',
-        license='MIT license',
-        packages=find_packages(),
-        zip_safe=False,
-        install_requires=[
-            'jax>=0.4.8',
-            'numpyro>=0.11.0',
-            'dm-haiku>=0.0.5',
-            'matplotlib>=3.1'
-        ],
-        classifiers=['Programming Language :: Python',
-                     'Development Status :: 3 - Alpha',
-                     'Intended Audience :: Science/Research',
-                     'Operating System :: POSIX :: Linux',
-                     'Operating System :: MacOS :: MacOS X',
-                     'Topic :: Scientific/Engineering']
-    )