Fix invert for _AutoregressiveEmissionsMixin

Addresses issue #76
lindermanlab · Dec 13, 2019 · 0e397ac · 0e397ac
1 parent f66316c
commit 0e397ac
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/ssm/emissions.py b/ssm/emissions.py
@@ -758,8 +758,9 @@ def log_likelihoods(self, data, input, mask, tag, x):
         return np.sum(lls * mask[:, None, :], axis=2)
 
     def invert(self, data, input=None, mask=None, tag=None):
-        pad = np.zeros((1, 1, self.N)) if self.single_subspace else np.zeros((1, self.K, self.N))
-        resid = data - np.concatenate((pad, self.As[None, :, :] * data[:-1, None, :]))
+        assert self.single_subspace, "Can only invert with a single emission model"
+        pad = np.zeros((1, self.N))
+        resid = data - np.concatenate((pad, self.As * data[:-1]))
         return self._invert(resid, input=input, mask=mask, tag=tag)
 
     def sample(self, z, x, input=None, tag=None):