fix trend model to work in symbolic mode (#578)

* fix trend model to work in symbolic mode * explicitly enable hybridization
awslabs · Jan 30, 2020 · 42f512a · 42f512a
1 parent 6693927
commit 42f512a
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/gluonts/model/deepstate/issm.py b/src/gluonts/model/deepstate/issm.py
@@ -201,8 +201,7 @@ def transition_coeff(
  F = getF(seasonal_indicators)
 
  _transition_coeff = (
- # TODO: this won't work if F == mx.sym
- F.array([[1, 1], [0, 1]])
+ (F.diag(F.ones(shape=(2,)), k=0) + F.diag(F.ones(shape=(1,)), k=1))
  .expand_dims(axis=0)
  .expand_dims(axis=0)
  )

diff --git a/test/model/test_deepstate_smoke.py b/test/model/test_deepstate_smoke.py
@@ -26,8 +26,11 @@
 common_estimator_hps = dict(
  freq="D",
  prediction_length=3,
- trainer=Trainer(epochs=3, num_batches_per_epoch=2, batch_size=1),
+ trainer=Trainer(
+ epochs=3, num_batches_per_epoch=2, batch_size=1, hybridize=True
+ ),
  past_length=10,
+ add_trend=True,
 )