I see, try y_pred = pm.Poisson("y_pred", mu=tt.square(f_pred) + 1e-6, shape=test_x.shape)
y_pred = pm.Poisson("y_pred", mu=tt.square(f_pred) + 1e-6, shape=test_x.shape)