Return value for 2n-dimensional ODE system

Thanks! I’m not sure whether I fully understand. I tried:

n=2

def odefunc(y, t, p):
    y1 = y[:n]
    y2 = y[n:]
    dy1 = -p[0]*y1*y2
    dy2 = p[0]*y1*y2 
    return np.concatenate([dy1, dy2]).T

pm.ode.DifferentialEquation(odefunc, times=[1], n_states=2*n, n_theta=1)

which throws an error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-152-11eea9a67857> in <module>
      8     return np.concatenate([dy1, dy2]).T
      9 
---> 10 pm.ode.DifferentialEquation(odefunc, times=[1], n_states=2*n, n_theta=1)

~/software/pymc3/pymc3/ode/ode.py in __init__(self, func, times, n_states, n_theta, t0)
     72         # Private
     73         self._augmented_times = np.insert(times, 0, t0).astype(floatX)
---> 74         self._augmented_func = utils.augment_system(func, self.n_states, self.n_theta)
     75         self._sens_ic = utils.make_sens_ic(self.n_states, self.n_theta, floatX)
     76 

~/software/pymc3/pymc3/ode/utils.py in augment_system(ode_func, n_states, n_theta)
     89 
     90     # Get symbolic representation of the ODEs by passing tensors for y, t and theta
---> 91     yhat = ode_func(t_y, t_t, t_p[n_states:])
     92     # Stack the results of the ode_func into a single tensor variable
     93     if not isinstance(yhat, (list, tuple)):

<ipython-input-152-11eea9a67857> in odefunc(y, t, p)
      6     dy1 = -p[0]*y1*y2
      7     dy2 = p[0]*y1*y2
----> 8     return np.concatenate([dy1, dy2]).T
      9 
     10 pm.ode.DifferentialEquation(odefunc, times=[1], n_states=2*n, n_theta=1)

<__array_function__ internals> in concatenate(*args, **kwargs)

ValueError: zero-dimensional arrays cannot be concatenated