From 6f894e06747d8551a698d9f4008ae633eec125a0 Mon Sep 17 00:00:00 2001 From: Roydon Date: Mon, 15 Sep 2025 15:24:31 +1200 Subject: [PATCH] With fresh installation of Python, latest version of packages BMI and JAX are incompatible. Tried installing `bmi` using `pip install benchmark-mi` with the latest version of jax (0.7.1) and a fresh Python install, and then running the suggested code in the getting started section: ``` import bmi task = bmi.benchmark.BENCHMARK_TASKS['1v1-normal-0.75'] print(f"Task {task.name} with dimensions {task.dim_x} and {task.dim_y}") print(f"Ground truth mutual information: {task.mutual_information:.2f}") X, Y = task.sample(1000, seed=42) cca = bmi.estimators.CCAMutualInformationEstimator() print(f"Estimate by CCA: {cca.estimate(X, Y):.2f}") ksg = bmi.estimators.KSGEnsembleFirstEstimator(neighborhoods=(5,)) print(f"Estimate by KSG: {ksg.estimate(X, Y):.2f}") ``` Get error, `AttributeError: jax.core.pp_eqn_rules was removed in JAX v0.4.34.`. This seems to be a package compatibility issue. Downgrading to jax==0.4.33 fixes the issue. --- README.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index b5cdb24..b489e86 100644 --- a/README.md +++ b/README.md @@ -31,7 +31,8 @@ $ pip install "bmi @ https://github.com/cbg-ethz/bmi" ``` Note: BMI uses [JAX](https://github.com/google/jax) and by default installs the CPU version of it. -If you have a device supporting CUDA, you can [install the CUDA version of JAX](https://github.com/google/jax#pip-installation-gpu-cuda-installed-via-pip-easier). +If you have a device supporting CUDA, you can [install the CUDA version of JAX](https://github.com/google/jax#pip-installation-gpu-cuda-installed-via-pip-easier). +Note BMI requires JAX==0.4.33. Now let's take one of the predefined distributions included in the benchmark (named "tasks") and sample 1,000 data points. Then, we will run two estimators on this task.