# Buggy code
def transition_coeff(
    self,
    seasonal_indicators: Tensor,  # (batch_size, time_length)
) -> Tensor:
    F = getF(seasonal_indicators)

    _transition_coeff = (
        (F.diag(F.ones(shape=(2,)), k=0) + F.diag(F.ones(shape=(1,)), k=1))
        .expand_dims(axis=0)
        .expand_dims(axis=0)
    )

    # get the right shape: (batch_size, seq_length, latent_dim, latent_dim)
    zeros = _broadcast_param(
        F.zeros_like(
            seasonal_indicators.slice_axis(axis=-1, begin=0, end=1).squeeze(axis=-1)
        ),
        axes=[2, 3],
        sizes=[self.latent_dim(), self.latent_dim()],
    )

    return _transition_coeff.broadcast_like(zeros)


# Traceback (mxnet.base.MXNetError)
Traceback (most recent call last):
File "D:/python/DeepSSM/deep ISSM.py", line 53, in <module>
predictor = estimator.train(dataset.train)
File "D:\python\DeepSSM\gluonts\model\estimator.py", line 254, in train
training_data, validation_data, num_workers, num_prefetch, **kwargs
File "D:\python\DeepSSM\gluonts\model\estimator.py", line 233, in train_model
validation_iter=validation_data_loader,
File "D:\python\DeepSSM\gluonts\trainer\_base.py", line 308, in __call__
epoch_loss = loop(epoch_no, train_iter)
File "D:\python\DeepSSM\gluonts\trainer\_base.py", line 242, in loop
output = net(*inputs)
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\gluon\block.py", line 540, in __call__
out = self.forward(*args)
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\gluon\block.py", line 917, in forward
return self.hybrid_forward(ndarray, x, *args, **params)
File "D:\python\DeepSSM\gluonts\model\deepstate\_network.py", line 193, in hybrid_forward
length=self.past_length,
File "D:\python\DeepSSM\gluonts\model\deepstate\_network.py", line 151, in compute_lds
) = self.issm.get_issm_coeff(seasonal_indicators)
File "D:\python\DeepSSM\gluonts\model\deepstate\issm.py", line 132, in get_issm_coeff
self.innovation_coeff(seasonal_indicators),
File "D:\python\DeepSSM\gluonts\model\deepstate\issm.py", line 247, in innovation_coeff
axis=2
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\ndarray\ndarray.py", line 1695, in squeeze
return op.squeeze(self, *args, **kwargs)
File "<string>", line 45, in squeeze
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\_ctypes\ndarray.py", line 92, in _imperative_invoke
ctypes.byref(out_stypes)))
File "C:\Anaconda\envs\time_series\lib\site-packages\mxnet\base.py", line 252, in check_call
raise MXNetError(py_str(_LIB.MXGetLastError()))
mxnet.base.MXNetError: [13:47:56] c:\jenkins\workspace\mxnet-tag\mxnet\src\operator\tensor\./matrix_op-inl.h:2178: Check failed: dshape[axes[i]] == 1 (2 vs. 1) cannot select an axis to squeeze out which has size=2 not equal to one

# Fixed code
def transition_coeff(
    self,
    feature: Tensor,  # (batch_size, time_length, 1)
) -> Tensor:
    F = getF(feature)

    _transition_coeff = (
        (F.diag(F.ones(shape=(2,)), k=0) + F.diag(F.ones(shape=(1,)), k=1))
        .expand_dims(axis=0)
        .expand_dims(axis=0)
    )

    # get the right shape: (batch_size, time_length, latent_dim, latent_dim)
    zeros = _broadcast_param(
        feature.squeeze(axis=2),
        axes=[2, 3],
        sizes=[self.latent_dim(), self.latent_dim()],
    )

    return _transition_coeff.broadcast_like(zeros)