class_docstr: Abstract base class for forecasting models.
  
  Derived classes must implement the :meth:`model` method.
class_docstr_tok: ["Abstract","base","class","for","forecasting","models",".","Derived","classes","must","implement","the",":","meth",":","`","model","`","method","."]
docstr: Prediction function, to be called by :meth:`model` implementations.
  
  This should be called outside of the  :meth:`time_plate`.
  
  This is similar to an observe statement in Pyro::
  
      pyro.sample("residual", noise_dist,
                  obs=(data - prediction))
  
  but with (1) additional reshaping logic to allow time-dependent
  ``noise_dist`` (most often a :class:`~pyro.distributions.GaussianHMM`
  or variant); and (2) additional logic to allow only a partial
  observation and forecast the remaining data.
  
  :param noise_dist: A noise distribution with ``.event_dim in {0,1,2}``.
      ``noise_dist`` is typically zero-mean or zero-median or zero-mode
      or somehow centered.
  :type noise_dist: ~pyro.distributions.Distribution
  :param prediction: A prediction for the data. This should have the same
      shape as ``data``, but broadcastable to full duration of the
      ``covariates``.
  :type prediction: ~torch.Tensor
docstr_tok: ["Prediction","function",",","to","be","called","by",":","meth",":","`","model","`","implementations",".","This","should","be","called","outside","of","the",":","meth",":","`","time_plate","`",".","This","is","similar","to","an","observe","statement","in","Pyro",":",":","pyro.sample","(","``","residual","''",",","noise_dist",",","obs=","(","data","-","prediction",")",")","but","with","(","1",")","additional","reshaping","logic","to","allow","time-dependent","``","noise_dist","``","(","most","often","a",":","class",":","`","~pyro.distributions.GaussianHMM","`","or","variant",")",";","and","(","2",")","additional","logic","to","allow","only","a","partial","observation","and","forecast","the","remaining","data",".",":","param","noise_dist",":","A","noise","distribution","with","``",".event_dim","in","{","0,1,2","}","``",".","``","noise_dist","``","is","typically","zero-mean","or","zero-median","or","zero-mode","or","somehow","centered",".",":","type","noise_dist",":","~pyro.distributions.Distribution",":","param","prediction",":","A","prediction","for","the","data",".","This","should","have","the","same","shape","as","``","data","``",",","but","broadcastable","to","full","duration","of","the","``","covariates","``",".",":","type","prediction",":","~torch.Tensor"]
code: def predict(self, noise_dist, prediction):
      """
      Prediction function, to be called by :meth:`model` implementations.
  
      This should be called outside of the  :meth:`time_plate`.
  
      This is similar to an observe statement in Pyro::
  
          pyro.sample("residual", noise_dist,
                      obs=(data - prediction))
  
      but with (1) additional reshaping logic to allow time-dependent
      ``noise_dist`` (most often a :class:`~pyro.distributions.GaussianHMM`
      or variant); and (2) additional logic to allow only a partial
      observation and forecast the remaining data.
  
      :param noise_dist: A noise distribution with ``.event_dim in {0,1,2}``.
          ``noise_dist`` is typically zero-mean or zero-median or zero-mode
          or somehow centered.
      :type noise_dist: ~pyro.distributions.Distribution
      :param prediction: A prediction for the data. This should have the same
          shape as ``data``, but broadcastable to full duration of the
          ``covariates``.
      :type prediction: ~torch.Tensor
      """
      assert (
          self._data is not None
      ), ".predict() called outside .model()"
      assert self._forecast is None, ".predict() called twice"
      assert isinstance(noise_dist, dist.Distribution)
      assert isinstance(prediction, torch.Tensor)
      if noise_dist.event_dim == 0:
          if noise_dist.batch_shape[-2:] != prediction.shape[-2:]:
              noise_dist = noise_dist.expand(
                  noise_dist.batch_shape[:-2] + prediction.shape[-2:]
              )
          noise_dist = noise_dist.to_event(2)
      elif noise_dist.event_dim == 1:
          if noise_dist.batch_shape[-1:] != prediction.shape[-2:-1]:
              noise_dist = noise_dist.expand(
                  noise_dist.batch_shape[:-1] + prediction.shape[-2:-1]
              )
          noise_dist = noise_dist.to_event(1)
      assert noise_dist.event_dim == 2
      assert noise_dist.event_shape == prediction.shape[-2:]
  
      # The following reshaping logic is required to reconcile batch and
      # event shapes. This would be unnecessary if Pyro used name dimensions
      # internally, e.g. using Funsor.
      #
      #     batch_shape                    | event_shape
      #     -------------------------------+----------------
      #  1. sample_shape + shape + (time,) | (obs_dim,)
      #  2.           sample_shape + shape | (time, obs_dim)
      #  3.    sample_shape + shape + (1,) | (time, obs_dim)
      #
      # Parameters like noise_dist.loc typically have shape as in 1.  However
      # calling .to_event(1) will shift the shapes resulting in 2., where
      # sample_shape+shape will be misaligned with other batch shapes in the
      # trace. To fix this the following logic "unsqueezes" the distribution,
      # resulting in correctly aligned shapes 3. Note the "time" dimension is
      # effectively moved from a batch dimension to an event dimension.
      noise_dist = reshape_batch(
          noise_dist, noise_dist.batch_shape + (1,)
      )
      data = pyro.subsample(self._data.unsqueeze(-3), event_dim=2)
      prediction = prediction.unsqueeze(-3)
  
      # Create a sample site.
      t_obs = data.size(-2)
      t_cov = prediction.size(-2)
      if t_obs == t_cov:  # training
          pyro.sample("residual", noise_dist, obs=data - prediction)
          self._forecast = data.new_zeros(
              data.shape[:-2] + (0,) + data.shape[-1:]
          )
      else:  # forecasting
          left_pred = prediction[..., :t_obs, :]
          right_pred = prediction[..., t_obs:, :]
  
          # This prefix_condition indirection is needed to ensure that
          # PrefixConditionMessenger is handled outside of the .model() call.
          self._prefix_condition_data["residual"] = data - left_pred
          noise = pyro.sample("residual", noise_dist)
          del self._prefix_condition_data["residual"]
  
          assert (
              noise.shape[-data.dim() :]
              == right_pred.shape[-data.dim() :]
          )
          self._forecast = right_pred + noise
  
      # Move the "time" batch dim back to its original place.
      assert self._forecast.size(-3) == 1
      self._forecast = self._forecast.squeeze(-3)
  
imports: import logging
  from abc import ABCMeta, abstractmethod
  from contextlib import ExitStack
  import torch
  import torch.nn
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.infer import MCMC, NUTS, SVI, Trace_ELBO
  from pyro.infer.autoguide import AutoNormal, init_to_sample
  from pyro.infer.predictive import _guess_max_plate_nesting
  from pyro.nn.module import PyroModule
  from pyro.optim import DCTAdam
  from .util import MarkDCTParamMessenger, PrefixConditionMessenger, PrefixReplayMessenger, PrefixWarmStartMessenger, reshape_batch, time_reparam_dct, time_reparam_haar
label_desc: Use image node_id 4 for calling the ForecastingModel obj's underlying member method code with example usage: obj.predict(noise_dist, prediction) without return types