# Buggy code
def sample(
    self, num_samples: Optional[int] = None, scale: Optional[Tensor] = None
) -> Tensor:
    r"""
    Generates samples from the LDS: p(z_1, z_2, \ldots, z_{`seq_length`}).

    Parameters
    ----------
    num_samples
        Number of samples to generate
    scale
        Scale of each sequence in x, shape (batch_size, output_dim)

    Returns
    -------
    Tensor
        Samples, shape (num_samples, batch_size, seq_length, output_dim)
    """
    F = self.F

    # Note on shapes: here we work with tensors of the following shape
    # in each time step t: (num_samples, batch_size, dim, dim),
    # where dim can be obs_dim or latent_dim or a constant 1 to facilitate
    # generalized matrix multiplication (gemm2)

    # Sample observation noise for all time steps
    # noise_std: (batch_size, seq_length, obs_dim, 1)
    noise_std = F.stack(*self.noise_std, axis=1).expand_dims(axis=-1)

    # samples_eps_obs[t]: (num_samples, batch_size, obs_dim, 1)
    samples_eps_obs = (
        Gaussian(noise_std.zeros_like(), noise_std)
        .sample(num_samples)
        .split(axis=-3, num_outputs=self.seq_length, squeeze_axis=True)
    )

    # Sample standard normal for all time steps
    # samples_eps_std_normal[t]: (num_samples, batch_size, obs_dim, 1)
    samples_std_normal = (
        Gaussian(noise_std.zeros_like(), noise_std.ones_like())
        .sample(num_samples)
        .split(axis=-3, num_outputs=self.seq_length, squeeze_axis=True)
    )

    # Sample the prior state.
    # samples_lat_state: (num_samples, batch_size, latent_dim, 1)
    # The prior covariance is observed to be slightly negative definite whenever there is
    # excessive zero padding at the beginning of the time series.
    # We add positive tolerance to the diagonal to avoid numerical issues.
    # Note that `jitter_cholesky` adds positive tolerance only if the decomposition without jitter fails.
    state = MultivariateGaussian(
        self.prior_mean,
        jitter_cholesky(F, self.prior_cov, self.latent_dim, float_type=np.float32),
    )
    samples_lat_state = state.sample(num_samples).expand_dims(axis=-1)

    samples_seq = []
    for t in range(self.seq_length):
        # Expand all coefficients to include samples in axis 0
        # emission_coeff_t: (num_samples, batch_size, obs_dim, latent_dim)
        # transition_coeff_t:
        #   (num_samples, batch_size, latent_dim, latent_dim)
        # innovation_coeff_t: (num_samples, batch_size, 1, latent_dim)
        emission_coeff_t, transition_coeff_t, innovation_coeff_t = [
            _broadcast_param(coeff, axes=[0], sizes=[num_samples])
            if num_samples is not None
            else coeff
            for coeff in [
                self.emission_coeff[t],
                self.transition_coeff[t],
                self.innovation_coeff[t],
            ]
        ]

        # Expand residuals as well
        # residual_t: (num_samples, batch_size, obs_dim, 1)
        residual_t = (
            _broadcast_param(
                self.residuals[t].expand_dims(axis=-1),
                axes=[0],
                sizes=[num_samples],
            )
            if num_samples is not None
            else self.residuals[t].expand_dims(axis=-1)
        )

        # (num_samples, batch_size, 1, obs_dim)
        samples_t = (
            F.linalg_gemm2(emission_coeff_t, samples_lat_state)
            + residual_t
            + samples_eps_obs[t]
        )
        samples_t = (
            samples_t.swapaxes(dim1=2, dim2=3)
            if num_samples is not None
            else samples_t.swapaxes(dim1=1, dim2=2)
        )
        samples_seq.append(samples_t)

        # sample next state: (num_samples, batch_size, latent_dim, 1)
        samples_lat_state = F.linalg_gemm2(
            transition_coeff_t, samples_lat_state
        ) + F.linalg_gemm2(innovation_coeff_t, samples_std_normal[t], transpose_a=True)

    # (num_samples, batch_size, seq_length, obs_dim)
    samples = F.concat(*samples_seq, dim=-2)
    return (
        samples
        if scale is None
        else F.broadcast_mul(
            samples,
            scale.expand_dims(axis=1).expand_dims(axis=0)
            if num_samples is not None
            else scale.expand_dims(axis=1),
        )
    )


# 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 sample(
    self, num_samples: Optional[int] = None, scale: Optional[Tensor] = None
) -> Tensor:
    r"""
    Generates samples from the LDS: p(z_1, z_2, \ldots, z_{`seq_length`}).

    Parameters
    ----------
    num_samples
        Number of samples to generate
    scale
        Scale of each sequence in x, shape (batch_size, output_dim)

    Returns
    -------
    Tensor
        Samples, shape (num_samples, batch_size, seq_length, output_dim)
    """
    F = self.F

    # Note on shapes: here we work with tensors of the following shape
    # in each time step t: (num_samples, batch_size, dim, dim),
    # where dim can be obs_dim or latent_dim or a constant 1 to facilitate
    # generalized matrix multiplication (gemm2)

    # Sample observation noise for all time steps
    # noise_std: (batch_size, seq_length, obs_dim, 1)
    noise_std = F.stack(*self.noise_std, axis=1).expand_dims(axis=-1)

    # samples_eps_obs[t]: (num_samples, batch_size, obs_dim, 1)
    samples_eps_obs = _safe_split(
        Gaussian(noise_std.zeros_like(), noise_std).sample(num_samples),
        axis=-3,
        num_outputs=self.seq_length,
        squeeze_axis=True,
    )

    # Sample standard normal for all time steps
    # samples_eps_std_normal[t]: (num_samples, batch_size, obs_dim, 1)
    samples_std_normal = _safe_split(
        Gaussian(noise_std.zeros_like(), noise_std.ones_like()).sample(num_samples),
        axis=-3,
        num_outputs=self.seq_length,
        squeeze_axis=True,
    )

    # Sample the prior state.
    # samples_lat_state: (num_samples, batch_size, latent_dim, 1)
    # The prior covariance is observed to be slightly negative definite whenever there is
    # excessive zero padding at the beginning of the time series.
    # We add positive tolerance to the diagonal to avoid numerical issues.
    # Note that `jitter_cholesky` adds positive tolerance only if the decomposition without jitter fails.
    state = MultivariateGaussian(
        self.prior_mean,
        jitter_cholesky(F, self.prior_cov, self.latent_dim, float_type=np.float32),
    )
    samples_lat_state = state.sample(num_samples).expand_dims(axis=-1)

    samples_seq = []
    for t in range(self.seq_length):
        # Expand all coefficients to include samples in axis 0
        # emission_coeff_t: (num_samples, batch_size, obs_dim, latent_dim)
        # transition_coeff_t:
        #   (num_samples, batch_size, latent_dim, latent_dim)
        # innovation_coeff_t: (num_samples, batch_size, 1, latent_dim)
        emission_coeff_t, transition_coeff_t, innovation_coeff_t = [
            _broadcast_param(coeff, axes=[0], sizes=[num_samples])
            if num_samples is not None
            else coeff
            for coeff in [
                self.emission_coeff[t],
                self.transition_coeff[t],
                self.innovation_coeff[t],
            ]
        ]

        # Expand residuals as well
        # residual_t: (num_samples, batch_size, obs_dim, 1)
        residual_t = (
            _broadcast_param(
                self.residuals[t].expand_dims(axis=-1),
                axes=[0],
                sizes=[num_samples],
            )
            if num_samples is not None
            else self.residuals[t].expand_dims(axis=-1)
        )

        # (num_samples, batch_size, 1, obs_dim)
        samples_t = (
            F.linalg_gemm2(emission_coeff_t, samples_lat_state)
            + residual_t
            + samples_eps_obs[t]
        )
        samples_t = (
            samples_t.swapaxes(dim1=2, dim2=3)
            if num_samples is not None
            else samples_t.swapaxes(dim1=1, dim2=2)
        )
        samples_seq.append(samples_t)

        # sample next state: (num_samples, batch_size, latent_dim, 1)
        samples_lat_state = F.linalg_gemm2(
            transition_coeff_t, samples_lat_state
        ) + F.linalg_gemm2(innovation_coeff_t, samples_std_normal[t], transpose_a=True)

    # (num_samples, batch_size, seq_length, obs_dim)
    samples = F.concat(*samples_seq, dim=-2)
    return (
        samples
        if scale is None
        else F.broadcast_mul(
            samples,
            scale.expand_dims(axis=1).expand_dims(axis=0)
            if num_samples is not None
            else scale.expand_dims(axis=1),
        )
    )