Fix fgraph_from_model with multivariate transformed variables (#6924)

* fix do with transformed

* Change test

---------

Co-authored-by: Ricardo Vieira <ricardo.vieira1994@gmail.com>

Author: ferrine

pymc/model/fgraph.py

@@ -70,7 +70,7 @@ def make_node(self, rv, value, *dims):
         dims = self._parse_dims(rv, *dims)
         if value is not None:
             assert isinstance(value, Variable)
-            assert rv.type.in_same_class(value.type)
+            assert rv.type.dtype == value.type.dtype
             return Apply(self, [rv, value, *dims], [rv.type(name=rv.name)])
 
 

tests/model/test_fgraph.py

@@ -351,3 +351,16 @@ def test_fgraph_rewrite(non_centered_rewrite):
         m_new.compile_logp()(ip),
         m_ref.compile_logp()(ip),
     )
+
+
+def test_multivariate_transform():
+    with pm.Model() as m:
+        x = pm.Dirichlet("x", a=[1, 1, 1])
+        y, *_ = pm.LKJCholeskyCov("y", n=4, eta=1, sd_dist=pm.Exponential.dist(1))
+
+    new_m = clone_model(m)
+
+    ip = m.initial_point()
+    new_ip = new_m.initial_point()
+    np.testing.assert_allclose(ip["x_simplex__"], new_ip["x_simplex__"])
+    np.testing.assert_allclose(ip["y_cholesky-cov-packed__"], new_ip["y_cholesky-cov-packed__"])