# Buggy code
def kalman_filter(self, targets: Tensor, observed: Tensor) -> Tuple[Tensor, ...]:
    """
    Performs Kalman filtering given observations.


    Parameters
    ----------
    targets
        Observations, shape (batch_size, seq_length, output_dim)
    observed
        Flag tensor indicating which observations are genuine (1.0) and
        which are missing (0.0)

    Returns
    -------
    Tensor
        Log probabilities, shape (batch_size, seq_length)
    Tensor
        Mean of p(l_T | l_{T-1}), where T is seq_length, with shape
        (batch_size, latent_dim)
    Tensor
        Covariance of p(l_T | l_{T-1}), where T is seq_length, with shape
        (batch_size, latent_dim, latent_dim)
    """
    F = self.F
    # targets[t]: (batch_size, obs_dim)
    targets = targets.split(axis=1, num_outputs=self.seq_length, squeeze_axis=True)

    log_p_seq = []

    mean = self.prior_mean
    cov = self.prior_cov

    observed = (
        observed.split(axis=1, num_outputs=self.seq_length, squeeze_axis=True)
        if observed is not None
        else None
    )

    for t in range(self.seq_length):
        # Compute the filtered distribution
        #   p(l_t | z_1, ..., z_{t + 1})
        # and log - probability
        #   log p(z_t | z_0, z_{t - 1})
        filtered_mean, filtered_cov, log_p = kalman_filter_step(
            F,
            target=targets[t],
            prior_mean=mean,
            prior_cov=cov,
            emission_coeff=self.emission_coeff[t],
            residual=self.residuals[t],
            noise_std=self.noise_std[t],
            latent_dim=self.latent_dim,
            output_dim=self.output_dim,
        )

        log_p_seq.append(log_p.expand_dims(axis=1))

        # Mean of p(l_{t+1} | l_t)
        mean = F.linalg_gemm2(
            self.transition_coeff[t],
            (
                filtered_mean.expand_dims(axis=-1)
                if observed is None
                else F.where(observed[t], x=filtered_mean, y=mean).expand_dims(axis=-1)
            ),
        ).squeeze(axis=-1)

        # Covariance of p(l_{t+1} | l_t)
        cov = F.linalg_gemm2(
            self.transition_coeff[t],
            F.linalg_gemm2(
                (
                    filtered_cov
                    if observed is None
                    else F.where(observed[t], x=filtered_cov, y=cov)
                ),
                self.transition_coeff[t],
                transpose_b=True,
            ),
        ) + F.linalg_gemm2(
            self.innovation_coeff[t],
            self.innovation_coeff[t],
            transpose_a=True,
        )

    # Return sequence of log likelihoods, as well as
    # final mean and covariance of p(l_T | l_{T-1} where T is seq_length
    return F.concat(*log_p_seq, dim=1), mean, 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 kalman_filter(self, targets: Tensor, observed: Tensor) -> Tuple[Tensor, ...]:
    """
    Performs Kalman filtering given observations.


    Parameters
    ----------
    targets
        Observations, shape (batch_size, seq_length, output_dim)
    observed
        Flag tensor indicating which observations are genuine (1.0) and
        which are missing (0.0)

    Returns
    -------
    Tensor
        Log probabilities, shape (batch_size, seq_length)
    Tensor
        Mean of p(l_T | l_{T-1}), where T is seq_length, with shape
        (batch_size, latent_dim)
    Tensor
        Covariance of p(l_T | l_{T-1}), where T is seq_length, with shape
        (batch_size, latent_dim, latent_dim)
    """
    F = self.F
    # targets[t]: (batch_size, obs_dim)
    targets = _safe_split(
        targets, axis=1, num_outputs=self.seq_length, squeeze_axis=True
    )

    log_p_seq = []

    mean = self.prior_mean
    cov = self.prior_cov

    observed = (
        _safe_split(
            observed,
            axis=1,
            num_outputs=self.seq_length,
            squeeze_axis=True,
        )
        if observed is not None
        else None
    )

    for t in range(self.seq_length):
        # Compute the filtered distribution
        #   p(l_t | z_1, ..., z_{t + 1})
        # and log - probability
        #   log p(z_t | z_0, z_{t - 1})
        filtered_mean, filtered_cov, log_p = kalman_filter_step(
            F,
            target=targets[t],
            prior_mean=mean,
            prior_cov=cov,
            emission_coeff=self.emission_coeff[t],
            residual=self.residuals[t],
            noise_std=self.noise_std[t],
            latent_dim=self.latent_dim,
            output_dim=self.output_dim,
        )

        log_p_seq.append(log_p.expand_dims(axis=1))

        # Mean of p(l_{t+1} | l_t)
        mean = F.linalg_gemm2(
            self.transition_coeff[t],
            (
                filtered_mean.expand_dims(axis=-1)
                if observed is None
                else F.where(observed[t], x=filtered_mean, y=mean).expand_dims(axis=-1)
            ),
        ).squeeze(axis=-1)

        # Covariance of p(l_{t+1} | l_t)
        cov = F.linalg_gemm2(
            self.transition_coeff[t],
            F.linalg_gemm2(
                (
                    filtered_cov
                    if observed is None
                    else F.where(observed[t], x=filtered_cov, y=cov)
                ),
                self.transition_coeff[t],
                transpose_b=True,
            ),
        ) + F.linalg_gemm2(
            self.innovation_coeff[t],
            self.innovation_coeff[t],
            transpose_a=True,
        )

    # Return sequence of log likelihoods, as well as
    # final mean and covariance of p(l_T | l_{T-1} where T is seq_length
    return F.concat(*log_p_seq, dim=1), mean, cov