Initializing values for multiple chains

Questions v5
1

Hello, I am trying to figure out how to initialize values for an MCMC sampler with multiple chains. I have no gradient information for my problem so at this stage I use DMetropolisZ. I found some information here but it does not really apply since it is for the NUTS sampler.

Here is the example code:

import numpy as np
import pymc as pm

chains = 4

with pm.Model(coords={"variable": ['a', 'b', 'c'],
              "property": ['x', 'y', 'z']}) as model:
    prior_x = pm.Uniform(name='x',
                         lower=0,
                         upper=1,
                         dims='variable')
    prior_y = pm.Uniform(name='y',
                         lower=1,
                         upper=2,
                         dims='variable')

    prior_z = pm.Uniform(name='z',
                         lower=2,
                         upper=3,
                         dims='variable')

    init_vals = {'x': [0.5]*3,
                 'y': [1.5]*3,
                 'z': [2.5]*3, }

    idata = pm.sample(draws=1000,
                      step=pm.DEMetropolisZ(),
                      initvals=init_vals,
                      chains=chains,
                      tune=500)

x_start = idata.posterior['x'].to_numpy()[:, 0, :]
y_start = idata.posterior['y'].to_numpy()[:, 0, :]
z_start = idata.posterior['z'].to_numpy()[:, 0, :]

assert np.all(x_start == 0.5)
assert np.all(y_start == 1.5)
assert np.all(z_start == 2.5)

First problem is that the init_values are not assigned (assertion fails).

Second problem is that I want to actually initialize values for each chain, so I would expect to be able to pass a matrix for each variable that has shape (3, 4). So (variable, chain). But using thie init dictionary:

    init_vals = {'x': np.ones((3, chains))*0.5,
                 'y': np.ones((3, chains))*1.5,
                 'z': np.ones((3, chains))*2.5}

gives the following error:

TypeError: Cannot convert Type TensorType(float64, (3, 4)) (of Variable Elemwise{add,no_inplace}.0) into Type TensorType(float64, (3,)). You can try to manually convert Elemwise{add,no_inplace}.0 into a TensorType(float64, (3,)).

The larger context of this is that I am running pymc with DMetropolisZ using about 40 chains on 20 variables with 3 coordinates each, and they always all initiate from the save value (center of the uniform prior distribution) and get stuck at local minima.

Thank you for the feedback!