diff --git a/.github/workflows/deploy.yml b/.github/workflows/deploy.yml
index a9bcabf..2d3c762 100644
--- a/.github/workflows/deploy.yml
+++ b/.github/workflows/deploy.yml
@@ -26,6 +26,11 @@ jobs:
         run: |
           pip install --upgrade pip
           pip install build
+      - name: Patch the README links to point to the correct files at the current tag
+        run: |
+          perl -i.bak -pe 's/\[(.*)\]\((?!http)\.?\/?(.*)\.([a-z]+)\)/[\1](https:\/\/raw.github.com\/idiap\/RawSpeechClassification\/${{ github.event.release.tag_name }}\/\2.\3)/g' README.md
+          perl -i.bak -pe 's/\[(.*)\]\((?!http)\.?\/?(.*)\)/[\1](https:\/\/github.com\/idiap\/RawSpeechClassification\/tree\/${{ github.event.release.tag_name }}\/\2)/g' README.md
+          rm README.md.bak
       - name: Package the project
         run: python -m build
       - name: Produce a GitHub actions artifact (the package)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index f01bb78..0812d67 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,11 @@ SPDX-License-Identifier: GPL-3.0-only
 
 # Changelog
 
+## September 2024
+
+- Add Jax backend
+- Make pip installable package on PyPi
+
 ## August 2024
 
 - Update the code for Keras 3 with PyTorch or Tensorflow backend
diff --git a/README.md b/README.md
index 9eab828..c976425 100644
--- a/README.md
+++ b/README.md
@@ -6,6 +6,8 @@ SPDX-License-Identifier: GPL-3.0-only
 
 # Raw Speech Classification
 
+[![PyPI package](https://shields.io/pypi/v/raw-speech-classification.svg?logo=pypi)](https://pypi.org/project/raw-speech-classification)
+
 Trains CNN (or any neural network based) classifiers from raw speech using Keras and
 tests them. The inputs are lists of wav files, where each file is labelled. It then
 creates fixed length signals and processes them. During testing, it computes scores at
@@ -14,7 +16,7 @@ the fixed length signals.
 
 ## Installation
 
-### From source, In a Conda environment
+### From source in a conda environment
 
 To install Keras 3 with PyTorch backend, run:
 
@@ -28,6 +30,12 @@ To install Keras 3 with TensorFlow backend, run:
 conda env create -f conda/rsclf-tensorflow.yaml
 ```
 
+To install Keras 3 with Jax backend, run:
+
+```bash
+conda env create -f conda/rsclf-jax.yaml
+```
+
 Then install the package in that environment (the default name is `rsclf`) with:
 
 ```bash
@@ -49,7 +57,20 @@ or
 pip install raw-speech-classification[tensorflow]
 ```
 
-You'll also need to set the `KERAS_BACKEND` environment variable to the correct backend
+or
+
+```bash
+pip install raw-speech-classification[jax]
+```
+
+If you already have an environment with PyTorch, TensorFlow, or Jax
+installed, you can simply run:
+
+```bash
+pip install raw-speech-classification
+```
+
+You will also need to set the `KERAS_BACKEND` environment variable to the correct backend
 before running `rsclf-train` or `rsclf-test` (see below), or globally for the current
 bash session with:
 
@@ -57,7 +78,7 @@ bash session with:
 export KERAS_BACKEND=torch
 ```
 
-Replace `torch` by `tensorflow` accordingly.
+Replace `torch` by `tensorflow` or `jax` accordingly.
 
 ## Using the code
 
@@ -68,7 +89,7 @@ Replace `torch` by `tensorflow` accordingly.
    `root` option could be `/home/bob/data/my_dataset` and the content of the files would
    then be like:
 
-   ```txt
+   ```text
    part1/file1.wav 1
    part1/file2.wav 0
    ```
@@ -127,7 +148,7 @@ obtain the following curve in `results/seg-f1/plot.png`:
    probabilities. If you need the results per speaker, configure it accordingly (see the
    script for details). The default output format is:
 
-   ```txt
+   ```text
     <speakerID> <label> [<posterior_probability_vector>]
    ```
 
diff --git a/conda/rsclf-jax.yaml b/conda/rsclf-jax.yaml
new file mode 100644
index 0000000..8416ea1
--- /dev/null
+++ b/conda/rsclf-jax.yaml
@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: Copyright © Idiap Research Institute <contact@idiap.ch>
+#
+# SPDX-License-Identifier: GPL-3.0-only
+
+name: rsclf
+dependencies:
+  - python=3.11
+  - pip
+  - pip:
+    - keras
+    - h5py
+    - scipy
+    - jax[cuda12]
+    - matplotlib
+    - numpy
+    - polars
diff --git a/pyproject.toml b/pyproject.toml
index f3f321a..a3745f4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ build-backend = "flit_core.buildapi"
 
 [project]
 name = "raw-speech-classification"
-version = "1.0.1"
+version = "1.0.2"
 license = {file = "LICENSES/GPL-3.0-only.txt"}
 authors = [
     { name = "S. Pavankumar Dubagunta" },
@@ -50,6 +50,9 @@ torch = [
 tensorflow = [
     "tensorflow[and-cuda]",
 ]
+jax = [
+    "jax[cuda12]",
+]
 dev = [
     "pre-commit",
 ]