Repository: awslabs/gluonts

Issue:
Issue in  DeepstateEstimator for "prediction length=1 " 
I am using DeepStateEstimator for forecasting purposes with Hourly data .But if I set prediction length=1,means prediction for next hour ,it gives me error and if I set prediction length>=2,its working fine.Please find the error below for prediction length=1.

If I run
``` 
forecasts = list(forecast_it)
tss = list(ts_it)
```
then it gives me error
```
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:

```



Discussion:
This is a minimal working example raising the bug:

```python
import pandas as pd
import mxnet as mx

from gluonts.dataset.common import ListDataset
from gluonts.model.deepstate import DeepStateEstimator
from gluonts.mx.trainer import Trainer

prediction_length = 1
freq = "1h"

train_ds = ListDataset(
    [{'target': [10.0] * 2000, 'start': pd.Timestamp("2020-01-01 00:00:00", freq=freq)}],
    freq=freq
)

estimator = DeepStateEstimator(
    prediction_length=prediction_length,
    cardinality=[1],
    use_feat_static_cat=False,
    freq=freq,
    trainer=Trainer(
        ctx=mx.context.cpu(),
        epochs=1,
        learning_rate=1E-3,
        hybridize=True,
        num_batches_per_epoch=2
    ),
)

predictor = estimator.train(train_ds)

predictions = list(predictor.predict(train_ds))
```

Resolving patch:
diff --git a/src/gluonts/mx/distribution/lds.py b/src/gluonts/mx/distribution/lds.py
--- a/src/gluonts/mx/distribution/lds.py
+++ b/src/gluonts/mx/distribution/lds.py
@@ -39,6 +39,41 @@ def __init__(self, lower, upper) -> None:
         self.upper = upper
 
 
+def _safe_split(x, num_outputs, axis, squeeze_axis, *args, **kwargs):
+    """
+    A type-stable wrapper around mx.nd.split.
+
+    Currently mx.nd.split behaves weirdly if num_outputs=1:
+
+        a = mx.nd.ones(shape=(1, 1, 2))
+        l = a.split(axis=1, num_outputs=1, squeeze_axis=False)
+        type(l)  # mx.NDArray
+        l.shape  # (1, 1, 2)
+
+    Compare that with the case num_outputs=2:
+
+        a = mx.nd.ones(shape=(1, 2, 2))
+        l = a.split(axis=1, num_outputs=2, squeeze_axis=False)
+        type(l)  # list
+        len(l)  # 2
+        l[0].shape  # (1, 1, 2)
+
+    This wrapper makes the behavior consistent by always returning a list
+    of length num_outputs, whose elements will have one less axis than x
+    in case x.shape[axis]==num_outputs and squeeze_axis==True, and the same
+    number of axes as x otherwise.
+    """
+    if num_outputs > 1:
+        return x.split(
+            axis=axis,
+            num_outputs=num_outputs,
+            squeeze_axis=squeeze_axis,
+            *args,
+            **kwargs
+        )
+    return [x.squeeze(axis=axis)] if squeeze_axis else [x]
+
+
 class LDS(Distribution):
     r"""
     Implements Linear Dynamical System (LDS) as a distribution.
@@ -111,28 +146,37 @@ def __init__(
 
         # 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
+        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 = innovation_coeff.split(
-            axis=1, num_outputs=self.seq_length, squeeze_axis=False
+        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 = transition_coeff.split(
-            axis=1, num_outputs=self.seq_length, squeeze_axis=True
+        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 = noise_std.split(
-            axis=1, num_outputs=self.seq_length, squeeze_axis=True
+        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 = residuals.split(
-            axis=1, num_outputs=self.seq_length, squeeze_axis=True
+        self.residuals = _safe_split(
+            residuals, axis=1, num_outputs=self.seq_length, squeeze_axis=True
         )
 
         self.prior_mean = prior_mean
@@ -219,8 +263,8 @@ def kalman_filter(
         """
         F = self.F
         # targets[t]: (batch_size, obs_dim)
-        targets = targets.split(
-            axis=1, num_outputs=self.seq_length, squeeze_axis=True
+        targets = _safe_split(
+            targets, axis=1, num_outputs=self.seq_length, squeeze_axis=True
         )
 
         log_p_seq = []
@@ -229,8 +273,11 @@ def kalman_filter(
         cov = self.prior_cov
 
         observed = (
-            observed.split(
-                axis=1, num_outputs=self.seq_length, squeeze_axis=True
+            _safe_split(
+                observed,
+                axis=1,
+                num_outputs=self.seq_length,
+                squeeze_axis=True,
             )
             if observed is not None
             else None
@@ -319,18 +366,22 @@ def sample(
         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)
+        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 = (
-            Gaussian(noise_std.zeros_like(), noise_std.ones_like())
-            .sample(num_samples)
-            .split(axis=-3, num_outputs=self.seq_length, squeeze_axis=True)
+        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.


Test patch:
diff --git a/test/model/deepstate/test_deepstate_smoke.py b/test/model/deepstate/test_deepstate_smoke.py
--- a/test/model/deepstate/test_deepstate_smoke.py
+++ b/test/model/deepstate/test_deepstate_smoke.py
@@ -20,7 +20,7 @@
 # First-party imports
 from gluonts.model.deepstate import DeepStateEstimator
 from gluonts.testutil.dummy_datasets import make_dummy_datasets_with_features
-from gluonts.trainer import Trainer
+from gluonts.mx.trainer import Trainer
 
 
 common_estimator_hps = dict(