Debugging pytensor shape issue

Author: aandorra-mia

pymc_extras/statespace/core/statespace.py

@@ -732,7 +732,7 @@ def _insert_random_variables(self):
 
         replacement_dict = {var: pymc_model[name] for name, var in self._name_to_variable.items()}
         self.subbed_ssm = vectorize_graph(matrices, replace=replacement_dict)
-        for name, matrix in zip(LONG_MATRIX_NAMES, self.subbed_ssm):
+        for name, matrix in zip(MATRIX_NAMES, self.subbed_ssm):
             matrix.name = name
 
     def _insert_data_variables(self):

pymc_extras/statespace/filters/distributions.py

@@ -375,7 +375,6 @@ def dist(cls, mus, covs, logp, **kwargs):
     def rv_op(cls, mus, covs, logp, size=None):
         # Batch dimensions (if any) will be on the far left, but scan requires time to be there instead
         mus_, covs_ = mus.type(), covs.type()
-        print(f"mus_.type.shape: {mus_.type.shape}, covs_.type.shape: {covs_.type.shape}")
 
         logp_ = logp.type()
         rng = pytensor.shared(np.random.default_rng())
@@ -385,7 +384,6 @@ def recursion(mus, covs, rng):
                 mus = pt.moveaxis(mus, -2, 0)
             if covs.ndim > 3:
                 covs = pt.moveaxis(covs, -3, 0)
-            print(f"mus.type.shape: {mus.type.shape}, covs.type.shape: {covs.type.shape}")
 
             def step(mu, cov, rng):
                 new_rng, mvn = pm.MvNormal.dist(mu=mu, cov=cov, rng=rng, method="svd").owner.outputs
@@ -394,32 +392,26 @@ def step(mu, cov, rng):
             mvn_seq, updates = pytensor.scan(
                 step, sequences=[mus, covs], non_sequences=[rng], strict=True, n_steps=mus.shape[0]
             )
-            print(f"mvn_seq.type.shape: {mvn_seq.type.shape}")
             mvn_seq = pt.specify_shape(mvn_seq, mus.type.shape)
 
             # Move time axis back to position -2 so batches are on the left
             if mvn_seq.ndim > 2:
                 mvn_seq = pt.moveaxis(mvn_seq, 0, -2)
-            print(f"mvn_seq.type.shape: {mvn_seq.type.shape}")
 
             (seq_mvn_rng,) = tuple(updates.values())
 
-            print(f"mvn_seq.type.shape: {mvn_seq.type.shape}")
-
             return [seq_mvn_rng, mvn_seq]
 
         mvn_seq_op = KalmanFilterRV(
             inputs=[mus_, covs_, logp_, rng], outputs=recursion(mus_, covs_, rng), ndim_supp=2
         )
 
         mvn_seq = mvn_seq_op(mus, covs, logp, rng)
-        print(f"mvn_seq.type.shape: {mvn_seq.type.shape}")
         return mvn_seq
 
 
 @_logprob.register(KalmanFilterRV)
 def sequence_mvnormal_logp(op, values, mus, covs, logp, rng, **kwargs):
-    print(values[0].type.shape, mus.type.shape, covs.type.shape)
     return check_parameters(
         logp,
         pt.eq(values[0].shape[-2], mus.shape[-2]),