Integrating MLNE with Pymc wrappers

[sheensu]

Hi all, I'm exploring how to integrate MNLE with pymc wrappers following this thread #987
And in particular this colab code https://colab.research.google.com/drive/1UmxuFYNo8VnlIHlFnbBI6DdyJiEzPExf?usp=sharing

I have 8 parameters in my model and [continuous value, discrete value] as the model output
Here are my codes just to get the pymc sampler run:

potential_fn = LikelihoodBasedPotential(mnle_estimator, prior, x_o=torch.tensor([[0.5, 0]], dtype=torch.float32))

def lp_f(x, track_gradients=True):
    with torch.set_grad_enabled(track_gradients):
        return potential_fn.__call__(x)

logl = PyMCPotential(lp_f, track_gradients=True, device="cpu")
dim = 8

with pymc.Model() as opmodel:

    theta = pm.Uniform( "theta", lower=lower, upper=upper, shape=dim)
    # use a Potential
    pm.Potential("likelihood", logl(theta))

    idata_grad = pm.sample(step=pm.NUTS(), tune=100, draws=500, chains=1, mp_ctx="spawn", initvals={"theta": np.array((0.0, 0.0, 0.002, 0.0, 5, 5, 5, 5))})

_ = az.plot_trace(idata_grad)

And I got these errors:

File ~/anaconda3/envs/sbi_env/lib/python3.10/site-packages/sbi/neural_nets/estimators/shape_handling.py:74, in reshape_to_batch_event(theta_or_x, event_shape)
     72 trailing_theta_or_x_shape = theta_or_x.shape[-event_shape_dim:]
     73 leading_theta_or_x_shape = theta_or_x.shape[:-event_shape_dim]
---> [74](https://file+.vscode-resource.vscode-cdn.net/Users/[]/Downloads/~/anaconda3/envs/sbi_env/lib/python3.10/site-packages/sbi/neural_nets/estimators/shape_handling.py:74) assert trailing_theta_or_x_shape == event_shape, (
     75     "The trailing dimensions of `theta_or_x` do not match the `event_shape`."
     76 )
     78 if len(leading_theta_or_x_shape) == 0:
     79     # A single datapoint is passed. Add batch artificially.
     80     return theta_or_x.unsqueeze(0)

AssertionError: The trailing dimensions of `theta_or_x` do not match the `event_shape`.
Apply node that caused the error: PyMCPotential(Composite{...}.2)
Toposort index: 1
Inputs types: [TensorType(float64, shape=(8,))]
Inputs shapes: [(8,)]
Inputs strides: [(8,)]
Inputs values: ['not shown']
Outputs clients: [[Add(Sum{axes=None}.0, likelihood)], [SpecifyShape(PyMCPotential.1, 8)]]

I've tried to implement different versions in the colab code, but I got similar errors. Does anyone have any suggestions? Thank you in advance!!

[satwiksps]

Hi @sheensu

I am not totally sure, but it looks like the error might come from a shape mismatch between what PyMC passes into the MNLE potential and what MNLE expects. PyMC seems to send a 1-D (8,) float64 array, while MNLE usuaaly expects a batched float32 tensor like (1, 8).

I had a similar issue before, and adding a small reshape helped:

def lp_f(x, track_gradients=True):
    with torch.set_grad_enabled(track_gradients):
        x = torch.tensor(x, dtype=torch.float32)
        if x.ndim == 1:
            x = x.unsqueeze(0)
        return potential_fn(x)

[sheensu]

Thank you so much @satwiksps, it's very helpful!!