cs = [c for _, c in global_RVs.items()]
shps = [shp for _, _, shp, _ in global_order.vmap]
c_g = [c * tt.ones(np.prod(shp)) for c, shp in zip(cs, shps)]
return inarray_global, uw_global, replace_global, c_g
After Change
// Replace RVs with reshaped subvectors of the joined vector
// The order of global_order is the same with that of global_RVs
subvecs = [reshape_t(inarray_global[slc], shp).astype(dtyp)
for _, slc, shp, dtyp in global_order.vmap]
replace_global = {v: subvec for v, subvec in zip(global_RVs, subvecs)}
// Weight vector
cs = [c for _, c in global_RVs.items()]
oness = [tt.ones(v.ravel().tag.test_value.shape) for v in global_RVs]