More static shapes

Author: jessegrabowski

pymc_extras/statespace/models/structural/core.py

@@ -551,6 +551,13 @@ def make_slice(name, x, o_x):
         obs_intercept.name = d.name
 
         transition = pt.linalg.block_diag(T, o_T)
+        transition = pt.specify_shape(
+            transition,
+            shape=[
+                sum(shapes) if not any([s is None for s in shapes]) else None
+                for shapes in zip(*[T.type.shape, o_T.type.shape])
+            ],
+        )
         transition.name = T.name
 
         design = join_tensors_by_dim_labels(
@@ -563,6 +570,13 @@ def make_slice(name, x, o_x):
         design.name = Z.name
 
         selection = pt.linalg.block_diag(R, o_R)
+        selection = pt.specify_shape(
+            selection,
+            shape=[
+                sum(shapes) if not any([s is None for s in shapes]) else None
+                for shapes in zip(*[R.type.shape, o_R.type.shape])
+            ],
+        )
         selection.name = R.name
 
         obs_cov = add_tensors_by_dim_labels(

pymc_extras/statespace/models/utilities.py

@@ -428,7 +428,8 @@ def reorder_from_labels(
     if indices.tolist() != list(range(n_out)):
         for axis in labeled_axis:
             idx = np.s_[tuple([slice(None, None) if i != axis else indices for i in range(x.ndim)])]
-            x = x[idx]
+            shape = x.type.shape
+            x = pt.specify_shape(x[idx], shape)
 
     return x
 
@@ -511,7 +512,11 @@ def ndim_pad_and_reorder(
 
     if n_missing > 0:
         pad_size = [(0, 0) if i not in labeled_axis else (0, n_missing) for i in range(x.ndim)]
-        x = pt.pad(x, pad_size, mode="constant", constant_values=0)
+        new_shape = [
+            shape + sum(size) if shape is not None else None
+            for shape, size in zip(x.type.shape, pad_size)
+        ]
+        x = pt.specify_shape(pt.pad(x, pad_size, mode="constant", constant_values=0), new_shape)
 
     return reorder_from_labels(x, labels, ordered_labels, labeled_axis)