You could try a lower higher target_accept, or adjusting the priors. These are the usual superficial solutions to divergences. Deeper solutions require thinking about why the model might not be a good fit for the data and making structural adjustments.
The nonsense with x0 can be improved to remove the concatenation and slicing, which might help. Note that the x_0 term in x_0 + (x_0 + x_1) + (x_0 + x_1 + x_2) can just be written as broadcast addition into all the accumulated terms, so all the concatenation stuff can just be reduced to: i_z_0 + i_z_t.cumsum(axis=0)