mmmm interesting 
so the reason why is basically because I’m doing CV prediction scores, and multiples predictions using different `obs_dataset’.
to be more precise my problem is a little bit more complicated. the shape of the variables mu and sigma is determined by some mutable_coords (there is a time coordinate that I use to do time series prediction). when doing cross-validation those coordinates change, and therefore the shape of mu and sigma change accordingly.
These variables, mu and sigma, are used in other parts of the model where the change in shape is needed, however, in this part, the variable “likelihood” that its a formula of the original mu, its only needed when fitting because the obs dataset its using just once, after that for the predictions I actually don’t need this value, however, because the mask is linked to the X variable (with mutable coords) makes a mess because you can’t index mu[mask] if mask is a variable.
does this makes sense ?
(thanks in advance for helping me!! 
)