Inference-Gym: double precision issue

[reubenharry]

I want to use inference gym models with 64 bit precision (for use in Blackjax samplers). I encounter complicated looking bugs, which I have reduced to the following minimal example (run interactively in Python):

import jax
jax.config.update("jax_enable_x64", True)                                                                                                                               
import inference_gym.using_jax as gym 
gym.targets.Banana()._unnormalized_log_prob(jax.numpy.array([1.0,1.0]))

results in the following error. (Apologies if this is a simple mistake on my end)

Traceback (most recent call last):
File "", line 1, in
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/inference_gym/targets/banana.py", line 133, in _unnormalized_log_prob
return self._banana.log_prob(value)

File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py", line 1287, in log_prob
return self._call_log_prob(value, name, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/distributions/distribution.py", line 1269, in _call_log_prob
return self._log_prob(value, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/distributions/transformed_distribution.py", line 364, in _log_prob
log_prob, _ = self.experimental_local_measure(
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/distributions/transformed_distribution.py", line 611, in experimental_local_measure
x = self.bijector.inverse(y, **bijector_kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/bijectors/bijector.py", line 1390, in inverse
return self._call_inverse(y, name, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/bijectors/bijector.py", line 1370, in _call_inverse
return self._cache.inverse(y, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/internal/cache_util.py", line 347, in inverse
return self._lookup(y, self._inverse_name, self._forward_name, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/internal/cache_util.py", line 493, in _lookup
self._invoke(input, forward_name, kwargs, attrs))
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/internal/cache_util.py", line 532, in _invoke
return getattr(self.bijector, fn_name)(input, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/substrates/jax/bijectors/masked_autoregressive.py", line 383, in _inverse
bijector = self._bijector_fn(y, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/inference_gym/targets/banana.py", line 88, in bijector_fn
shift = tf.concat(
^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/_utils.py", line 62, in wrap
return new_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/numpy_array.py", line 161, in _concat
values = _args_to_matching_arrays(values)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/numpy_array.py", line 149, in _args_to_matching_arrays
ret = [ops.convert_to_tensor(arg, dtype) for arg in args_list]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/numpy_array.py", line 149, in
ret = [ops.convert_to_tensor(arg, dtype) for arg in args_list]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/_utils.py", line 62, in wrap
return new_fn(*args, **kwargs)
^^^^^^^^^^^^^^^^^^^^^^^
File "/global/homes/r/reubenh/.conda/envs/jax2024/lib/python3.11/site-packages/tensorflow_probability/python/internal/backend/jax/ops.py", line 152, in _convert_to_tensor
raise TypeError(('Tensor conversion requested dtype {} for array with '
TypeError: Tensor conversion requested dtype float32 for array with dtype float64: [-2.97]

[reubenharry]

I will note that I do not encounter this issue for all inference gym models. For instance, with target = gym.targets.VectorModel(gym.targets.BrownianMotionUnknownScalesMissingMiddleObservations()), I have no problem.

[SiegeLordEx]

This will need some dedicated changes to work. It'll likely look like a dtype arg for each target's initializer.

[reubenharry]

Is this something you'd recommend I try, or are there plans for the inference-gym developers to address this? (Currently this is a pretty serious blocker for me to use inference-gym for jax benchmarking of samplers, since I can't use double precision).

[SiegeLordEx]

I'm almost done making the necessary changes, should be done in a few days.

[SiegeLordEx]

Ok, so I submitted the changes. To the best of my knowledge, they should be sufficient, but I haven't done extensive end-to-end tests. I'll hold off making the release for a little bit, so to test it you'll need to install it via something like pip install -e 'git+https://github.com/tensorflow/probability.git#egg=inference_gym&subdirectory=spinoffs/inference_gym'

[reubenharry]

Thanks! This specific command (pip install -e 'git+https://github.com/tensorflow/probability.git#egg=inference_gym&subdirectory=spinoffs/inference_gym') did not resolve the error (i.e. what is described in this issue), but I can try cloning and doing pip install -e . if you point me to the relevant branch

EDIT: ah, I see the commit is merged, so presumably I can clone main and try a local install from that. Let me know if I also need to reinstall other packages

[SiegeLordEx]

Just cloning main and installing from that should be fine.

[reubenharry]

If I do the following:

git clone https://github.com/tensorflow/probability.git
cd probability/spinoffs/inference_gym/
pip install -e .

I get

Successfully installed inference_gym-0.0.6.dev0

and then if I do

import jax
jax.config.update("jax_enable_x64", True)                                                                                                                               
import inference_gym.using_jax as gym 
gym.targets.Banana()._unnormalized_log_prob(jax.numpy.array([1.0,1.0]))

I get

TypeError: Tensor conversion requested dtype float32 for array with dtype float64: [-2.97]

[SiegeLordEx]

Apologies, I neglected to explain exactly what the change was. For 64 bit support, all of the targets got a new argument, dtype which is set to the dtype you wish to work with. This is done to enable to revert back to 32 bits even if jax_enable_x64 is enabled. So, in this case, what you need to do is:

gym.targets.Banana(dtype=jax.numpy.float64)._unnormalized_log_prob(jax.numpy.array([1.0,1.0]))
Array(-12.02591216, dtype=float64)

[reubenharry]

Ah yep, resolved. Thanks so much!