class_docstr: State space model for a local linear trend.
  
  A state space model (SSM) posits a set of latent (unobserved) variables that
  evolve over time with dynamics specified by a probabilistic transition model
  `p(z[t+1] | z[t])`. At each timestep, we observe a value sampled from an
  observation model conditioned on the current state, `p(x[t] | z[t])`. The
  special case where both the transition and observation models are Gaussians
  with mean specified as a linear function of the inputs, is known as a linear
  Gaussian state space model and supports tractable exact probabilistic
  calculations; see `tfp.distributions.LinearGaussianStateSpaceModel` for
  details.
  
  The local linear trend model is a special case of a linear Gaussian SSM, in
  which the latent state posits a `level` and `slope`, each evolving via a
  Gaussian random walk:
  
  ```python
  level[t] = level[t-1] + slope[t-1] + Normal(0., level_scale)
  slope[t] = slope[t-1] + Normal(0., slope_scale)
  ```
  
  The latent state is the two-dimensional tuple `[level, slope]`. The
  `level` is observed at each timestep.
  
  The parameters `level_scale`, `slope_scale`, and `observation_noise_scale`
  are each (a batch of) scalars. The batch shape of this `Distribution` is the
  broadcast batch shape of these parameters and of the `initial_state_prior`.
  
  #### Mathematical Details
  
  The linear trend model implements a
  `tfp.distributions.LinearGaussianStateSpaceModel` with `latent_size = 2`
  and `observation_size = 1`, following the transition model:
  
  ```
  transition_matrix = [[1., 1.]
                       [0., 1.]]
  transition_noise ~ N(loc=0., scale=diag([level_scale, slope_scale]))
  ```
  
  which implements the evolution of `[level, slope]` described above, and
  the observation model:
  
  ```
  observation_matrix = [[1., 0.]]
  observation_noise ~ N(loc=0, scale=observation_noise_scale)
  ```
  
  which picks out the first latent component, i.e., the `level`, as the
  observation at each timestep.
  
  #### Examples
  
  A simple model definition:
  
  ```python
  linear_trend_model = LocalLinearTrendStateSpaceModel(
      num_timesteps=50,
      level_scale=0.5,
      slope_scale=0.5,
      initial_state_prior=tfd.MultivariateNormalDiag(scale_diag=[1., 1.]))
  
  y = linear_trend_model.sample() # y has shape [50, 1]
  lp = linear_trend_model.log_prob(y) # log_prob is scalar
  ```
  
  Passing additional parameter dimensions constructs a batch of models. The
  overall batch shape is the broadcast batch shape of the parameters:
  
  ```python
  linear_trend_model = LocalLinearTrendStateSpaceModel(
      num_timesteps=50,
      level_scale=tf.ones([10]),
      slope_scale=0.5,
      initial_state_prior=tfd.MultivariateNormalDiag(
        scale_diag=tf.ones([10, 10, 2])))
  
  y = linear_trend_model.sample(5) # y has shape [5, 10, 10, 50, 1]
  lp = linear_trend_model.log_prob(y) # has shape [5, 10, 10]
  ```
class_docstr_tok: ["State","space","model","for","a","local","linear","trend",".","A","state","space","model","(","SSM",")","posits","a","set","of","latent","(","unobserved",")","variables","that","evolve","over","time","with","dynamics","specified","by","a","probabilistic","transition","model","`","p","(","z","[","t+1","]","|","z","[","t","]",")","`",".","At","each","timestep",",","we","observe","a","value","sampled","from","an","observation","model","conditioned","on","the","current","state",",","`","p","(","x","[","t","]","|","z","[","t","]",")","`",".","The","special","case","where","both","the","transition","and","observation","models","are","Gaussians","with","mean","specified","as","a","linear","function","of","the","inputs",",","is","known","as","a","linear","Gaussian","state","space","model","and","supports","tractable","exact","probabilistic","calculations",";","see","`","tfp.distributions.LinearGaussianStateSpaceModel","`","for","details",".","The","local","linear","trend","model","is","a","special","case","of","a","linear","Gaussian","SSM",",","in","which","the","latent","state","posits","a","`","level","`","and","`","slope","`",",","each","evolving","via","a","Gaussian","random","walk",":","``","`","python","level","[","t","]","=","level","[","t-1","]","+","slope","[","t-1","]","+","Normal","(","0.",",","level_scale",")","slope","[","t","]","=","slope","[","t-1","]","+","Normal","(","0.",",","slope_scale",")","``","`","The","latent","state","is","the","two-dimensional","tuple","`","[","level",",","slope","]","`",".","The","`","level","`","is","observed","at","each","timestep",".","The","parameters","`","level_scale","`",",","`","slope_scale","`",",","and","`","observation_noise_scale","`","are","each","(","a","batch","of",")","scalars",".","The","batch","shape","of","this","`","Distribution","`","is","the","broadcast","batch","shape","of","these","parameters","and","of","the","`","initial_state_prior","`",".","#","#","#","#","Mathematical","Details","The","linear","trend","model","implements","a","`","tfp.distributions.LinearGaussianStateSpaceModel","`","with","`","latent_size","=","2","`","and","`","observation_size","=","1","`",",","following","the","transition","model",":","``","`","transition_matrix","=","[","[","1.",",","1",".","]","[","0.",",","1",".","]","]","transition_noise","~","N","(","loc=0.",",","scale=diag","(","[","level_scale",",","slope_scale","]",")",")","``","`","which","implements","the","evolution","of","`","[","level",",","slope","]","`","described","above",",","and","the","observation","model",":","``","`","observation_matrix","=","[","[","1.",",","0",".","]","]","observation_noise","~","N","(","loc=0",",","scale=observation_noise_scale",")","``","`","which","picks","out","the","first","latent","component",",","i.e.",",","the","`","level","`",",","as","the","observation","at","each","timestep",".","#","#","#","#","Examples","A","simple","model","definition",":","``","`","python","linear_trend_model","=","LocalLinearTrendStateSpaceModel","(","num_timesteps=50",",","level_scale=0.5",",","slope_scale=0.5",",","initial_state_prior=tfd.MultivariateNormalDiag","(","scale_diag=","[","1.",",","1",".","]",")",")","y","=","linear_trend_model.sample","(",")","#","y","has","shape","[","50",",","1","]","lp","=","linear_trend_model.log_prob","(","y",")","#","log_prob","is","scalar","``","`","Passing","additional","parameter","dimensions","constructs","a","batch","of","models",".","The","overall","batch","shape","is","the","broadcast","batch","shape","of","the","parameters",":","``","`","python","linear_trend_model","=","LocalLinearTrendStateSpaceModel","(","num_timesteps=50",",","level_scale=tf.ones","(","[","10","]",")",",","slope_scale=0.5",",","initial_state_prior=tfd.MultivariateNormalDiag","(","scale_diag=tf.ones","(","[","10",",","10",",","2","]",")",")",")","y","=","linear_trend_model.sample","(","5",")","#","y","has","shape","[","5",",","10",",","10",",","50",",","1","]","lp","=","linear_trend_model.log_prob","(","y",")","#","has","shape","[","5",",","10",",","10","]","``","`"]
docstr: Build a state space model implementing a local linear trend.
  
  Args:
    num_timesteps: Scalar `int` `Tensor` number of timesteps to model
      with this distribution.
    level_scale: Scalar (any additional dimensions are treated as batch
      dimensions) `float` `Tensor` indicating the standard deviation of the
      level transitions.
    slope_scale: Scalar (any additional dimensions are treated as batch
      dimensions) `float` `Tensor` indicating the standard deviation of the
      slope transitions.
    initial_state_prior: instance of `tfd.MultivariateNormal`
      representing the prior distribution on latent states; must
      have event shape `[2]`.
    observation_noise_scale: Scalar (any additional dimensions are
      treated as batch dimensions) `float` `Tensor` indicating the standard
      deviation of the observation noise.
    name: Python `str` name prefixed to ops created by this class.
      Default value: "LocalLinearTrendStateSpaceModel".
    **linear_gaussian_ssm_kwargs: Optional additional keyword arguments to
      to the base `tfd.LinearGaussianStateSpaceModel` constructor.
docstr_tok: ["Build","a","state","space","model","implementing","a","local","linear","trend",".","Args",":","num_timesteps",":","Scalar","`","int","`","`","Tensor","`","number","of","timesteps","to","model","with","this","distribution",".","level_scale",":","Scalar","(","any","additional","dimensions","are","treated","as","batch","dimensions",")","`","float","`","`","Tensor","`","indicating","the","standard","deviation","of","the","level","transitions",".","slope_scale",":","Scalar","(","any","additional","dimensions","are","treated","as","batch","dimensions",")","`","float","`","`","Tensor","`","indicating","the","standard","deviation","of","the","slope","transitions",".","initial_state_prior",":","instance","of","`","tfd.MultivariateNormal","`","representing","the","prior","distribution","on","latent","states",";","must","have","event","shape","`","[","2","]","`",".","observation_noise_scale",":","Scalar","(","any","additional","dimensions","are","treated","as","batch","dimensions",")","`","float","`","`","Tensor","`","indicating","the","standard","deviation","of","the","observation","noise",".","name",":","Python","`","str","`","name","prefixed","to","ops","created","by","this","class",".","Default","value",":","``","LocalLinearTrendStateSpaceModel","''",".","*","*","linear_gaussian_ssm_kwargs",":","Optional","additional","keyword","arguments","to","to","the","base","`","tfd.LinearGaussianStateSpaceModel","`","constructor","."]
code: def __init__(
      self,
      num_timesteps,
      level_scale,
      slope_scale,
      initial_state_prior,
      observation_noise_scale=0.0,
      name=None,
      **linear_gaussian_ssm_kwargs,
  ):
      """Build a state space model implementing a local linear trend.
  
      Args:
        num_timesteps: Scalar `int` `Tensor` number of timesteps to model
          with this distribution.
        level_scale: Scalar (any additional dimensions are treated as batch
          dimensions) `float` `Tensor` indicating the standard deviation of the
          level transitions.
        slope_scale: Scalar (any additional dimensions are treated as batch
          dimensions) `float` `Tensor` indicating the standard deviation of the
          slope transitions.
        initial_state_prior: instance of `tfd.MultivariateNormal`
          representing the prior distribution on latent states; must
          have event shape `[2]`.
        observation_noise_scale: Scalar (any additional dimensions are
          treated as batch dimensions) `float` `Tensor` indicating the standard
          deviation of the observation noise.
        name: Python `str` name prefixed to ops created by this class.
          Default value: "LocalLinearTrendStateSpaceModel".
        **linear_gaussian_ssm_kwargs: Optional additional keyword arguments to
          to the base `tfd.LinearGaussianStateSpaceModel` constructor.
      """
      parameters = dict(locals())
      parameters.update(linear_gaussian_ssm_kwargs)
      del parameters["linear_gaussian_ssm_kwargs"]
      with tf.name_scope(
          name or "LocalLinearTrendStateSpaceModel"
      ) as name:
          # The initial state prior determines the dtype of sampled values.
          # Other model parameters must have the same dtype.
          dtype = initial_state_prior.dtype
  
          level_scale = tf.convert_to_tensor(
              value=level_scale, name="level_scale", dtype=dtype
          )
          slope_scale = tf.convert_to_tensor(
              value=slope_scale, name="slope_scale", dtype=dtype
          )
          observation_noise_scale = tf.convert_to_tensor(
              value=observation_noise_scale,
              name="observation_noise_scale",
              dtype=dtype,
          )
  
          # Explicitly broadcast all parameters to the same batch shape. This
          # allows us to use `tf.stack` for a compact model specification.
          broadcast_batch_shape = dist_util.get_broadcast_shape(
              level_scale, slope_scale
          )
          broadcast_ones = tf.ones(broadcast_batch_shape, dtype=dtype)
  
          self._level_scale = level_scale
          self._slope_scale = slope_scale
          self._observation_noise_scale = observation_noise_scale
  
          # Construct a linear Gaussian state space model implementing the
          # local linear trend model. See "Mathematical Details" in the
          # class docstring for further explanation.
          super(LocalLinearTrendStateSpaceModel, self).__init__(
              num_timesteps=num_timesteps,
              transition_matrix=tf.constant(
                  [[1.0, 1.0], [0.0, 1.0]],
                  dtype=dtype,
                  name="transition_matrix",
              ),
              transition_noise=mvn_diag.MultivariateNormalDiag(
                  scale_diag=tf.stack(
                      [
                          level_scale * broadcast_ones,
                          slope_scale * broadcast_ones,
                      ],
                      axis=-1,
                  ),
                  name="transition_noise",
              ),
              observation_matrix=tf.constant(
                  [[1.0, 0.0]], dtype=dtype, name="observation_matrix"
              ),
              observation_noise=mvn_diag.MultivariateNormalDiag(
                  scale_diag=observation_noise_scale[..., tf.newaxis],
                  name="observation_noise",
              ),
              initial_state_prior=initial_state_prior,
              name=name,
              **linear_gaussian_ssm_kwargs,
          )
          self._parameters = parameters
  
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import softplus
  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.distributions import normal
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import samplers
  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 LocalLinearTrendStateSpaceModel object from inherited base classes: linear_gaussian_ssm with example: obj = LocalLinearTrendStateSpaceModel(num_timesteps, level_scale, slope_scale, initial_state_prior, observation_noise_scale, name)