diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index c6c5559..4778c7d 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -87,7 +87,7 @@ repos:
       - id: mypy
         files: src|tests
         additional_dependencies:
-          - jax>=0.4.20
+          - jax>=0.4.34
           - numpy
           - packaging
           - pandas-stubs
diff --git a/README.md b/README.md
index 1eec254..6fb2132 100644
--- a/README.md
+++ b/README.md
@@ -5,17 +5,23 @@
 [![image](https://codecov.io/gh/OpenSourceEconomics/lcm/branch/main/graph/badge.svg)](https://codecov.io/gh/OpenSourceEconomics/lcm)
 
 This package aims to generalize and facilitate the specification, solution, and
-estimation of dynamic choice models.
+simulation of discrete-continuous dynamic choice models.
 
 ## Installation
 
-`LCM` currently needs to be installed from GitHub and requires `jax` and `jaxlib`. If
-you aim to run `LCM` on a GPU, you need to install `jaxlib` with CUDA support (for
+`LCM` currently needs to be installed from GitHub and requires `jax` and `jaxlib`.
+
+> [!IMPORTANT]
+> We recommend installing at least version "0.4.34" of `jax` and `jaxlib`, due to memory
+> issues with previous versions.
+
+If you aim to run `LCM` on a GPU, you need to install `jaxlib` with CUDA support (for
 Linux) or with support for AMD GPUs / ARM-based Silicon GPUs (for MacOS). In any case,
 for installation of `jax` and `jaxlib`, please consult the `jax`
 [docs](https://jax.readthedocs.io/en/latest/installation.html#supported-platforms).
 
-> **Note:** We currently test GPU support for `LCM` only on Linux with CUDA 12.
+> [!NOTE]
+> We currently test GPU support for `LCM` only on Linux with CUDA 12.
 
 ## License