We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[bug] bart.py example consistently fails with cholesky error with default arguments
I'm just attempting to get started using pyro, and was trying to work through examples when I found that the example listed on the example page under "Multivariate Forecasting" reliably fails for me.

I see that there is a related open issue which is quite old https://github.com/pyro-ppl/pyro/issues/2017 which has some suggestions for the broader problem which seemed promising, but the PR following the approach discussed (https://github.com/pyro-ppl/pyro/pull/2019) was never approved.

It's not impossible that this is due to some environmental factor. However, I created a clean virtualenv to explore pyro. If there's some environmental contributing factor which I am not aware of, please document it, or even better, add a helper method to health-check a given environment.

If you can confirm/reproduce the failure, I would respectfully suggest that either 
- a solution should be pursued
- a working example of a workaround should be exhibited and linked
- or at a minimum non-functioning "examples" should be removed from the examples page

### Issue Description

`examples/contrib/forecast/bart.py` fails with cholesky error when run with default params (no args)

Note that this example attempts to use the `backtest` method, which trains a model several times over different time windows. The first several such windows succeed.

The error appears as follows:

```
Traceback (most recent call last):
  File "examples/contrib/forecast/bart.py", line 180, in <module>
    main(args)
  File "examples/contrib/forecast/bart.py", line 156, in main
    forecaster_options=forecaster_options,
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/evaluate.py", line 205, in backtest
    batch_size=batch_size,
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/forecaster.py", line 361, in __call__
    return super().__call__(data, covariates, num_samples, batch_size)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/forecaster.py", line 390, in forward
    return self.model(data, covariates)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/nn/module.py", line 426, in __call__
    return super().__call__(*args, **kwargs)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/forecaster.py", line 185, in forward
    self.model(zero_data, covariates)
  File "examples/contrib/forecast/bart.py", line 121, in model
    self.predict(noise_model, prediction)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/forecaster.py", line 157, in predict
    noise = pyro.sample("residual", noise_dist)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/primitives.py", line 163, in sample
    apply_stack(msg)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/poutine/runtime.py", line 218, in apply_stack
    default_process_message(msg)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/poutine/runtime.py", line 179, in default_process_message
    msg["value"] = msg["fn"](*msg["args"], **msg["kwargs"])
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/distributions/torch_distribution.py", line 49, in __call__
    if self.has_rsample
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/distributions/hmm.py", line 584, in rsample
    z = _sequential_gaussian_filter_sample(self._init, trans, sample_shape)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/distributions/hmm.py", line 144, in _sequential_gaussian_filter_sample
    contracted = joint.marginalize(left=state_dim)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/ops/gaussian.py", line 244, in marginalize
    P_b = cholesky(P_bb)
  File "/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/ops/tensor_utils.py", line 399, in cholesky
    return torch.linalg.cholesky(x)
RuntimeError: torch.linalg.cholesky: (Batch element 255): The factorization could not be completed because the input is not positive-definite (the leading minor of order 2 is not positive-definite).
```

### Environment

| | value |
|--|--|
| OS | macOS Big Sur (11.5.2) (intell)|
| python version |  3.7.9| 
| pytorch version | 1.10.1|
| pyro version | 1.8.0|

Note, I get the same behavior on linux in docker.

### Code Snippet
copy-pasted the example here: https://pyro.ai/examples/forecast_simple.html / https://github.com/pyro-ppl/pyro/blob/dev/examples/contrib/forecast/bart.py and simply ran:

```
python bart.py
```


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/contrib/forecast/bart.py`
Content:
```
1 # Copyright Contributors to the Pyro project.
2 # SPDX-License-Identifier: Apache-2.0
3 
4 import argparse
5 import logging
6 
7 import numpy as np
8 import torch
9 
10 import pyro
11 import pyro.distributions as dist
12 from pyro.contrib.examples.bart import load_bart_od
13 from pyro.contrib.forecast import ForecastingModel, backtest
14 from pyro.ops.tensor_utils import periodic_cumsum, periodic_repeat
15 
16 logging.getLogger("pyro").setLevel(logging.DEBUG)
17 logging.getLogger("pyro").handlers[0].setLevel(logging.DEBUG)
18 
19 
20 def preprocess(args):
21     """
22     Extract a tensor of (arrivals,departures) to Embarcadero station.
23     """
24     print("Loading data")
25     dataset = load_bart_od()
26 
27     # The full dataset has all station->station ridership counts for all of 50
28     # train stations. In this simple example we will model only the aggretate
29     # counts to and from a single station, Embarcadero.
30     i = dataset["stations"].index("EMBR")
31     arrivals = dataset["counts"][:, :, i].sum(-1)
32     departures = dataset["counts"][:, i, :].sum(-1)
33     data = torch.stack([arrivals, departures], dim=-1)
34 
35     # This simple example uses no covariates, so we will construct a
36     # zero-element tensor of the correct length as empty covariates.
37     covariates = torch.zeros(len(data), 0)
38 
39     return data, covariates
40 
41 
42 # We define a model by subclassing the ForecastingModel class and implementing
43 # a single .model() method.
44 class Model(ForecastingModel):
45     # The .model() method inputs two tensors: a fake tensor zero_data that is
46     # the same size and dtype as the real data (but of course the generative
47     # model shouldn't depend on the value of the data it generates!), and a
48     # tensor of covariates. Our simple model depends on no covariates, so we
49     # simply pass in an empty tensor (see  the preprocess() function above).
50     def model(self, zero_data, covariates):
51         period = 24 * 7
52         duration, dim = zero_data.shape[-2:]
53         assert dim == 2  # Data is bivariate: (arrivals, departures).
54 
55         # Sample global parameters.
56         noise_scale = pyro.sample(
57             "noise_scale", dist.LogNormal(torch.full((dim,), -3.0), 1.0).to_event(1)
58         )
59         assert noise_scale.shape[-1:] == (dim,)
60         trans_timescale = pyro.sample(
61             "trans_timescale", dist.LogNormal(torch.zeros(dim), 1).to_event(1)
62         )
63         assert trans_timescale.shape[-1:] == (dim,)
64 
65         trans_loc = pyro.sample("trans_loc", dist.Cauchy(0, 1 / period))
66         trans_loc = trans_loc.unsqueeze(-1).expand(trans_loc.shape + (dim,))
67         assert trans_loc.shape[-1:] == (dim,)
68         trans_scale = pyro.sample(
69             "trans_scale", dist.LogNormal(torch.zeros(dim), 0.1).to_event(1)
70         )
71         trans_corr = pyro.sample("trans_corr", dist.LKJCholesky(dim, torch.ones(())))
72         trans_scale_tril = trans_scale.unsqueeze(-1) * trans_corr
73         assert trans_scale_tril.shape[-2:] == (dim, dim)
74 
75         obs_scale = pyro.sample(
76             "obs_scale", dist.LogNormal(torch.zeros(dim), 0.1).to_event(1)
77         )
78         obs_corr = pyro.sample("obs_corr", dist.LKJCholesky(dim, torch.ones(())))
79         obs_scale_tril = obs_scale.unsqueeze(-1) * obs_corr
80         assert obs_scale_tril.shape[-2:] == (dim, dim)
81 
82         # Note the initial seasonality should be sampled in a plate with the
83         # same dim as the time_plate, dim=-1. That way we can repeat the dim
84         # below using periodic_repeat().
85         with pyro.plate("season_plate", period, dim=-1):
86             season_init = pyro.sample(
87                 "season_init", dist.Normal(torch.zeros(dim), 1).to_event(1)
88             )
89             assert season_init.shape[-2:] == (period, dim)
90 
91         # Sample independent noise at each time step.
92         with self.time_plate:
93             season_noise = pyro.sample(
94                 "season_noise", dist.Normal(0, noise_scale).to_event(1)
95             )
96             assert season_noise.shape[-2:] == (duration, dim)
97 
98         # Construct a prediction. This prediction has an exactly repeated
99         # seasonal part plus slow seasonal drift. We use two deterministic,
100         # linear functions to transform our diagonal Normal noise to nontrivial
101         # samples from a Gaussian process.
102         prediction = periodic_repeat(season_init, duration, dim=-2) + periodic_cumsum(
103             season_noise, period, dim=-2
104         )
105         assert prediction.shape[-2:] == (duration, dim)
106 
107         # Construct a joint noise model. This model is a GaussianHMM, whose
108         # .rsample() and .log_prob() methods are parallelized over time; this
109         # this entire model is parallelized over time.
110         init_dist = dist.Normal(torch.zeros(dim), 100).to_event(1)
111         trans_mat = trans_timescale.neg().exp().diag_embed()
112         trans_dist = dist.MultivariateNormal(trans_loc, scale_tril=trans_scale_tril)
113         obs_mat = torch.eye(dim)
114         obs_dist = dist.MultivariateNormal(torch.zeros(dim), scale_tril=obs_scale_tril)
115         noise_model = dist.GaussianHMM(
116             init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=duration
117         )
118         assert noise_model.event_shape == (duration, dim)
119 
120         # The final statement registers our noise model and prediction.
121         self.predict(noise_model, prediction)
122 
123 
124 def main(args):
125     data, covariates = preprocess(args)
126 
127     # We will model positive count data by log1p-transforming it into real
128     # valued data.  But since we want to evaluate back in the count domain, we
129     # will also define a transform to apply during evaluation, transforming
130     # from real back to count-valued data. Truth is mapped by the log1p()
131     # inverse expm1(), but the prediction will be sampled from a Poisson
132     # distribution.
133     data = data.log1p()
134 
135     def transform(pred, truth):
136         pred = torch.poisson(pred.clamp(min=1e-4).expm1())
137         truth = truth.expm1()
138         return pred, truth
139 
140     # The backtest() function automatically trains and evaluates our model on
141     # different windows of data.
142     forecaster_options = {
143         "num_steps": args.num_steps,
144         "learning_rate": args.learning_rate,
145         "log_every": args.log_every,
146         "dct_gradients": args.dct,
147     }
148     metrics = backtest(
149         data,
150         covariates,
151         Model,
152         train_window=args.train_window,
153         test_window=args.test_window,
154         stride=args.stride,
155         num_samples=args.num_samples,
156         forecaster_options=forecaster_options,
157     )
158 
159     for name in ["mae", "rmse", "crps"]:
160         values = [m[name] for m in metrics]
161         mean = np.mean(values)
162         std = np.std(values)
163         print("{} = {:0.3g} +- {:0.3g}".format(name, mean, std))
164     return metrics
165 
166 
167 if __name__ == "__main__":
168     assert pyro.__version__.startswith("1.8.0")
169     parser = argparse.ArgumentParser(description="Bart Ridership Forecasting Example")
170     parser.add_argument("--train-window", default=2160, type=int)
171     parser.add_argument("--test-window", default=336, type=int)
172     parser.add_argument("--stride", default=168, type=int)
173     parser.add_argument("-n", "--num-steps", default=501, type=int)
174     parser.add_argument("-lr", "--learning-rate", default=0.05, type=float)
175     parser.add_argument("--dct", action="store_true")
176     parser.add_argument("--num-samples", default=100, type=int)
177     parser.add_argument("--log-every", default=50, type=int)
178     parser.add_argument("--seed", default=1234567890, type=int)
179     args = parser.parse_args()
180     main(args)
181
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/contrib/forecast/bart.py b/examples/contrib/forecast/bart.py
--- a/examples/contrib/forecast/bart.py
+++ b/examples/contrib/forecast/bart.py
@@ -171,7 +171,7 @@
     parser.add_argument("--test-window", default=336, type=int)
     parser.add_argument("--stride", default=168, type=int)
     parser.add_argument("-n", "--num-steps", default=501, type=int)
-    parser.add_argument("-lr", "--learning-rate", default=0.05, type=float)
+    parser.add_argument("-lr", "--learning-rate", default=0.01, type=float)
     parser.add_argument("--dct", action="store_true")
     parser.add_argument("--num-samples", default=100, type=int)
     parser.add_argument("--log-every", default=50, type=int)

{"golden_diff": "diff --git a/examples/contrib/forecast/bart.py b/examples/contrib/forecast/bart.py\n--- a/examples/contrib/forecast/bart.py\n+++ b/examples/contrib/forecast/bart.py\n@@ -171,7 +171,7 @@\n     parser.add_argument(\"--test-window\", default=336, type=int)\n     parser.add_argument(\"--stride\", default=168, type=int)\n     parser.add_argument(\"-n\", \"--num-steps\", default=501, type=int)\n-    parser.add_argument(\"-lr\", \"--learning-rate\", default=0.05, type=float)\n+    parser.add_argument(\"-lr\", \"--learning-rate\", default=0.01, type=float)\n     parser.add_argument(\"--dct\", action=\"store_true\")\n     parser.add_argument(\"--num-samples\", default=100, type=int)\n     parser.add_argument(\"--log-every\", default=50, type=int)\n", "issue": "[bug] bart.py example consistently fails with cholesky error with default arguments\nI'm just attempting to get started using pyro, and was trying to work through examples when I found that the example listed on the example page under \"Multivariate Forecasting\" reliably fails for me.\r\n\r\nI see that there is a related open issue which is quite old https://github.com/pyro-ppl/pyro/issues/2017 which has some suggestions for the broader problem which seemed promising, but the PR following the approach discussed (https://github.com/pyro-ppl/pyro/pull/2019) was never approved.\r\n\r\nIt's not impossible that this is due to some environmental factor. However, I created a clean virtualenv to explore pyro. If there's some environmental contributing factor which I am not aware of, please document it, or even better, add a helper method to health-check a given environment.\r\n\r\nIf you can confirm/reproduce the failure, I would respectfully suggest that either \r\n- a solution should be pursued\r\n- a working example of a workaround should be exhibited and linked\r\n- or at a minimum non-functioning \"examples\" should be removed from the examples page\r\n\r\n### Issue Description\r\n\r\n`examples/contrib/forecast/bart.py` fails with cholesky error when run with default params (no args)\r\n\r\nNote that this example attempts to use the `backtest` method, which trains a model several times over different time windows. The first several such windows succeed.\r\n\r\nThe error appears as follows:\r\n\r\n```\r\nTraceback (most recent call last):\r\n  File \"examples/contrib/forecast/bart.py\", line 180, in <module>\r\n    main(args)\r\n  File \"examples/contrib/forecast/bart.py\", line 156, in main\r\n    forecaster_options=forecaster_options,\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/evaluate.py\", line 205, in backtest\r\n    batch_size=batch_size,\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/forecaster.py\", line 361, in __call__\r\n    return super().__call__(data, covariates, num_samples, batch_size)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/torch/nn/modules/module.py\", line 1102, in _call_impl\r\n    return forward_call(*input, **kwargs)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/forecaster.py\", line 390, in forward\r\n    return self.model(data, covariates)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/nn/module.py\", line 426, in __call__\r\n    return super().__call__(*args, **kwargs)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/torch/nn/modules/module.py\", line 1102, in _call_impl\r\n    return forward_call(*input, **kwargs)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/forecaster.py\", line 185, in forward\r\n    self.model(zero_data, covariates)\r\n  File \"examples/contrib/forecast/bart.py\", line 121, in model\r\n    self.predict(noise_model, prediction)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/contrib/forecast/forecaster.py\", line 157, in predict\r\n    noise = pyro.sample(\"residual\", noise_dist)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/primitives.py\", line 163, in sample\r\n    apply_stack(msg)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/poutine/runtime.py\", line 218, in apply_stack\r\n    default_process_message(msg)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/poutine/runtime.py\", line 179, in default_process_message\r\n    msg[\"value\"] = msg[\"fn\"](*msg[\"args\"], **msg[\"kwargs\"])\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/distributions/torch_distribution.py\", line 49, in __call__\r\n    if self.has_rsample\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/distributions/hmm.py\", line 584, in rsample\r\n    z = _sequential_gaussian_filter_sample(self._init, trans, sample_shape)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/distributions/hmm.py\", line 144, in _sequential_gaussian_filter_sample\r\n    contracted = joint.marginalize(left=state_dim)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/ops/gaussian.py\", line 244, in marginalize\r\n    P_b = cholesky(P_bb)\r\n  File \"/Users/aaron/.pyenv/versions/vpyro/lib/python3.7/site-packages/pyro/ops/tensor_utils.py\", line 399, in cholesky\r\n    return torch.linalg.cholesky(x)\r\nRuntimeError: torch.linalg.cholesky: (Batch element 255): The factorization could not be completed because the input is not positive-definite (the leading minor of order 2 is not positive-definite).\r\n```\r\n\r\n### Environment\r\n\r\n| | value |\r\n|--|--|\r\n| OS | macOS Big Sur (11.5.2) (intell)|\r\n| python version |  3.7.9| \r\n| pytorch version | 1.10.1|\r\n| pyro version | 1.8.0|\r\n\r\nNote, I get the same behavior on linux in docker.\r\n\r\n### Code Snippet\r\ncopy-pasted the example here: https://pyro.ai/examples/forecast_simple.html / https://github.com/pyro-ppl/pyro/blob/dev/examples/contrib/forecast/bart.py and simply ran:\r\n\r\n```\r\npython bart.py\r\n```\r\n\n", "before_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nimport argparse\nimport logging\n\nimport numpy as np\nimport torch\n\nimport pyro\nimport pyro.distributions as dist\nfrom pyro.contrib.examples.bart import load_bart_od\nfrom pyro.contrib.forecast import ForecastingModel, backtest\nfrom pyro.ops.tensor_utils import periodic_cumsum, periodic_repeat\n\nlogging.getLogger(\"pyro\").setLevel(logging.DEBUG)\nlogging.getLogger(\"pyro\").handlers[0].setLevel(logging.DEBUG)\n\n\ndef preprocess(args):\n    \"\"\"\n    Extract a tensor of (arrivals,departures) to Embarcadero station.\n    \"\"\"\n    print(\"Loading data\")\n    dataset = load_bart_od()\n\n    # The full dataset has all station->station ridership counts for all of 50\n    # train stations. In this simple example we will model only the aggretate\n    # counts to and from a single station, Embarcadero.\n    i = dataset[\"stations\"].index(\"EMBR\")\n    arrivals = dataset[\"counts\"][:, :, i].sum(-1)\n    departures = dataset[\"counts\"][:, i, :].sum(-1)\n    data = torch.stack([arrivals, departures], dim=-1)\n\n    # This simple example uses no covariates, so we will construct a\n    # zero-element tensor of the correct length as empty covariates.\n    covariates = torch.zeros(len(data), 0)\n\n    return data, covariates\n\n\n# We define a model by subclassing the ForecastingModel class and implementing\n# a single .model() method.\nclass Model(ForecastingModel):\n    # The .model() method inputs two tensors: a fake tensor zero_data that is\n    # the same size and dtype as the real data (but of course the generative\n    # model shouldn't depend on the value of the data it generates!), and a\n    # tensor of covariates. Our simple model depends on no covariates, so we\n    # simply pass in an empty tensor (see  the preprocess() function above).\n    def model(self, zero_data, covariates):\n        period = 24 * 7\n        duration, dim = zero_data.shape[-2:]\n        assert dim == 2  # Data is bivariate: (arrivals, departures).\n\n        # Sample global parameters.\n        noise_scale = pyro.sample(\n            \"noise_scale\", dist.LogNormal(torch.full((dim,), -3.0), 1.0).to_event(1)\n        )\n        assert noise_scale.shape[-1:] == (dim,)\n        trans_timescale = pyro.sample(\n            \"trans_timescale\", dist.LogNormal(torch.zeros(dim), 1).to_event(1)\n        )\n        assert trans_timescale.shape[-1:] == (dim,)\n\n        trans_loc = pyro.sample(\"trans_loc\", dist.Cauchy(0, 1 / period))\n        trans_loc = trans_loc.unsqueeze(-1).expand(trans_loc.shape + (dim,))\n        assert trans_loc.shape[-1:] == (dim,)\n        trans_scale = pyro.sample(\n            \"trans_scale\", dist.LogNormal(torch.zeros(dim), 0.1).to_event(1)\n        )\n        trans_corr = pyro.sample(\"trans_corr\", dist.LKJCholesky(dim, torch.ones(())))\n        trans_scale_tril = trans_scale.unsqueeze(-1) * trans_corr\n        assert trans_scale_tril.shape[-2:] == (dim, dim)\n\n        obs_scale = pyro.sample(\n            \"obs_scale\", dist.LogNormal(torch.zeros(dim), 0.1).to_event(1)\n        )\n        obs_corr = pyro.sample(\"obs_corr\", dist.LKJCholesky(dim, torch.ones(())))\n        obs_scale_tril = obs_scale.unsqueeze(-1) * obs_corr\n        assert obs_scale_tril.shape[-2:] == (dim, dim)\n\n        # Note the initial seasonality should be sampled in a plate with the\n        # same dim as the time_plate, dim=-1. That way we can repeat the dim\n        # below using periodic_repeat().\n        with pyro.plate(\"season_plate\", period, dim=-1):\n            season_init = pyro.sample(\n                \"season_init\", dist.Normal(torch.zeros(dim), 1).to_event(1)\n            )\n            assert season_init.shape[-2:] == (period, dim)\n\n        # Sample independent noise at each time step.\n        with self.time_plate:\n            season_noise = pyro.sample(\n                \"season_noise\", dist.Normal(0, noise_scale).to_event(1)\n            )\n            assert season_noise.shape[-2:] == (duration, dim)\n\n        # Construct a prediction. This prediction has an exactly repeated\n        # seasonal part plus slow seasonal drift. We use two deterministic,\n        # linear functions to transform our diagonal Normal noise to nontrivial\n        # samples from a Gaussian process.\n        prediction = periodic_repeat(season_init, duration, dim=-2) + periodic_cumsum(\n            season_noise, period, dim=-2\n        )\n        assert prediction.shape[-2:] == (duration, dim)\n\n        # Construct a joint noise model. This model is a GaussianHMM, whose\n        # .rsample() and .log_prob() methods are parallelized over time; this\n        # this entire model is parallelized over time.\n        init_dist = dist.Normal(torch.zeros(dim), 100).to_event(1)\n        trans_mat = trans_timescale.neg().exp().diag_embed()\n        trans_dist = dist.MultivariateNormal(trans_loc, scale_tril=trans_scale_tril)\n        obs_mat = torch.eye(dim)\n        obs_dist = dist.MultivariateNormal(torch.zeros(dim), scale_tril=obs_scale_tril)\n        noise_model = dist.GaussianHMM(\n            init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=duration\n        )\n        assert noise_model.event_shape == (duration, dim)\n\n        # The final statement registers our noise model and prediction.\n        self.predict(noise_model, prediction)\n\n\ndef main(args):\n    data, covariates = preprocess(args)\n\n    # We will model positive count data by log1p-transforming it into real\n    # valued data.  But since we want to evaluate back in the count domain, we\n    # will also define a transform to apply during evaluation, transforming\n    # from real back to count-valued data. Truth is mapped by the log1p()\n    # inverse expm1(), but the prediction will be sampled from a Poisson\n    # distribution.\n    data = data.log1p()\n\n    def transform(pred, truth):\n        pred = torch.poisson(pred.clamp(min=1e-4).expm1())\n        truth = truth.expm1()\n        return pred, truth\n\n    # The backtest() function automatically trains and evaluates our model on\n    # different windows of data.\n    forecaster_options = {\n        \"num_steps\": args.num_steps,\n        \"learning_rate\": args.learning_rate,\n        \"log_every\": args.log_every,\n        \"dct_gradients\": args.dct,\n    }\n    metrics = backtest(\n        data,\n        covariates,\n        Model,\n        train_window=args.train_window,\n        test_window=args.test_window,\n        stride=args.stride,\n        num_samples=args.num_samples,\n        forecaster_options=forecaster_options,\n    )\n\n    for name in [\"mae\", \"rmse\", \"crps\"]:\n        values = [m[name] for m in metrics]\n        mean = np.mean(values)\n        std = np.std(values)\n        print(\"{} = {:0.3g} +- {:0.3g}\".format(name, mean, std))\n    return metrics\n\n\nif __name__ == \"__main__\":\n    assert pyro.__version__.startswith(\"1.8.0\")\n    parser = argparse.ArgumentParser(description=\"Bart Ridership Forecasting Example\")\n    parser.add_argument(\"--train-window\", default=2160, type=int)\n    parser.add_argument(\"--test-window\", default=336, type=int)\n    parser.add_argument(\"--stride\", default=168, type=int)\n    parser.add_argument(\"-n\", \"--num-steps\", default=501, type=int)\n    parser.add_argument(\"-lr\", \"--learning-rate\", default=0.05, type=float)\n    parser.add_argument(\"--dct\", action=\"store_true\")\n    parser.add_argument(\"--num-samples\", default=100, type=int)\n    parser.add_argument(\"--log-every\", default=50, type=int)\n    parser.add_argument(\"--seed\", default=1234567890, type=int)\n    args = parser.parse_args()\n    main(args)\n", "path": "examples/contrib/forecast/bart.py"}], "after_files": [{"content": "# Copyright Contributors to the Pyro project.\n# SPDX-License-Identifier: Apache-2.0\n\nimport argparse\nimport logging\n\nimport numpy as np\nimport torch\n\nimport pyro\nimport pyro.distributions as dist\nfrom pyro.contrib.examples.bart import load_bart_od\nfrom pyro.contrib.forecast import ForecastingModel, backtest\nfrom pyro.ops.tensor_utils import periodic_cumsum, periodic_repeat\n\nlogging.getLogger(\"pyro\").setLevel(logging.DEBUG)\nlogging.getLogger(\"pyro\").handlers[0].setLevel(logging.DEBUG)\n\n\ndef preprocess(args):\n    \"\"\"\n    Extract a tensor of (arrivals,departures) to Embarcadero station.\n    \"\"\"\n    print(\"Loading data\")\n    dataset = load_bart_od()\n\n    # The full dataset has all station->station ridership counts for all of 50\n    # train stations. In this simple example we will model only the aggretate\n    # counts to and from a single station, Embarcadero.\n    i = dataset[\"stations\"].index(\"EMBR\")\n    arrivals = dataset[\"counts\"][:, :, i].sum(-1)\n    departures = dataset[\"counts\"][:, i, :].sum(-1)\n    data = torch.stack([arrivals, departures], dim=-1)\n\n    # This simple example uses no covariates, so we will construct a\n    # zero-element tensor of the correct length as empty covariates.\n    covariates = torch.zeros(len(data), 0)\n\n    return data, covariates\n\n\n# We define a model by subclassing the ForecastingModel class and implementing\n# a single .model() method.\nclass Model(ForecastingModel):\n    # The .model() method inputs two tensors: a fake tensor zero_data that is\n    # the same size and dtype as the real data (but of course the generative\n    # model shouldn't depend on the value of the data it generates!), and a\n    # tensor of covariates. Our simple model depends on no covariates, so we\n    # simply pass in an empty tensor (see  the preprocess() function above).\n    def model(self, zero_data, covariates):\n        period = 24 * 7\n        duration, dim = zero_data.shape[-2:]\n        assert dim == 2  # Data is bivariate: (arrivals, departures).\n\n        # Sample global parameters.\n        noise_scale = pyro.sample(\n            \"noise_scale\", dist.LogNormal(torch.full((dim,), -3.0), 1.0).to_event(1)\n        )\n        assert noise_scale.shape[-1:] == (dim,)\n        trans_timescale = pyro.sample(\n            \"trans_timescale\", dist.LogNormal(torch.zeros(dim), 1).to_event(1)\n        )\n        assert trans_timescale.shape[-1:] == (dim,)\n\n        trans_loc = pyro.sample(\"trans_loc\", dist.Cauchy(0, 1 / period))\n        trans_loc = trans_loc.unsqueeze(-1).expand(trans_loc.shape + (dim,))\n        assert trans_loc.shape[-1:] == (dim,)\n        trans_scale = pyro.sample(\n            \"trans_scale\", dist.LogNormal(torch.zeros(dim), 0.1).to_event(1)\n        )\n        trans_corr = pyro.sample(\"trans_corr\", dist.LKJCholesky(dim, torch.ones(())))\n        trans_scale_tril = trans_scale.unsqueeze(-1) * trans_corr\n        assert trans_scale_tril.shape[-2:] == (dim, dim)\n\n        obs_scale = pyro.sample(\n            \"obs_scale\", dist.LogNormal(torch.zeros(dim), 0.1).to_event(1)\n        )\n        obs_corr = pyro.sample(\"obs_corr\", dist.LKJCholesky(dim, torch.ones(())))\n        obs_scale_tril = obs_scale.unsqueeze(-1) * obs_corr\n        assert obs_scale_tril.shape[-2:] == (dim, dim)\n\n        # Note the initial seasonality should be sampled in a plate with the\n        # same dim as the time_plate, dim=-1. That way we can repeat the dim\n        # below using periodic_repeat().\n        with pyro.plate(\"season_plate\", period, dim=-1):\n            season_init = pyro.sample(\n                \"season_init\", dist.Normal(torch.zeros(dim), 1).to_event(1)\n            )\n            assert season_init.shape[-2:] == (period, dim)\n\n        # Sample independent noise at each time step.\n        with self.time_plate:\n            season_noise = pyro.sample(\n                \"season_noise\", dist.Normal(0, noise_scale).to_event(1)\n            )\n            assert season_noise.shape[-2:] == (duration, dim)\n\n        # Construct a prediction. This prediction has an exactly repeated\n        # seasonal part plus slow seasonal drift. We use two deterministic,\n        # linear functions to transform our diagonal Normal noise to nontrivial\n        # samples from a Gaussian process.\n        prediction = periodic_repeat(season_init, duration, dim=-2) + periodic_cumsum(\n            season_noise, period, dim=-2\n        )\n        assert prediction.shape[-2:] == (duration, dim)\n\n        # Construct a joint noise model. This model is a GaussianHMM, whose\n        # .rsample() and .log_prob() methods are parallelized over time; this\n        # this entire model is parallelized over time.\n        init_dist = dist.Normal(torch.zeros(dim), 100).to_event(1)\n        trans_mat = trans_timescale.neg().exp().diag_embed()\n        trans_dist = dist.MultivariateNormal(trans_loc, scale_tril=trans_scale_tril)\n        obs_mat = torch.eye(dim)\n        obs_dist = dist.MultivariateNormal(torch.zeros(dim), scale_tril=obs_scale_tril)\n        noise_model = dist.GaussianHMM(\n            init_dist, trans_mat, trans_dist, obs_mat, obs_dist, duration=duration\n        )\n        assert noise_model.event_shape == (duration, dim)\n\n        # The final statement registers our noise model and prediction.\n        self.predict(noise_model, prediction)\n\n\ndef main(args):\n    data, covariates = preprocess(args)\n\n    # We will model positive count data by log1p-transforming it into real\n    # valued data.  But since we want to evaluate back in the count domain, we\n    # will also define a transform to apply during evaluation, transforming\n    # from real back to count-valued data. Truth is mapped by the log1p()\n    # inverse expm1(), but the prediction will be sampled from a Poisson\n    # distribution.\n    data = data.log1p()\n\n    def transform(pred, truth):\n        pred = torch.poisson(pred.clamp(min=1e-4).expm1())\n        truth = truth.expm1()\n        return pred, truth\n\n    # The backtest() function automatically trains and evaluates our model on\n    # different windows of data.\n    forecaster_options = {\n        \"num_steps\": args.num_steps,\n        \"learning_rate\": args.learning_rate,\n        \"log_every\": args.log_every,\n        \"dct_gradients\": args.dct,\n    }\n    metrics = backtest(\n        data,\n        covariates,\n        Model,\n        train_window=args.train_window,\n        test_window=args.test_window,\n        stride=args.stride,\n        num_samples=args.num_samples,\n        forecaster_options=forecaster_options,\n    )\n\n    for name in [\"mae\", \"rmse\", \"crps\"]:\n        values = [m[name] for m in metrics]\n        mean = np.mean(values)\n        std = np.std(values)\n        print(\"{} = {:0.3g} +- {:0.3g}\".format(name, mean, std))\n    return metrics\n\n\nif __name__ == \"__main__\":\n    assert pyro.__version__.startswith(\"1.8.0\")\n    parser = argparse.ArgumentParser(description=\"Bart Ridership Forecasting Example\")\n    parser.add_argument(\"--train-window\", default=2160, type=int)\n    parser.add_argument(\"--test-window\", default=336, type=int)\n    parser.add_argument(\"--stride\", default=168, type=int)\n    parser.add_argument(\"-n\", \"--num-steps\", default=501, type=int)\n    parser.add_argument(\"-lr\", \"--learning-rate\", default=0.01, type=float)\n    parser.add_argument(\"--dct\", action=\"store_true\")\n    parser.add_argument(\"--num-samples\", default=100, type=int)\n    parser.add_argument(\"--log-every\", default=50, type=int)\n    parser.add_argument(\"--seed\", default=1234567890, type=int)\n    args = parser.parse_args()\n    main(args)\n", "path": "examples/contrib/forecast/bart.py"}]}