diff --git a/.github/workflows/pypi-publish.yml b/.github/workflows/pypi-publish.yml
index 2efa91bc..cff7f5d9 100644
--- a/.github/workflows/pypi-publish.yml
+++ b/.github/workflows/pypi-publish.yml
@@ -10,7 +10,7 @@ jobs:
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       with:
         python-version: '3.x'
     - name: Check consistency between the package version and release tag
@@ -24,11 +24,11 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install setuptools wheel twine
+        pip install setuptools wheel twine build
     - name: Build and publish
       env:
         TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
         TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
       run: |
-        python setup.py sdist bdist_wheel
+        python -m build
         twine upload dist/*
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 96f30df0..bb9c9f14 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -24,7 +24,7 @@ jobs:
         with:
             python-version: "${{ matrix.python-version }}"
             cache: "pip"
-            cache-dependency-path: '**/requirements*.txt'
+            cache-dependency-path: 'pyproject.toml'
       - name: Run CI tests
         run: bash test.sh
         shell: bash
diff --git a/pyproject.toml b/pyproject.toml
index b1401d1e..855e0e87 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,2 +1,85 @@
-# Empty file in preparation for PR #513.
-# See https://github.com/deepmind/optax/pull/513 for more information.
+[build-system]
+requires = ["flit_core >=3.2,<4"]
+build-backend = "flit_core.buildapi"
+
+[project]
+name = "optax"
+dynamic = ["version"]
+description = "A gradient processing and optimisation library in JAX."
+readme = "README.md"
+license = { file = "LICENSE" }
+requires-python = ">=3.8"
+authors = [
+    {name = "DeepMind", email = "optax-dev@google.com"},
+]
+keywords = [
+    "python", 
+    "machine learning", 
+    "reinforcement-learning"
+]
+classifiers = [
+    "Environment :: Console",
+    "Programming Language :: Python",
+    "Intended Audience :: Developers",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Intended Audience :: Science/Research",
+    "Development Status :: 4 - Beta",
+    "License :: OSI Approved :: Apache Software License",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Software Development :: Libraries :: Python Modules",
+]
+dependencies = [
+  "absl-py>=0.7.1",
+  "chex>=0.1.5",
+  "jax>=0.1.55",
+  "jaxlib>=0.1.37",
+  "numpy>=1.18.0",
+]
+
+[project.urls]
+homepage = "https://github.com/deepmind/optax"
+repository = "https://github.com/deepmind/optax"
+documentation = "https://optax.readthedocs.io/"
+
+[project.optional-dependencies]
+test = [
+    "dm-haiku>=0.0.3",
+    "dm-tree>=0.1.7",
+    "flax==0.5.3"
+]
+
+examples = [
+  "dm-haiku>=0.0.3",
+  "tensorflow-datasets>=4.2.0",
+  "tensorflow>=2.4.0"
+]
+
+docs = [
+    "sphinx==4.5.0",
+    "sphinx-book-theme==0.3.3",
+    "sphinxcontrib-katex==0.9.0",
+    "sphinxcontrib-bibtex==2.4.2",
+    "sphinx-autodoc-typehints==1.11.1",
+    "IPython==7.16.3",
+    "ipykernel==5.3.4",
+    "pandoc==1.0.2",
+    "myst_nb==0.13.1",
+    "docutils==0.16",
+    "matplotlib==3.5.0",
+    "dm-haiku==0.0.8"
+]
+
+dp-accounting = [
+    "absl-py>=1.0.0",
+    "attrs>=21.4.0",
+    "mpmath>=1.2.1",
+    "numpy>=1.21.4",
+    "scipy>=1.7.1"
+]
+
+[tool.setuptools]
+zip_safe = false
+
+[tool.setuptools.packages.find]
+exclude = ["*_test.py"]
diff --git a/readthedocs.yml b/readthedocs.yml
index 6d3cdf10..5d53d443 100644
--- a/readthedocs.yml
+++ b/readthedocs.yml
@@ -11,7 +11,11 @@ sphinx:
 python:
   version: 3.8
   install:
-    - requirements: requirements/requirements-docs.txt
-    - requirements: requirements/requirements.txt
-    - method: setuptools
+    # Equivalent to 'pip install .'
+    - method: pip
       path: .
+    # Equivalent to 'pip install .[docs]'
+    - method: pip
+      path: .
+      extra_requirements:
+        - docs
diff --git a/test.sh b/test.sh
index 00845c5d..70a8bc48 100755
--- a/test.sh
+++ b/test.sh
@@ -26,16 +26,14 @@ python --version
 # Install dependencies.
 pip install --upgrade pip setuptools wheel
 pip install flake8 pytest-xdist pytype pylint pylint-exit
-pip install -r requirements/requirements.txt
-pip install -r requirements/requirements-test.txt
-pip install -r requirements/requirements-examples.txt
+pip install -e ".[test, examples]"
 
 # Dp-accounting specifies exact minor versions as requirements which sometimes
 # become incompatible with other libraries optax needs. We therefore install
 # dependencies for dp-accounting manually.
 # TODO(b/239416992): Remove this workaround if dp-accounting switches to minimum
 # version requirements.
-pip install -r requirements/minimum-requirements-dp-accounting.txt
+pip install -e ".[dp-accounting]"
 pip install "dp-accounting>=0.1.1" --no-deps
 
 # Ensure optax was not installed by one of the dependencies above,
@@ -60,7 +58,8 @@ pylint --rcfile=.pylintrc `find optax examples -name '*_test.py' | xargs` -d W02
 rm .pylintrc
 
 # Build the package.
-python setup.py sdist
+pip install build
+python -m build
 pip wheel --verbose --no-deps --no-clean dist/optax*.tar.gz
 pip install optax*.whl
 
@@ -78,7 +77,7 @@ python -m pytest -n auto .
 cd ..
 
 # Build Sphinx docs.
-pip install -r requirements/requirements-docs.txt
+pip install -e ".[docs]"
 cd docs && make html
 cd ..