Hierarchical model - Shape of observed data

Questions
1

Hello,

I am trying to fit a hierarchical model for the parameters of an SIR model. I am using the ‘sunode’ model to solve the ODE and assuming that I have two groups in my data the data structure returned from sunode (‘res’) is [num x 2]. My question here is:

Do I not need to flatten this data structure and the data structure of the observed data? Or, can I pass both of them as a [num x 2] data structure?

with pm.Model() as model4:
            sigma = pm.HalfCauchy('sigma', self.likelihood['sigma'])
            prior_lam = pm.Normal('prior_lam', 1.5, 1.5)
            prior_mu = pm.Normal('prior_mu', 1.5, 1.5)
            prior_lam_std = pm.HalfNormal('prior_lam_std', 5)
            prior_mu_std = pm.HalfNormal('prior_mu_std', 5)
          
            lam_mu = prior_lam
            mu_mu = prior_mu
            lam = pm.Lognormal('lambda', lam_mu , prior_lam_std, shape=2) 
            mu = pm.Lognormal('mu', mu_mu, prior_mu_std, shape=2)           

            res, _, problem, solver, _, _ = sunode.wrappers.as_theano.solve_ivp(
            y0={
                'S': (self.S_init, (2,)),
                'I': (self.I_init, (2,)),},
            params={
                'lam': (lam, (2,)),
                'mu': (mu, (2,)),
                '_dummy': (np.array(1.), ())},
            rhs=self.SIR_sunode,
            # The time points where we want to access the solution
            tvals=self.time_range,
            t0=self.time_range[0]
            )
            
            self.debug = res['I']
           
            if(likelihood['distribution'] == 'lognormal'):
                I = pm.Lognormal('I', mu=res['I'], sigma=sigma, observed=self.cases_obs_scaled)

Thanks!