# Buggy code
def __init__(
    self,
    emission_coeff: Tensor,
    transition_coeff: Tensor,
    innovation_coeff: Tensor,
    noise_std: Tensor,
    residuals: Tensor,
    prior_mean: Tensor,
    prior_cov: Tensor,
    latent_dim: int,
    output_dim: int,
    seq_length: int,
) -> None:
    self.latent_dim = latent_dim
    self.output_dim = output_dim
    self.seq_length = seq_length

    # Split coefficients along time axis for easy access
    # emission_coef[t]: (batch_size, obs_dim, latent_dim)
    self.emission_coeff = emission_coeff.split(
        axis=1, num_outputs=self.seq_length, squeeze_axis=True
    )

    # innovation_coef[t]: (batch_size, latent_dim)
    self.innovation_coeff = innovation_coeff.split(
        axis=1, num_outputs=self.seq_length, squeeze_axis=False
    )

    # transition_coeff: (batch_size, latent_dim, latent_dim)
    self.transition_coeff = transition_coeff.split(
        axis=1, num_outputs=self.seq_length, squeeze_axis=True
    )

    # noise_std[t]: (batch_size, obs_dim)
    self.noise_std = noise_std.split(
        axis=1, num_outputs=self.seq_length, squeeze_axis=True
    )

    # residuals[t]: (batch_size, obs_dim)
    self.residuals = residuals.split(
        axis=1, num_outputs=self.seq_length, squeeze_axis=True
    )

    self.prior_mean = prior_mean
    self.prior_cov = prior_cov


# Traceback (MXNetError)
MXNetError                                Traceback (most recent call last)
<ipython-input-38-19e7d512cf80> in <module>
----> 1 forecasts = list(forecast_it)
2 tss = list(ts_it)

c:\users\appdata\local\programs\python\python37\lib\site-packages\gluonts\model\predictor.py in predict(self, dataset, num_samples, num_workers, num_prefetch, **kwargs)
340             freq=self.freq,
341             output_transform=self.output_transform,
--> 342             num_samples=num_samples,
343         )
344

c:\users\appdata\local\programs\python\python37\lib\site-packages\gluonts\model\forecast_generator.py in __call__(self, inference_data_loader, prediction_net, input_names, freq, output_transform, num_samples, **kwargs)
195         for batch in inference_data_loader:
196             inputs = [batch[k] for k in input_names]
--> 197             outputs = prediction_net(*inputs).asnumpy()
198             if output_transform is not None:
199                 outputs = output_transform(batch, outputs)

c:\users\appdata\local\programs\python\python37\lib\site-packages\mxnet\gluon\block.py in __call__(self, *args)
691             hook(self, args)
692
--> 693         out = self.forward(*args)
694
695         for hook in self._forward_hooks.values():

c:\users\appdata\local\programs\python\python37\lib\site-packages\mxnet\gluon\block.py in forward(self, x, *args)
1156                     params = {k: v.data(ctx) for k, v in self._reg_params.items()}
1157
-> 1158                 return self.hybrid_forward(ndarray, x, *args, **params)
1159
1160         params = {i: j.var() for i, j in self._reg_params.items()}

c:\users\appdata\local\programs\python\python37\lib\site-packages\gluonts\model\deepstate\_network.py in hybrid_forward(self, F, feat_static_cat, past_observed_values, past_seasonal_indicators, past_time_feat, past_target, future_seasonal_indicators, future_time_feat)
263
264         samples = lds_prediction.sample(
--> 265             num_samples=self.num_parallel_samples, scale=scale
266         )
267

c:\users\appdata\local\programs\python\python37\lib\site-packages\gluonts\mx\distribution\lds.py in sample(self, num_samples, scale)
382                 F.linalg_gemm2(emission_coeff_t, samples_lat_state)
383                 + residual_t
--> 384                 + samples_eps_obs[t]
385             )
386             samples_t = (

c:\users\appdata\local\programs\python\python37\lib\site-packages\mxnet\ndarray\register.py in linalg_gemm2(A, B, transpose_a, transpose_b, alpha, axis, out, name, **kwargs)

c:\users\appdata\local\programs\python\python37\lib\site-packages\mxnet\_ctypes\ndarray.py in _imperative_invoke(handle, ndargs, keys, vals, out, is_np_op)
105         c_str_array(keys),
106         c_str_array([str(s) for s in vals]),
--> 107         ctypes.byref(out_stypes)))
108
109     create_ndarray_fn = _np_ndarray_cls if is_np_op else _ndarray_cls

c:\users\appdata\local\programs\python\python37\lib\site-packages\mxnet\base.py in check_call(ret)
253     """
254     if ret != 0:
--> 255         raise MXNetError(py_str(_LIB.MXGetLastError()))
256
257

MXNetError: [22:14:24] c:\jenkins\workspace\mxnet-tag\mxnet\src\imperative\./imperative_utils.h:146: Operator _linalg_gemm2 inferring shapes failed.
input shapes:
[100,1,33]
[100,1,33,1]
output shapes:
None
operator attributes:

# Fixed code
def __init__(
    self,
    emission_coeff: Tensor,
    transition_coeff: Tensor,
    innovation_coeff: Tensor,
    noise_std: Tensor,
    residuals: Tensor,
    prior_mean: Tensor,
    prior_cov: Tensor,
    latent_dim: int,
    output_dim: int,
    seq_length: int,
) -> None:
    self.latent_dim = latent_dim
    self.output_dim = output_dim
    self.seq_length = seq_length

    # Split coefficients along time axis for easy access
    # emission_coef[t]: (batch_size, obs_dim, latent_dim)
    self.emission_coeff = _safe_split(
        emission_coeff,
        axis=1,
        num_outputs=self.seq_length,
        squeeze_axis=True,
    )

    # innovation_coef[t]: (batch_size, latent_dim)
    self.innovation_coeff = _safe_split(
        innovation_coeff,
        axis=1,
        num_outputs=self.seq_length,
        squeeze_axis=False,
    )

    # transition_coeff: (batch_size, latent_dim, latent_dim)
    self.transition_coeff = _safe_split(
        transition_coeff,
        axis=1,
        num_outputs=self.seq_length,
        squeeze_axis=True,
    )

    # noise_std[t]: (batch_size, obs_dim)
    self.noise_std = _safe_split(
        noise_std, axis=1, num_outputs=self.seq_length, squeeze_axis=True
    )

    # residuals[t]: (batch_size, obs_dim)
    self.residuals = _safe_split(
        residuals, axis=1, num_outputs=self.seq_length, squeeze_axis=True
    )

    self.prior_mean = prior_mean
    self.prior_cov = prior_cov