class_docstr: Represents an autoregressive integrated moving-average (ARIMA) model.
  
  An [autoregressive moving-average (ARMA)](
  https://en.wikipedia.org/wiki/Autoregressive_moving_average_model) process is
  defined by the recursion
  
  ```
  level[t + 1] = (
      level_drift
      + noise[t + 1]
      + sum(ar_coefficients * levels[t : t - order : -1])
      + sum(ma_coefficients * noise[t : t - order : -1]))
    noise[t + 1] ~ Normal(0., scale=level_scale)
    ```
  
  where `noise` is an iid noise process. An integrated ([ARIMA](
  https://en.wikipedia.org/wiki/Autoregressive_integrated_moving_average))
  process corresponds to an ARMA model of the
  `integration_degree`th-order differences of a sequence, or equivalently,
  taking `integration_degree` cumulative sums of an underlying ARMA process.
docstr: Specifies an ARIMA(p=ar_order, d=integration_degree, q=ma_order) model.
  
  Args:
    ar_order: scalar Python positive `int` specifying the order of the
      autoregressive process (`p` in `ARIMA(p, d, q)`).
    ma_order: scalar Python positive `int` specifying the order of the
      moving-average process (`q` in `ARIMA(p, d, q)`).
    integration_degree: scalar Python positive `int` specifying the number
      of times to integrate an ARMA process. (`d` in `ARIMA(p, d, q)`).
      Default value: `0`.
    ar_coefficients_prior: optional `tfd.Distribution` instance specifying a
      prior on the `ar_coefficients` parameter. If `None`, a default standard
      normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([ar_order]))`)
      prior is used.
      Default value: `None`.
    ma_coefficients_prior: optional `tfd.Distribution` instance specifying a
      prior on the `ma_coefficients` parameter. If `None`, a default standard
      normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([ma_order]))`)
      prior is used.
      Default value: `None`.
    level_drift_prior: optional `tfd.Distribution` instance specifying a prior
      on the `level_drift` parameter. If `None`, the parameter is not inferred
      and is instead fixed to zero.
      Default value: `None`.
    level_scale_prior: optional `tfd.Distribution` instance specifying a prior
      on the `level_scale` parameter. If `None`, a heuristic default prior is
      constructed based on the provided `observed_time_series`.
      Default value: `None`.
    initial_state_prior: optional `tfd.Distribution` instance specifying a
      prior on the initial state, corresponding to the values of the process
      at a set of size `order` of imagined timesteps before the initial step.
      If `None`, a heuristic default prior is constructed based on the
      provided `observed_time_series`.
      Default value: `None`.
    ar_coefficient_constraining_bijector: optional `tfb.Bijector` instance
      representing a constraining mapping for the autoregressive coefficients.
      For example, `tfb.Tanh()` constrains the coefficients to lie in
      `(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
      `tfb.Identity()` implies no constraint. If `None`, the default behavior
      constrains the coefficients to lie in `(-1, 1)` using a `Tanh` bijector.
      Default value: `None`.
    ma_coefficient_constraining_bijector: optional `tfb.Bijector` instance
      representing a constraining mapping for the moving average coefficients.
      For example, `tfb.Tanh()` constrains the coefficients to lie in
      `(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
      `tfb.Identity()` implies no constraint. If `None`, the default behavior
      is to apply no constraint.
      Default value: `None`.
    observed_time_series: optional `float` `Tensor` of shape
      `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
      supported when `T > 1`), specifying an observed time series. Any `NaN`s
      are interpreted as missing observations; missingness may be also be
      explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
      Any priors not explicitly set will be given default values according to
      the scale of the observed time series (or batch of time series).
      Default value: `None`.
    name: the name of this model component.
      Default value: 'ARIMA'.
docstr_tok: ["Specifies","an","ARIMA","(","p=ar_order",",","d=integration_degree",",","q=ma_order",")","model",".","Args",":","ar_order",":","scalar","Python","positive","`","int","`","specifying","the","order","of","the","autoregressive","process","(","`","p","`","in","`","ARIMA","(","p",",","d",",","q",")","`",")",".","ma_order",":","scalar","Python","positive","`","int","`","specifying","the","order","of","the","moving-average","process","(","`","q","`","in","`","ARIMA","(","p",",","d",",","q",")","`",")",".","integration_degree",":","scalar","Python","positive","`","int","`","specifying","the","number","of","times","to","integrate","an","ARMA","process",".","(","`","d","`","in","`","ARIMA","(","p",",","d",",","q",")","`",")",".","Default","value",":","`","0","`",".","ar_coefficients_prior",":","optional","`","tfd.Distribution","`","instance","specifying","a","prior","on","the","`","ar_coefficients","`","parameter",".","If","`","None","`",",","a","default","standard","normal","(","`","tfd.MultivariateNormalDiag","(","scale_diag=tf.ones","(","[","ar_order","]",")",")","`",")","prior","is","used",".","Default","value",":","`","None","`",".","ma_coefficients_prior",":","optional","`","tfd.Distribution","`","instance","specifying","a","prior","on","the","`","ma_coefficients","`","parameter",".","If","`","None","`",",","a","default","standard","normal","(","`","tfd.MultivariateNormalDiag","(","scale_diag=tf.ones","(","[","ma_order","]",")",")","`",")","prior","is","used",".","Default","value",":","`","None","`",".","level_drift_prior",":","optional","`","tfd.Distribution","`","instance","specifying","a","prior","on","the","`","level_drift","`","parameter",".","If","`","None","`",",","the","parameter","is","not","inferred","and","is","instead","fixed","to","zero",".","Default","value",":","`","None","`",".","level_scale_prior",":","optional","`","tfd.Distribution","`","instance","specifying","a","prior","on","the","`","level_scale","`","parameter",".","If","`","None","`",",","a","heuristic","default","prior","is","constructed","based","on","the","provided","`","observed_time_series","`",".","Default","value",":","`","None","`",".","initial_state_prior",":","optional","`","tfd.Distribution","`","instance","specifying","a","prior","on","the","initial","state",",","corresponding","to","the","values","of","the","process","at","a","set","of","size","`","order","`","of","imagined","timesteps","before","the","initial","step",".","If","`","None","`",",","a","heuristic","default","prior","is","constructed","based","on","the","provided","`","observed_time_series","`",".","Default","value",":","`","None","`",".","ar_coefficient_constraining_bijector",":","optional","`","tfb.Bijector","`","instance","representing","a","constraining","mapping","for","the","autoregressive","coefficients",".","For","example",",","`","tfb.Tanh","(",")","`","constrains","the","coefficients","to","lie","in","`","(","-1",",","1",")","`",",","while","`","tfb.Softplus","(",")","`","constrains","them","to","be","positive",",","and","`","tfb.Identity","(",")","`","implies","no","constraint",".","If","`","None","`",",","the","default","behavior","constrains","the","coefficients","to","lie","in","`","(","-1",",","1",")","`","using","a","`","Tanh","`","bijector",".","Default","value",":","`","None","`",".","ma_coefficient_constraining_bijector",":","optional","`","tfb.Bijector","`","instance","representing","a","constraining","mapping","for","the","moving","average","coefficients",".","For","example",",","`","tfb.Tanh","(",")","`","constrains","the","coefficients","to","lie","in","`","(","-1",",","1",")","`",",","while","`","tfb.Softplus","(",")","`","constrains","them","to","be","positive",",","and","`","tfb.Identity","(",")","`","implies","no","constraint",".","If","`","None","`",",","the","default","behavior","is","to","apply","no","constraint",".","Default","value",":","`","None","`",".","observed_time_series",":","optional","`","float","`","`","Tensor","`","of","shape","`","batch_shape","+","[","T",",","1","]","`","(","omitting","the","trailing","unit","dimension","is","also","supported","when","`","T",">","1","`",")",",","specifying","an","observed","time","series",".","Any","`","NaN","`","s","are","interpreted","as","missing","observations",";","missingness","may","be","also","be","explicitly","specified","by","passing","a","`","tfp.sts.MaskedTimeSeries","`","instance",".","Any","priors","not","explicitly","set","will","be","given","default","values","according","to","the","scale","of","the","observed","time","series","(","or","batch","of","time","series",")",".","Default","value",":","`","None","`",".","name",":","the","name","of","this","model","component",".","Default","value",":","'ARIMA","'","."]
code: def __init__(
      self,
      ar_order,
      ma_order,
      integration_degree=0,
      ar_coefficients_prior=None,
      ma_coefficients_prior=None,
      level_drift_prior=None,
      level_scale_prior=None,
      initial_state_prior=None,
      ar_coefficient_constraining_bijector=None,
      ma_coefficient_constraining_bijector=None,
      observed_time_series=None,
      name=None,
  ):
      """Specifies an ARIMA(p=ar_order, d=integration_degree, q=ma_order) model.
  
      Args:
        ar_order: scalar Python positive `int` specifying the order of the
          autoregressive process (`p` in `ARIMA(p, d, q)`).
        ma_order: scalar Python positive `int` specifying the order of the
          moving-average process (`q` in `ARIMA(p, d, q)`).
        integration_degree: scalar Python positive `int` specifying the number
          of times to integrate an ARMA process. (`d` in `ARIMA(p, d, q)`).
          Default value: `0`.
        ar_coefficients_prior: optional `tfd.Distribution` instance specifying a
          prior on the `ar_coefficients` parameter. If `None`, a default standard
          normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([ar_order]))`)
          prior is used.
          Default value: `None`.
        ma_coefficients_prior: optional `tfd.Distribution` instance specifying a
          prior on the `ma_coefficients` parameter. If `None`, a default standard
          normal (`tfd.MultivariateNormalDiag(scale_diag=tf.ones([ma_order]))`)
          prior is used.
          Default value: `None`.
        level_drift_prior: optional `tfd.Distribution` instance specifying a prior
          on the `level_drift` parameter. If `None`, the parameter is not inferred
          and is instead fixed to zero.
          Default value: `None`.
        level_scale_prior: optional `tfd.Distribution` instance specifying a prior
          on the `level_scale` parameter. If `None`, a heuristic default prior is
          constructed based on the provided `observed_time_series`.
          Default value: `None`.
        initial_state_prior: optional `tfd.Distribution` instance specifying a
          prior on the initial state, corresponding to the values of the process
          at a set of size `order` of imagined timesteps before the initial step.
          If `None`, a heuristic default prior is constructed based on the
          provided `observed_time_series`.
          Default value: `None`.
        ar_coefficient_constraining_bijector: optional `tfb.Bijector` instance
          representing a constraining mapping for the autoregressive coefficients.
          For example, `tfb.Tanh()` constrains the coefficients to lie in
          `(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
          `tfb.Identity()` implies no constraint. If `None`, the default behavior
          constrains the coefficients to lie in `(-1, 1)` using a `Tanh` bijector.
          Default value: `None`.
        ma_coefficient_constraining_bijector: optional `tfb.Bijector` instance
          representing a constraining mapping for the moving average coefficients.
          For example, `tfb.Tanh()` constrains the coefficients to lie in
          `(-1, 1)`, while `tfb.Softplus()` constrains them to be positive, and
          `tfb.Identity()` implies no constraint. If `None`, the default behavior
          is to apply no constraint.
          Default value: `None`.
        observed_time_series: optional `float` `Tensor` of shape
          `batch_shape + [T, 1]` (omitting the trailing unit dimension is also
          supported when `T > 1`), specifying an observed time series. Any `NaN`s
          are interpreted as missing observations; missingness may be also be
          explicitly specified by passing a `tfp.sts.MaskedTimeSeries` instance.
          Any priors not explicitly set will be given default values according to
          the scale of the observed time series (or batch of time series).
          Default value: `None`.
        name: the name of this model component.
          Default value: 'ARIMA'.
      """
      init_parameters = dict(locals())
      with tf.name_scope(name or "ARIMA") as name:
          masked_time_series = None
          if observed_time_series is not None:
              masked_time_series = (
                  sts_util.canonicalize_observed_time_series_with_mask(
                      observed_time_series
                  )
              )
          dtype = dtype_util.common_dtype(
              [
                  (
                      masked_time_series.time_series
                      if masked_time_series is not None
                      else None
                  ),
                  ar_coefficients_prior,
                  ma_coefficients_prior,
                  level_scale_prior,
                  initial_state_prior,
              ],
              dtype_hint=tf.float32,
          )
  
          if observed_time_series is not None:
              for _ in range(integration_degree):
                  # Compute statistics using `integration_order`-order differences.
                  masked_time_series = missing_values_util.differentiate_masked_time_series(
                      masked_time_series
                  )
              (
                  _,
                  observed_stddev,
                  observed_initial,
              ) = sts_util.empirical_statistics(masked_time_series)
          else:
              observed_stddev, observed_initial = (
                  tf.convert_to_tensor(value=1.0, dtype=dtype),
                  tf.convert_to_tensor(value=0.0, dtype=dtype),
              )
          batch_ones = ps.ones(
              ps.concat(
                  [ps.shape(observed_initial), [1]],  # Batch shape
                  axis=0,
              ),
              dtype=dtype,
          )
  
          # Heuristic default priors. Overriding these may dramatically
          # change inference performance and results.
          if ar_coefficients_prior is None:
              ar_coefficients_prior = mvn_diag.MultivariateNormalDiag(
                  scale_diag=batch_ones * ps.ones([ar_order])
              )
          if ma_coefficients_prior is None:
              ma_coefficients_prior = mvn_diag.MultivariateNormalDiag(
                  scale_diag=batch_ones * ps.ones([ma_order])
              )
          if level_scale_prior is None:
              level_scale_prior = lognormal.LogNormal(
                  loc=tf.math.log(0.05 * observed_stddev), scale=3.0
              )
  
          if (
              ar_coefficients_prior.event_shape.is_fully_defined()
              and ar_order != ar_coefficients_prior.event_shape[0]
          ):
              raise ValueError(
                  "Autoregressive prior dimension {} doesn't match order {}.".format(
                      ar_coefficients_prior.event_shape[0], ar_order
                  )
              )
          if (
              ma_coefficients_prior.event_shape.is_fully_defined()
              and ma_order != ma_coefficients_prior.event_shape[0]
          ):
              raise ValueError(
                  "Moving average prior dimension {} doesn't match order {}.".format(
                      ma_coefficients_prior.event_shape[0], ma_order
                  )
              )
  
          latent_size = (
              ps.maximum(ar_order, ma_order + 1) + integration_degree
          )
          if initial_state_prior is None:
              initial_state_prior = mvn_diag.MultivariateNormalDiag(
                  loc=sts_util.pad_tensor_with_trailing_zeros(
                      observed_initial[..., tf.newaxis] * batch_ones,
                      num_zeros=latent_size - 1,
                  ),
                  scale_diag=sts_util.pad_tensor_with_trailing_zeros(
                      (tf.abs(observed_initial) + observed_stddev)[
                          ..., tf.newaxis
                      ]
                      * batch_ones,
                      num_zeros=latent_size - 1,
                  ),
              )
  
          self._ar_order = ar_order
          self._ma_order = ma_order
          self._integration_degree = integration_degree
          self._ar_coefficients_prior = ar_coefficients_prior
          self._ma_coefficients_prior = ma_coefficients_prior
          self._level_scale_prior = level_scale_prior
          self._initial_state_prior = initial_state_prior
  
          parameters = []
          if ar_order > 0:
              parameters.append(
                  Parameter(
                      "ar_coefficients",
                      ar_coefficients_prior,
                      (
                          ar_coefficient_constraining_bijector
                          if ar_coefficient_constraining_bijector
                          else tanh.Tanh()
                      ),
                  )
              )
          if ma_order > 0:
              parameters.append(
                  Parameter(
                      "ma_coefficients",
                      ma_coefficients_prior,
                      (
                          ma_coefficient_constraining_bijector
                          if ma_coefficient_constraining_bijector
                          else identity.Identity()
                      ),
                  )
              )
          if level_drift_prior is not None:
              parameters.append(
                  Parameter(
                      "level_drift",
                      level_drift_prior,
                      chain.Chain(
                          [
                              scale.Scale(scale=observed_stddev),
                              (
                                  level_drift_prior.experimental_default_event_space_bijector()
                              ),
                          ]
                      ),
                  )
              )
          super(AutoregressiveIntegratedMovingAverage, self).__init__(
              parameters=parameters
              + [
                  Parameter(
                      "level_scale",
                      level_scale_prior,
                      chain.Chain(
                          [
                              scale.Scale(scale=observed_stddev),
                              softplus.Softplus(
                                  low=dtype_util.eps(dtype)
                              ),
                          ]
                      ),
                  )
              ],
              latent_size=latent_size,
              init_parameters=init_parameters,
              name=name,
          )
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import identity
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.bijectors import tanh
  from tensorflow_probability.python.distributions import linear_gaussian_ssm
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.sts.components.autoregressive_moving_average import AutoregressiveMovingAverageStateSpaceModel
  from tensorflow_probability.python.sts.internal import missing_values_util
  from tensorflow_probability.python.sts.internal import util
  from tensorflow_probability.python.sts.structural_time_series import Parameter
  from tensorflow_probability.python.sts.structural_time_series import StructuralTimeSeries
label_desc: Use image node_id 1 to create a new AutoregressiveIntegratedMovingAverage object from inherited base classes: StructuralTimeSeries with example: obj = AutoregressiveIntegratedMovingAverage(ar_order, ma_order, integration_degree, ar_coefficients_prior, ma_coefficients_prior, level_drift_prior, level_scale_prior, initial_state_prior, ar_coefficient_constraining_bijector, ma_coefficient_constraining_bijector, observed_time_series, name)