The problems occurs only at prediction. I have no issues sampling from the prior, posterior, or intermediate random variables. This is because at prediction the shape of X changes, but since I allocate output based on self.N and R, I am unable to update inplace
output*pm.math.matmul(feature_test(X[:,d],self.N,self.M),factors[d])
as output and pm.math.matmul(feature_test(X[:,d],self.N,self.M),factors[d]) have different shapes: output is self.N x R, while pm.math.matmul(feature_test(X[:,d],self.N,self.M),factors[d]) is new_shape_of_X x R. All I need is to correctly update self.N which is the shape of ‘X’ when predicting. But I don’t know what the best practice is to do so.