Skip to content

Latest JAX and BMI versions are incompatible - #181

Open
RGoldsack wants to merge 3 commits into
cbg-ethz:mainfrom
RGoldsack:Install-issues
Open

RGoldsack wants to merge 3 commits into
cbg-ethz:mainfrom
RGoldsack:Install-issues

Conversation

@RGoldsack

Copy link
Copy Markdown

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.

RGoldsack and others added 3 commits September 15, 2025 15:24
… 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.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants