inference gym / numpyro compatibility

[kylejcaron]

trafficstars

I'm looking to use Inference Gym targets in numpyro, but I'm running into issues I believe because there are numpy arrays in the Inference Gym model init, which causes tracer conversion errors in numpyro/jax

Any ideas how to get around this? I cant tell how the inference_gym.using_jax module works, but I was hoping that it would change the arrays to be initialized as jax arrays and not numpy

import jax
import numpyro
from numpyro import distributions as dist
from numpyro.infer import MCMC, NUTS
from inference_gym import using_jax as gym

class Banana(dist.Distribution):
    arg_constraints = {"ndims": dist.constraints.positive_integer, "curvature": dist.constraints.real}
    support = dist.constraints.real_vector
    pytree_data_fields = ("ndims", "curvature")

    def __init__(self, ndims, curvature):
        self.ndims = ndims
        self.curvature = curvature
        self.gym_dist = gym.targets.Banana(ndims=ndims, curvature=curvature)
        super().__init__(event_shape=(ndims,))
    
    def sample(self, key, sample_shape=()):
        return self.gym_dist.sample(seed=key, sample_shape=sample_shape)
    
    def log_prob(self, value):
        return self.gym_dist._unnormalized_log_prob(value)
    

samples = Banana(ndims=3, curvature=0.03).sample(jax.random.PRNGKey(0), sample_shape=(100,))

def model(X):
    curvature = numpyro.sample("curvature", dist.Beta(1,30))
    return numpyro.sample("obs", Banana(ndims=3, curvature=curvature), obs=X)

kernel = NUTS(model)
mcmc = MCMC(kernel, num_warmup=500, num_samples=1000)
mcmc.run(jax.random.PRNGKey(0), X=samples)

Here's the full traceback

Traceback

File "/var/folders/1p/v01fvg3j1cz1hzv988ygj8m00000gn/T/ipykernel_28885/3778384098.py", line 37, in mcmc.run(jax.random.PRNGKey(0), X=samples) File "/.venv/lib/python3.11/site-packages/numpyro/infer/mcmc.py", line 702, in run states_flat, last_state = partial_map_fn(map_args) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/mcmc.py", line 465, in _single_chain_mcmc new_init_state = self.sampler.init( ^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/hmc.py", line 749, in init init_params = self._init_state( ^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/hmc.py", line 693, in _init_state ) = initialize_model( ^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 713, in initialize_model (init_params, pe, grad), is_valid = find_valid_initial_params( ^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 447, in find_valid_initial_params (init_params, pe, z_grad), is_valid = _find_valid_params( ^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 433, in _find_valid_params _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state) ^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 417, in body_fn pe, z_grad = value_and_grad(potential_fn)(params) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/traceback_util.py", line 180, in reraise_with_filtered_traceback return fun(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/api.py", line 468, in value_and_grad_f ans, vjp_py = _vjp(f_partial, *dyn_args) ^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/api.py", line 1975, in _vjp out_primals, vjp = ad.vjp(flat_fun, primals_flat) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/ad.py", line 252, in vjp out_primals, pvals, jaxpr, consts = linearize(traceable, *primals) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/ad.py", line 237, in linearize jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/profiler.py", line 333, in wrapper return func(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py", line 574, in trace_to_jaxpr_nounits jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals) ^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/linear_util.py", line 192, in call_wrapped return self.f_transformed(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py", line 587, in trace_to_subjaxpr_nounits out_tracers, jaxpr, out_consts, env = _trace_to_subjaxpr_nounits( ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/partial_eval.py", line 616, in _trace_to_subjaxpr_nounits ans = f(*in_args) ^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/api_util.py", line 72, in flatten_fun ans = f(*py_args, **py_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/ad.py", line 78, in jvpfun out_primals, out_tangents = f(tag, primals, tangents) ^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/interpreters/ad.py", line 115, in jvp_subtrace ans = f(*in_tracers) ^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/api_util.py", line 88, in flatten_fun_nokwargs ans = f(*py_args) ^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/api_util.py", line 292, in _argnums_partial return fun(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 299, in potential_energy log_joint, model_trace = log_density( ^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/infer/util.py", line 70, in log_density model_trace = trace(model).get_trace(*model_args, **model_kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/handlers.py", line 186, in get_trace self(*args, **kwargs) File "/.venv/lib/python3.11/site-packages/numpyro/primitives.py", line 105, in call return self.fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/primitives.py", line 105, in call return self.fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/primitives.py", line 105, in call return self.fn(*args, **kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^ [Previous line repeated 3 more times] File "/var/folders/1p/v01fvg3j1cz1hzv988ygj8m00000gn/T/ipykernel_28885/3778384098.py", line 31, in model return numpyro.sample("obs", Banana(ndims=3, curvature=curvature), obs=X) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/numpyro/distributions/distribution.py", line 100, in call return super().call(*args, kwargs) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/var/folders/1p/v01fvg3j1cz1hzv988ygj8m00000gn/T/ipykernel_28885/3778384098.py", line 17, in init self.gym_dist = gym.targets.Banana(ndims=ndims, curvature=curvature) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/inference_gym/targets/banana.py", line 116, in init [10.] + [np.sqrt(1. + 2 * curvature2 * 10.**4)] + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.venv/lib/python3.11/site-packages/jax/_src/core.py", line 692, in array raise TracerArrayConversionError(self)

[SiegeLordEx]

While this particular issue is fixable in principle (see below), the core issue is that this Inference Gym targets are not intended to be used this way. The targets are high level constructs that can do IO and other things not compatible with jitted computation in their initializer. They're not "distributions" in the sense of being building blocks to constructing larger probabilistic models.

If you want a local fix, edit the /.venv/lib/python3.11/site-packages/inference_gym/targets/banana.py source code on that line to look like:

ground_truth_standard_deviation=tf.constant(
   [10.] + [tf.sqrt(1. + 2 * curvature**2 * 10.**4)] +
   [1.] * (ndims - 2)),

That's just a local workaround. It's unclear what the proper solution would look like because there's a lot of corner cases.