Dispatching issue around non-linear node (linearization)

[bartvanerp]

The following model does not run, because of incorrect dispatching:

g(x, z) = x .* z

@meta function test_meta()
    g() -> Linearization()
end

# Model Creation
@model function test_model(z, y)
    x ~ NormalMeanVariance(1.0, 1.0)
    u := g(x, z)
    y ~ MvNormalMeanPrecision(u, diageye(2))

end

results = infer(
    model = test_model(),
	data = (z = [1, 2], y = [1, 2]),
    meta = test_meta(),
)

It throws the following error:

DimensionMismatch: jacobian(f, x) expects that x is an array. Perhaps you meant derivative(f, x)?