Add jax/jaxlib version recommendation to README; update jax version i…

…n pre-commit hooks

.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

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