Potential mask and single_subspace issue with Gaussian Emissions: "assert value.shape == (1, N, D) if self.single_subspace else (K, N, D)"

[smestern]

I was running some hyperparameter scans, and ran into a weird issue. When fitting an SLDS with Gaussian emission, a single subspace, and masks, the m-step fails because it produces a non-single-subspace.
Stack Trace:

  File "c:\users\smest\source\ssm\ssm\lds.py", line 790, in fit
    elbos = _fitting_methods[method](
  File "c:\users\smest\source\ssm\ssm\lds.py", line 699, in _fit_laplace_em
    self._fit_laplace_em_params_update(
  File "c:\users\smest\source\ssm\ssm\lds.py", line 620, in _fit_laplace_em_params_update
    self.emissions.m_step(discrete_expectations, continuous_samples,
  File "c:\users\smest\source\ssm\ssm\emissions.py", line 438, in m_step
    self.Cs = np.array(Cs)
  File "c:\users\smest\source\ssm\ssm\emissions.py", line 127, in Cs
    assert value.shape == (1, N, D) if self.single_subspace else (K, N, D)
AssertionError

Minimal Code Example:

import numpy as np
import ssm
data = np.random.randn(100, 10)  # 100 time steps, 10 observed dimensions
masks = np.ones_like(data, dtype=bool) 
#set a random subset of the data to be missing
masks[20:30, 0:5] = False  # Mask out some data points
inputs = np.zeros((100, 0))  # No inputs in this example

#does not fail
slds = ssm.SLDS(N=10, K=3, D=2, M=0, single_subspace=True, emissions="gaussian_orthog")
slds.initialize(data, inputs=inputs, masks=masks, verbose=0)
q_elbos, q = slds.fit(data, inputs=inputs, masks=masks, verbose=1, method="laplace_em", num_iters=100, num_samples=10, initialize=False)


#fails
slds = ssm.SLDS(N=10, K=3, D=2, M=0, single_subspace=True, emissions="gaussian")
slds.initialize(data, inputs=inputs, masks=masks, verbose=0)
q_elbos, q = slds.fit(data, inputs=inputs, masks=masks, verbose=1, method="laplace_em", num_iters=100, num_samples=10, initialize=False)

The issue does not happen when the mask is all true. It appears the issue happens around line 438 in emissions.py

      if self.single_subspace and all([np.all(mask) for mask in masks]): #<--Has single subspace but fails np.all(mask)
          # Return exact m-step updates for C, F, d, and inv_etas
          CF, d, Sigma = fit_linear_regression(
              Xs, ys,
              prior_ExxT=1e-4 * np.eye(self.D + self.M + 1),
              prior_ExyT=np.zeros((self.D + self.M + 1, self.N)))
          self.Cs = CF[None, :, :self.D]
          self.Fs = CF[None, :, self.D:]
          self.ds = d[None, :]
          self.inv_etas = np.log(np.diag(Sigma))[None, :]
      else:
          Cs, Fs, ds, inv_etas = [], [], [], []
          for k in range(self.K):
              CF, d, Sigma = fit_linear_regression(
                  Xs, ys, weights=[w[:, k] for w in ws],
                  prior_ExxT=1e-4 * np.eye(self.D + self.M + 1),
                  prior_ExyT=np.zeros((self.D + self.M + 1, self.N)))
              Cs.append(CF[:, :self.D])
              Fs.append(CF[:, self.D:])
              ds.append(d)
              inv_etas.append(np.log(np.diag(Sigma)))
          self.Cs = np.array(Cs) # <--- Attempts to assign a non-single subspace emission matrix in the shape of (K, N, D)
          self.Fs = np.array(Fs)
          self.ds = np.array(ds)
          self.inv_etas = np.array(inv_etas)

I am not sure if this is intended.