`fit_laplace` errors when deterministics or data do not have dims

[jessegrabowski]

Reported here. The following model fails:

import pymc as pm
import pymc_extras as pmx
import numpy as np


with pm.Model() as m:
    X = pm.Data('X', np.random.normal(size=(100, 3)))

    alpha = pm.Normal("alpha", 100, 5, shape=(1,))
    beta = pm.Normal("beta", mu=0, sigma=3, shape=(3,))

    mu = pm.Deterministic("mu", alpha + X @ beta)
    sigma = pm.Exponential("sigma", 1, shape=(1,))

    y_hat = pm.Normal("y_hat", mu=mu, sigma=sigma)
    
    data = pm.sample_prior_predictive(draws=1).prior.y_hat.values.squeeze()
    
with pm.observe(m, {'y_hat': data}):
    pmx.fit_laplace()

Full Traceback

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[64], line 15
     12     data = pm.sample_prior_predictive(draws=1).prior.y_hat.values.squeeze()
     14 with pm.observe(m, {'y_hat': data}):
---> 15     pmx.fit_laplace()

File ~/mambaforge/envs/grabowski_phd/lib/python3.12/site-packages/pymc_extras/inference/laplace_approx/laplace.py:428, in fit_laplace(optimize_method, model, use_grad, use_hessp, use_hess, initvals, random_seed, jitter_rvs, progressbar, include_transformed, gradient_backend, chains, draws, optimizer_kwargs, compile_kwargs)
    413     H_inv = _compute_inverse_hessian(
    414         optimizer_result=None,
    415         optimal_point=idata.fit.mean_vector.values,
   (...)    419         method=optimize_method,
    420     )
    422     idata.fit["covariance_matrix"] = xr.DataArray(
    423         H_inv,
    424         dims=("rows", "columns"),
    425         coords={"rows": unpacked_variable_names, "columns": unpacked_variable_names},
    426     )
--> 428 with model_to_laplace_approx(model, unpacked_variable_names, chains, draws) as laplace_model:
    429     new_posterior = (
    430         pm.sample_posterior_predictive(
    431             idata.fit.expand_dims(chain=[0], draw=[0]),
   (...)    441         .rename({"temp_chain": "chain", "temp_draw": "draw"})
    442     )
    444     if include_transformed:

File ~/mambaforge/envs/grabowski_phd/lib/python3.12/site-packages/pymc_extras/inference/laplace_approx/laplace.py:223, in model_to_laplace_approx(model, unpacked_variable_names, chains, draws)
    220             initval = initial_point.get(name, None)
    221             dim_shapes = initval.shape if initval is not None else batched_rv.type.shape[2:]
    222             laplace_model.add_coords(
--> 223                 {name: np.arange(shape) for name, shape in zip(dims[2:], dim_shapes)}
    224             )
    226         pm.Deterministic(name, batched_rv, dims=dims)
    228 return laplace_model

TypeError: unsupported operand type(s) for -: 'NoneType' and 'int'