Importing inference gym on colab

[Charles Margossian]

Hi all,

This morning colleagues and I have tried importing inference gym on colab. We tried both our notebook and the inference gym tutorial and got the following error message:

--------------------------------------------------------------------------- RuntimeError Traceback (most recent call last) RuntimeError: module compiled against API version 0x10 but this version of numpy is 0xf . Check the section C-API incompatibility at the Troubleshooting ImportError section at https://numpy.org/devdocs/user/troubleshooting-importerror.html#c-api-incompatibility for indications on how to solve this problem .

--------------------------------------------------------------------------- RuntimeError Traceback (most recent call last) RuntimeError: module compiled against API version 0x10 but this version of numpy is 0xf . Check the section C-API incompatibility at the Troubleshooting ImportError section at https://numpy.org/devdocs/user/troubleshooting-importerror.html#c-api-incompatibility for indications on how to solve this problem .

--------------------------------------------------------------------------- RuntimeError Traceback (most recent call last) RuntimeError: module compiled against API version 0x10 but this version of numpy is 0xf . Check the section C-API incompatibility at the Troubleshooting ImportError section at https://numpy.org/devdocs/user/troubleshooting-importerror.html#c-api-incompatibility for indications on how to solve this problem .

--------------------------------------------------------------------------- ImportError Traceback (most recent call last) ImportError: numpy.core._multiarray_umath failed to import

--------------------------------------------------------------------------- ImportError Traceback (most recent call last) ImportError: numpy.core.umath failed to import

--------------------------------------------------------------------------- TypeError Traceback (most recent call last) <ipython-input-2-7acf21dc9221> in <cell line: 3>() 1 import functools 2 ----> 3 import inference_gym.using_jax as gym 4 import jax 5 from jax import lax

---

15 frames

---

/usr/local/lib/python3.10/dist-packages/tensorflow/python/framework/dtypes.py in <module> 35 from tensorflow.tsl.python.lib.core import pywrap_ml_dtypes 36 ---> 37 _np_bfloat16 = pywrap_ml_dtypes.bfloat16() 38 _np_float8_e4m3fn = pywrap_ml_dtypes.float8_e4m3fn() 39 _np_float8_e5m2 = pywrap_ml_dtypes.float8_e5m2() TypeError: Unable to convert function return value to a Python type! The signature was () -> handle

Colab AI suggested uninstalling TensorFlow to avoid a conflict with Jax. We tried that but the error message persisted.

[Colin Carroll]

Hey Charles -- 

tl;dr: Change the installation line to just

`!pip install -U inference_gym tfds_nightly`

My best guess is that somehow the new tensorflow release bumped the numpy version and that hasn't quite worked itself out. The repro worked, and adding the line

!pip install -U "numpy~=1.23"

after the other installs allows importing inference gym along with all the nightlies. Embarrassingly, it breaks arviz (I think by breaking numba).

There are also three calls to `jax.tree_multimap` that need to be change to `jax.tree_util.tree_map`.

--
You received this message because you are subscribed to the Google Groups "TensorFlow Probability" group.
To unsubscribe from this group and stop receiving emails from it, send an email to tfprobabilit...@tensorflow.org.
To view this discussion on the web visit https://groups.google.com/a/tensorflow.org/d/msgid/tfprobability/64407265-22fa-48c2-b452-233100e7161cn%40tensorflow.org.

[Manny Mokel]

Hello,

I have been having similar issues to Charles using TFP on JAX in colab with inference gym as well as several other modules (I can't import tfp.distributions and others). I'm currently cloning the TFP repository from git, then executing '!pip install tf-nightly tfp-nightly jax jaxlib'. 

I was having these issues Friday as well, and seemed to be able to temporarily fix them by installing TensorFlow 2.13.0 in colab. I also tried changing the version of numpy to the one above, but it did not do anything for me. The consistent error is 

'Unable to convert function return value to a Python type! The signature was

() -> handle'

[Colin Carroll]

Hi-- Sorry for the install troubles! 

It is harder for me to check on an environment like this, but a few things:

1. Are you running in an isolated environment (virtualenv, conda, docker)? (the answer here should be "yes", otherwise you might be getting imports from a bunch of wild places!)

2. Are you using both jax and tensorflow in the same environment? If no, do things work with just `!pip install tfp-nightly[jax] jax jaxlib`? That should be a faster install, as well.

3. I ended up looking at this to figure out the right numpy version to install. I seem to recall a similar error coming to yours when I was troubleshooting. Can you import numpy and like, multiply matrices from the environment you made? That would isolate the problem as being with a version mismatch for numpy.

Feel free to share more detail about how you're installing things, the output of `pip list`, and a more specific stack trace!

--C

[Manny Mokel (Intern)]

Hi Colin,

Uninstalling TensorFlow and running '!pip install tfp-nightly[jax] jax jaxlib' seemed to fix things'. Thank you!

[Manny Mokel]

I've just gone to run this again, and I'm now getting a different error (but on all the same import lines): 

"ImportError: cannot import name 'control_flow_case' from 'tensorflow.python.ops' (/usr/local/lib/python3.10/dist-packages/tensorflow/python/ops/__init__.py)"

This is with running "!pip install tfp-nightly[jax] jax jaxlib" from a Google Colab notebook, and uninstalling TensorFlow 

It's not just with inference gym, but even on import tensorflow_probability as _tfp

[Colin Carroll]

Maybe you need to be `import tensorflow_probability.substrates.jax as tfp` and `import inference_gym.using_jax as gym`, which should avoid importing tensorflow at all. If you can share more about how you're running it and a more detailed stack trace I may be more helpful!

[Manny Mokel]

Here's an example notebook of the things I'm trying to import: https://colab.research.google.com/drive/1VEAiLwIFz6IQVZApCbrLFCuL2U03AWK3?usp=sharing 

Here's an example of an error I get when trying to import inference_gym. I get this error on most import lines listed in the document:


ImportError Traceback (most recent call last) <ipython-input-6-eb0b88b493cb> in <cell line: 7>() 5 from jax import numpy as jnp 6 ----> 7 from inference_gym import using_jax as gym 8 9 from fun_mc import using_jax as fun_mcmc

---

16 frames

---

/content/inference_gym/using_jax.py in <module> 16 # Need to register the rewrite hooks. 17 from inference_gym.backends import rewrite ---> 18 from inference_gym.dynamic.backend_jax import targets 19 del rewrite 20 /content/inference_gym/targets/__init__.py in <module> 18 with util.silence_nonrewritten_import_errors(): 19 # pylint: disable=g-import-not-at-top ---> 20 from inference_gym.targets.banana import Banana 21 from inference_gym.targets.bayesian_model import BayesianModel 22 from inference_gym.targets.brownian_motion import BrownianMotion /content/inference_gym/targets/banana.py in <module> 15 16 import numpy as np ---> 17 import tensorflow.compat.v2 as tf 18 import tensorflow_probability as tfp 19 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/internal/lazy_loader.py in __getattr__(self, item) 51 52 def __getattr__(self, item): ---> 53 module = self._load() 54 return getattr(module, item) 55 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/internal/lazy_loader.py in _load(self) 38 self._on_first_access = None 39 # Import the target module and insert it into the parent's namespace ---> 40 module = importlib.import_module(self.__name__) 41 if self._parent_module_globals is not None: 42 self._parent_module_globals[self._local_name] = module /usr/lib/python3.10/importlib/__init__.py in import_module(name, package) 124 break 125 level += 1 --> 126 return _bootstrap._gcd_import(name[level:], package, level) 127 128 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/substrates/jax/__init__.py in <module> 37 del _ensure_jax_install # Cleanup symbol to avoid polluting namespace. 38 ---> 39 from tensorflow_probability.python.version import __version__ 40 # from tensorflow_probability.substrates.jax.google import staging # DisableOnExport # pylint:disable=line-too-long 41 from tensorflow_probability.substrates.jax import bijectors /usr/local/lib/python3.10/dist-packages/tensorflow_probability/__init__.py in <module> 21 # from tensorflow_probability.google import staging # DisableOnExport 22 # from tensorflow_probability.google import tfp_google # DisableOnExport ---> 23 from tensorflow_probability.python import * # pylint: disable=wildcard-import 24 from tensorflow_probability.python.version import __version__ 25 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/__init__.py in <module> 136 # Non-lazy load of packages that register with tensorflow or keras. 137 for pkg_name in _maybe_nonlazy_load: --> 138 dir(globals()[pkg_name]) # Forces loading the package from its lazy loader. 139 140 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/internal/lazy_loader.py in __dir__(self) 55 56 def __dir__(self): ---> 57 module = self._load() 58 return dir(module) 59 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/internal/lazy_loader.py in _load(self) 38 self._on_first_access = None 39 # Import the target module and insert it into the parent's namespace ---> 40 module = importlib.import_module(self.__name__) 41 if self._parent_module_globals is not None: 42 self._parent_module_globals[self._local_name] = module /usr/lib/python3.10/importlib/__init__.py in import_module(name, package) 124 break 125 level += 1 --> 126 return _bootstrap._gcd_import(name[level:], package, level) 127 128 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/experimental/__init__.py in <module> 29 30 from tensorflow_probability.python.experimental import auto_batching ---> 31 from tensorflow_probability.python.experimental import bayesopt 32 from tensorflow_probability.python.experimental import bijectors 33 from tensorflow_probability.python.experimental import distribute /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/experimental/bayesopt/__init__.py in <module> 15 """TensorFlow Probability experimental Bayesopt package.""" 16 ---> 17 from tensorflow_probability.python.experimental.bayesopt import acquisition 18 from tensorflow_probability.python.internal import all_util 19 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/experimental/bayesopt/acquisition/__init__.py in <module> 15 """Acquisition Functions.""" 16 ---> 17 from tensorflow_probability.python.experimental.bayesopt.acquisition.acquisition_function import AcquisitionFunction 18 from tensorflow_probability.python.experimental.bayesopt.acquisition.acquisition_function import MCMCReducer 19 from tensorflow_probability.python.experimental.bayesopt.acquisition.expected_improvement import GaussianProcessExpectedImprovement /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/experimental/bayesopt/acquisition/acquisition_function.py in <module> 20 21 from tensorflow_probability.python.internal import dtype_util ---> 22 from tensorflow_probability.python.internal import prefer_static as ps 23 from tensorflow_probability.python.internal import tensor_util 24 /usr/local/lib/python3.10/dist-packages/tensorflow_probability/python/internal/prefer_static.py in <module> 27 28 from tensorflow.python.framework import tensor_util # pylint: disable=g-direct-tensorflow-import ---> 29 from tensorflow.python.ops import control_flow_case # pylint: disable=g-direct-tensorflow-import 30 from tensorflow.python.util import tf_inspect # pylint: disable=g-direct-tensorflow-import 31 ImportError: cannot import name 'control_flow_case' from 'tensorflow.python.ops' (/usr/local/lib/python3.10/dist-packages/tensorflow/python/ops/__init__.py)