tensorflow / probability

Probabilistic reasoning and statistical analysis in TensorFlow
https://www.tensorflow.org/probability/
Apache License 2.0
4.16k stars 1.08k forks source link

autobnn error #1804

Closed 1kaiser closed 3 months ago

1kaiser commented 3 months ago

https://colab.research.google.com/github/tensorflow/probability/blob/main/discussion/examples/Forecasting_With_AutoBNN.ipynb

already satisfied: fastprogress>=1.0.0 in /usr/local/lib/python3.10/dist-packages (from blackjax->bayeux-ml->autobnn) (1.0.3) Requirement already satisfied: equinox>=0.10.6 in /usr/local/lib/python3.10/dist-packages (from flowmc->bayeux-ml->autobnn) (0.11.3) Requirement already satisfied: evosax>=0.1.4 in /usr/local/lib/python3.10/dist-packages (from flowmc->bayeux-ml->autobnn) (0.1.6) Requirement already satisfied: tqdm in /usr/local/lib/python3.10/dist-packages (from flowmc->bayeux-ml->autobnn) (4.66.2) Requirement already satisfied: multipledispatch in /usr/local/lib/python3.10/dist-packages (from numpyro->bayeux-ml->autobnn) (1.0.0) Requirement already satisfied: lineax>=0.0.4 in /usr/local/lib/python3.10/dist-packages (from optimistix->bayeux-ml->autobnn) (0.0.4) Requirement already satisfied: cached_property in /usr/local/lib/python3.10/dist-packages (from orbax->flax->autobnn) (1.5.2) Requirement already satisfied: importlib_resources in /usr/local/lib/python3.10/dist-packages (from orbax->flax->autobnn) (6.4.0) Requirement already satisfied: etils in /usr/local/lib/python3.10/dist-packages (from orbax->flax->autobnn) (1.7.0) Requirement already satisfied: dotmap in /usr/local/lib/python3.10/dist-packages (from evosax>=0.1.4->flowmc->bayeux-ml->autobnn) (1.3.30) Requirement already satisfied: h5py in /usr/local/lib/python3.10/dist-packages (from h5netcdf>=1.0.2->arviz->bayeux-ml->autobnn) (3.9.0) Requirement already satisfied: mdurl~=0.1 in /usr/local/lib/python3.10/dist-packages (from markdown-it-py>=2.2.0->rich>=11.1->flax->autobnn) (0.1.2)

AttributeError Traceback (most recent call last) in <cell line: 9>() 7 import matplotlib.pyplot as plt 8 ----> 9 from autobnn import estimators 10 from autobnn import operators 11 from autobnn import kernels

4 frames /usr/local/lib/python3.10/dist-packages/flax/configurations.py in define_boolstate(name, default, help) 38 on runtime with `flax.config.update('flax', )`. 39 """ ---> 40 return jax_config.define_boolstate('flax' + name, default, help) 41 42

AttributeError: 'Config' object has no attribute 'define_bool_state'

ColCarroll commented 3 months ago

Hey! Thanks for opening this issue!

Can you clarify how you set up this environment? If I go to https://colab.research.google.com/github/tensorflow/probability/blob/main/discussion/examples/Forecasting_With_AutoBNN.ipynb it executes just fine.

Specifically, this looks like it uses an older version of flax -- if you have a reproducible example, I could probably add a version guard.

1kaiser commented 3 months ago

thanks ,

We recommend using a TPU runtime for this colab. this made me think to select TPU before running

which is connected by default

image