Multivariate semi-parametric Cox proportional hazard

Hello PyMC3 World!

I just started using PyMC3 and thus have to excuse for my very beginners questions in advance.

I am trying to use a semi-parametric Cox Regression model for survival analysis as described in this post by @AustinRochford :

https://docs.pymc.io/en/v3/pymc-examples/examples/survival_analysis/survival_analysis.html

In order to perform multivariate analysis I was trying to use the shape parameter for beta to assign all regression coefficients in a single stochastic variable. The model looks something like this:

covariates = [covariate1, covariate2]  

with pm.Model(coords=coords) as model:
      lambda0 = pm.Gamma("lambda0", 0.01, 0.01, dims="intervals")

      beta = pm.Normal("beta", 0, sigma=1000, shape=len(covariates))

      exponent = pm.math.dot(data[covariates].to_numpy(dtype=float), beta)

      lambda_ = pm.Deterministic("lambda_",
                                 T.outer(T.exp(exponent), lambda0))
      
      mu = pm.Deterministic("mu", exposure * lambda_)

      obs = pm.Poisson("obs", mu, observed=death)

wehre data is a pandas dataframe which holds the survival time as well as the covariate values in its columns.

After sampling I expected a beta from the posterior of shape (2,) with two regression coefficients for each covariate under inspection.
It turns out that both beta values are equal which is not what I expected.

I think there must be something wrong in the exponent evaluation but I can’t find a good solution besides splitting beta into two separate variables like beta0 and beta1 which is not the way I want to implement it since there are more covariates which I want to consider in the future.

Thanks in advance and again sorry for such a newbie question.

Best,
Paul