add dispatch for identity Op, use static shapes for parameters

Author: ferrine

pymc/sampling/jax.py

@@ -45,6 +45,7 @@
 from pymc.distributions.multivariate import PosDefMatrix
 from pymc.initial_point import StartDict
 from pymc.logprob.utils import CheckParameterValue
+from pymc.pytensorf import IdentityOp
 from pymc.sampling.mcmc import _init_jitter
 from pymc.util import (
     RandomSeed,
@@ -67,6 +68,14 @@
 )
 
 
+@jax_funcify.register(IdentityOp)
+def jax_funcify_Identity(op, **kwargs):
+    def identity_fn(value):
+        return value
+
+    return identity_fn
+
+
 @jax_funcify.register(Assert)
 @jax_funcify.register(CheckParameterValue)
 def jax_funcify_Assert(op, **kwargs):

pymc/variational/approximations.py

@@ -93,8 +93,8 @@ def create_shared_params(self, start=None, start_sigma=None):
         rho = rho1
 
         return {
-            "mu": pytensor.shared(pm.floatX(start), "mu"),
-            "rho": pytensor.shared(pm.floatX(rho), "rho"),
+            "mu": pytensor.shared(pm.floatX(start), "mu", shape=start.shape),
+            "rho": pytensor.shared(pm.floatX(rho), "rho", shape=rho.shape),
         }
 
     @node_property
@@ -137,7 +137,10 @@ def create_shared_params(self, start=None):
         start = self._prepare_start(start)
         n = self.ddim
         L_tril = np.eye(n)[np.tril_indices(n)].astype(pytensor.config.floatX)
-        return {"mu": pytensor.shared(start, "mu"), "L_tril": pytensor.shared(L_tril, "L_tril")}
+        return {
+            "mu": pytensor.shared(start, "mu", shape=start.shape),
+            "L_tril": pytensor.shared(L_tril, "L_tril", shape=L_tril.shape),
+        }
 
     @node_property
     def L(self):