`pymc.sampling_jax.sample_blackjax_nuts` crashes with `chains=1`

[aussetg]

Description of your problem

pymc.sampling_jax.sample_blackjax_nuts crashes with chains=1

Please provide a minimal, self-contained, and reproducible example.

# %%
import numpy as np
import pandas as pd
import pymc as pm
import pymc.sampling_jax
import pymc.distributions.transforms
import aesara.tensor as at
import aesara
np.set_printoptions(2)

# %%
data = pd.DataFrame(np.random.randn(100, 23))

# %%
K = 15
n_f = len(data.columns)

# %%
with pm.Model() as model:
    alpha = pm.Gamma("alpha", 0.7, 1.0)
    w = pm.StickBreakingWeights("w", alpha=alpha, K=K-1)

    cov = np.identity(n_f)

    mu = pm.MvNormal(
        "mu",
        mu=data.mean(axis=0),
        cov=cov,
        shape=(K, n_f),
    )

    components = [
        pm.MvNormal.dist(mu=mu[i, :], cov=cov, shape=(n_f,))
        for i in range(K)
    ]

    obs = pm.Mixture(
        "obs",
        w=w,
        comp_dists=components,
        observed=data,
    )

# %%
with model:
    trace = pymc.sampling_jax.sample_blackjax_nuts(target_accept=0.9, chains=1)

Please provide the full traceback.

Compiling...
Traceback (most recent call last):
  File "/usr/lib64/python3.9/code.py", line 90, in runcode
    exec(code, self.locals)
  File "<input>", line 2, in <module>
  File "/home/guillaume/.cache/pypoetry/virtualenvs/tabular-T_K5n0PG-py3.9/lib/python3.9/site-packages/pymc/sampling_jax.py", line 301, in sample_blackjax_nuts
    init_params = [np.stack(init_params)]
  File "<__array_function__ internals>", line 180, in stack
  File "/home/guillaume/.cache/pypoetry/virtualenvs/tabular-T_K5n0PG-py3.9/lib64/python3.9/site-packages/numpy/core/shape_base.py", line 426, in stack
    raise ValueError('all input arrays must have the same shape')
ValueError: all input arrays must have the same shape

Versions and main components

• PyMC/PyMC3 Version: PyMC4.1.1
• Aesara/Theano Version: 2.7.4
• Python Version: 3.9
• Operating system: Linux ( Fedora 36 )
• How did you install PyMC/PyMC3: (conda/pip) Pip ( poetry )

[bherwerth]

Hi @aussetg ,

I can reproduce the error and can run your code with chains=2.

The error comes from here (line 301):

pymc/pymc/sampling_jax.py

Lines 300 to 302 in ab40dcc

	if chains == 1:
	init_params = [np.stack(init_params)]
	init_params = [np.stack(init_state) for init_state in zip(*init_params)]

and init_params is returned from _get_batched_jittered_initial_points, where the case chains == 1 is treated here:

pymc/pymc/sampling_jax.py

Lines 168 to 171 in ab40dcc

	if chains == 1:
	initial_points = initial_points[0]
	else:
	initial_points = [np.stack(init_state) for init_state in zip(*initial_points)]

Could it thus be that np.stack in line 301 is not needed and it would be enough to "undo" the indexing of line 169 by using [init_params] in line 301 (or removing line 301 and using zip(init_params) in line 302)?

[aussetg]

Yes it works with chains > 1 but chains > 1 doesn't work with a GPU as it complains about the number of chains being greater than the number of XLA devices.

[bherwerth]

I did some further investigation and I think that indeed line 301 in sampling_jax.py causes a problem when init_params comes from priors don't have the same shape.

Here is a simpler example that fails, it's basically the test case in test_sampling_jax.py but sigma with shape (2,):

# %%
import numpy as np
import pandas as pd
import pymc as pm
import pymc.sampling_jax
import pymc.distributions.transforms
import aesara.tensor as at
import aesara

sampler = pymc.sampling_jax.sample_blackjax_nuts
chains = 1
postprocessing_backend = "cpu"

aesara.config.on_opt_error = "raise"
np.random.seed(13244)

obs = np.random.normal(10, 2, size=100)
obs_at = aesara.shared(obs, borrow=True, name="obs")
# %%
with pm.Model() as model:
    a = pm.Uniform("a", -20, 20)
    sigma2 = pm.HalfNormal("sigma", shape=(2,))
    sigma = 0.5*(sigma2[0] + sigma2[1])
    b = pm.Normal("b", a, sigma=sigma, observed=obs_at)

    trace = sampler(
        chains=chains,
        random_seed=1322,
        keep_untransformed=True,
        postprocessing_backend=postprocessing_backend,
    )

The code runs when I change remove line 301 in sampling_jax.py and use the following for line 302:

        init_params = [np.stack(init_state) for init_state in zip(init_params)]

If there's agreement about the problem with line 301 & 302, I am happy to propose a PR with a fix.