Importing inference gym

Skip to first unread message

Charles Margossian

Jun 30, 2023, 6:13:50 PM6/30/23
to TensorFlow Probability
I'm using a notebook that runs experiments on the inference gym and I'm having issues importing the gym. The code works on Colab and I had it working earlier on my local machine. But something broke when I fired up a new notebook...

The only additional step I took before starting the new notebook is running the following commandpip install tensorflow_probability[tfds]

which was recommended after I attempted to load the German Credit example, and which is required to load real datasets (and that worked just fine). But now I wonder if this may have broken something.

Here's the error:

----> 6 from inference_gym import using_jax as gym
and the error message at the end:

ValueError: Arg specs do not match: original=FullArgSpec(args=['input', 'out_type', 'name'], varargs=None, varkw=None, defaults=(None, None), kwonlyargs=[], kwonlydefaults=None, annotations={}), new=FullArgSpec(args=['input', 'out_type', 'name'], varargs=None, varkw=None, defaults=(tf.int32, None), kwonlyargs=[], kwonlydefaults=None, annotations={}), fn=<function size_v2 at 0x7feeed7c78b0>
The full message is much no longer and I can post it if it helps. I'm looking for some guidance interpreting the error message.

Thank you for your help!Charles

Charles Margossian

Jun 30, 2023, 6:34:21 PM6/30/23
to TensorFlow Probability, Charles Margossian

So I uninstalled TFP,

pip uninstall tensorflow_probability

and installed the nightly version via

!rm -Rf probability
!rm -Rf fun_mc
!rm -Rf inference_gym
!git clone
!mv probability/spinoffs/fun_mc/fun_mc .
!mv probability/spinoffs/inference_gym/inference_gym .
!pip install tf-nightly tfp-nightly jax jaxlib

This removes the error I had previously, but now importing the gym returns the following error

AttributeError: module 'tensorflow_probability.substrates.jax' has no attribute 'tf2jax'

Pavel Sountsov

Jun 30, 2023, 7:42:20 PM6/30/23
to Charles Margossian, TensorFlow Probability
I just tried the tutorial notebook and that still appeared to work. Note that after the pip command succeeded, I restarted the runtime since I think that's necessary to pick up the newly installed packages rather than the builin TensorFlow etc.

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
To view this discussion on the web visit

Charles Margossian

Jul 5, 2023, 11:21:23 AM7/5/23
to TensorFlow Probability,, TensorFlow Probability, Charles Margossian
Thanks for the pointer Pavel. I reinstalled the packages and then started a new kernel, and this seems to solve the problem. I'll keep testing on more examples.
Reply all
Reply to author
0 new messages