PyMC3 Dirichlet Process Multivariate Gaussian Mixture Model

Question

I'm having trouble getting my shapes to work for a Dirichlet Process Gaussian Mixture Model. My data observations has shape (number of samples, number of dimensions). Each Gaussian's mean should be drawn from an isotropic prior, and each Gaussian's covariance should be the identity matrix. I thought I set this up correctly, but I'm getting the following error:

Input dimension mis-match. (input[0].shape[1] = 13, input[1].shape[1] = 2)

My code is:

import numpy as np
import pymc3 as pm
import theano.tensor as tt

num_obs, obs_dim = observations.shape
max_num_clusters = 13

def stick_breaking(beta):
    portion_remaining = tt.concatenate([[1], tt.extra_ops.cumprod(1 - beta)[:-1]])
    return beta * portion_remaining

with pm.Model() as model:
    w = pm.Deterministic("w", stick_breaking(beta))
    cluster_means = pm.MvNormal(f'cluster_means',
                                mu=pm.floatX(np.zeros(obs_dim)),
                                cov=pm.floatX(gaussian_mean_prior_cov_scaling * np.eye(obs_dim)),
                                shape=(max_num_clusters, obs_dim))

    comp_dists = pm.MvNormal.dist(mu=cluster_means,
                                  cov=gaussian_cov_scaling * np.eye(obs_dim),
                                  shape=(max_num_clusters, obs_dim))

    obs = pm.Mixture(
        "obs",
        w=w,
        comp_dists=comp_dists,
        observed=observations,
        shape=obs_dim)

Can someone clarify how to get the shapes to work?