class_docstr: Formal representation of a seasonal effect model.
  
  A seasonal effect model posits a fixed set of recurring, discrete 'seasons',
  each of which is active for a fixed number of timesteps and, while active,
  contributes a different effect to the time series. These are generally not
  meteorological seasons, but represent regular recurring patterns such as
  hour-of-day or day-of-week effects. Each season lasts for a fixed number of
  timesteps. The effect of each season drifts from one occurrence to the next
  following a Gaussian random walk:
  
  ```python
  effects[season, occurrence[i]] = (
    effects[season, occurrence[i-1]] + Normal(loc=0., scale=drift_scale))
  ```
  
  The `drift_scale` parameter governs the standard deviation of the random walk;
  for example, in a day-of-week model it governs the change in effect from this
  Monday to next Monday.
  
  #### Examples
  
  A seasonal effect model representing day-of-week seasonality on hourly data:
  
  ```python
  day_of_week = tfp.sts.Seasonal(num_seasons=7,
                                 num_steps_per_season=24,
                                 observed_time_series=y,
                                 name='day_of_week')
  ```
  
  A seasonal effect model representing month-of-year seasonality on daily data,
  with explicit priors:
  
  ```python
  month_of_year = tfp.sts.Seasonal(
    num_seasons=12,
    num_steps_per_season=[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
    drift_scale_prior=tfd.LogNormal(loc=-1., scale=0.1),
    initial_effect_prior=tfd.Normal(loc=0., scale=5.),
    name='month_of_year')
  ```
  
  Note that this version works over time periods not involving a leap year. A
  general implementation of month-of-year seasonality would require additional
  logic:
  
  ```python
  num_days_per_month = np.array(
    [[31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
     [31, 29, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],  # year with leap day
     [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31],
     [31, 28, 31, 30, 30, 31, 31, 31, 30, 31, 30, 31]])
  
  month_of_year = tfp.sts.Seasonal(
    num_seasons=12,
    num_steps_per_season=num_days_per_month,
    drift_scale_prior=tfd.LogNormal(loc=-1., scale=0.1),
    initial_effect_prior=tfd.Normal(loc=0., scale=5.),
    name='month_of_year')
  ```
  
  A model representing both day-of-week and hour-of-day seasonality, on hourly
  data:
  
  ```
  day_of_week = tfp.sts.Seasonal(num_seasons=7,
                                 num_steps_per_season=24,
                                 observed_time_series=y,
                                 name='day_of_week')
  hour_of_day = tfp.sts.Seasonal(num_seasons=24,
                                 num_steps_per_season=1,
                                 observed_time_series=y,
                                 name='hour_of_day')
  model = tfp.sts.Sum(components=[day_of_week, hour_of_day],
                      observed_time_series=y)
  ```
class_docstr_tok: ["Formal","representation","of","a","seasonal","effect","model",".","A","seasonal","effect","model","posits","a","fixed","set","of","recurring",",","discrete","'seasons","'",",","each","of","which","is","active","for","a","fixed","number","of","timesteps","and",",","while","active",",","contributes","a","different","effect","to","the","time","series",".","These","are","generally","not","meteorological","seasons",",","but","represent","regular","recurring","patterns","such","as","hour-of-day","or","day-of-week","effects",".","Each","season","lasts","for","a","fixed","number","of","timesteps",".","The","effect","of","each","season","drifts","from","one","occurrence","to","the","next","following","a","Gaussian","random","walk",":","``","`","python","effects","[","season",",","occurrence","[","i","]","]","=","(","effects","[","season",",","occurrence","[","i-1","]","]","+","Normal","(","loc=0.",",","scale=drift_scale",")",")","``","`","The","`","drift_scale","`","parameter","governs","the","standard","deviation","of","the","random","walk",";","for","example",",","in","a","day-of-week","model","it","governs","the","change","in","effect","from","this","Monday","to","next","Monday",".","#","#","#","#","Examples","A","seasonal","effect","model","representing","day-of-week","seasonality","on","hourly","data",":","``","`","python","day_of_week","=","tfp.sts.Seasonal","(","num_seasons=7",",","num_steps_per_season=24",",","observed_time_series=y",",","name='day_of_week","'",")","``","`","A","seasonal","effect","model","representing","month-of-year","seasonality","on","daily","data",",","with","explicit","priors",":","``","`","python","month_of_year","=","tfp.sts.Seasonal","(","num_seasons=12",",","num_steps_per_season=","[","31",",","28",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]",",","drift_scale_prior=tfd.LogNormal","(","loc=-1.",",","scale=0.1",")",",","initial_effect_prior=tfd.Normal","(","loc=0.",",","scale=5",".",")",",","name='month_of_year","'",")","``","`","Note","that","this","version","works","over","time","periods","not","involving","a","leap","year",".","A","general","implementation","of","month-of-year","seasonality","would","require","additional","logic",":","``","`","python","num_days_per_month","=","np.array","(","[","[","31",",","28",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]",",","[","31",",","29",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]",",","#","year","with","leap","day","[","31",",","28",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]",",","[","31",",","28",",","31",",","30",",","30",",","31",",","31",",","31",",","30",",","31",",","30",",","31","]","]",")","month_of_year","=","tfp.sts.Seasonal","(","num_seasons=12",",","num_steps_per_season=num_days_per_month",",","drift_scale_prior=tfd.LogNormal","(","loc=-1.",",","scale=0.1",")",",","initial_effect_prior=tfd.Normal","(","loc=0.",",","scale=5",".",")",",","name='month_of_year","'",")","``","`","A","model","representing","both","day-of-week","and","hour-of-day","seasonality",",","on","hourly","data",":","``","`","day_of_week","=","tfp.sts.Seasonal","(","num_seasons=7",",","num_steps_per_season=24",",","observed_time_series=y",",","name='day_of_week","'",")","hour_of_day","=","tfp.sts.Seasonal","(","num_seasons=24",",","num_steps_per_season=1",",","observed_time_series=y",",","name='hour_of_day","'",")","model","=","tfp.sts.Sum","(","components=","[","day_of_week",",","hour_of_day","]",",","observed_time_series=y",")","``","`"]
docstr: Returns a posterior sample of `drift scale` given the latents and prior.
  
  This is experimental, and compatibility may be broken in the future.
  
  Args:
    latents: A value convertible to a tensor of shape [batch shape, timeseries
      length, latent size].
    sample_shape: Shape of the sample.
    seed: The seed to sample with.
docstr_tok: ["Returns","a","posterior","sample","of","`","drift","scale","`","given","the","latents","and","prior",".","This","is","experimental",",","and","compatibility","may","be","broken","in","the","future",".","Args",":","latents",":","A","value","convertible","to","a","tensor","of","shape","[","batch","shape",",","timeseries","length",",","latent","size","]",".","sample_shape",":","Shape","of","the","sample",".","seed",":","The","seed","to","sample","with","."]
code: def experimental_resample_drift_scale(
      self, latents, sample_shape=(), seed=None
  ):
      """Returns a posterior sample of `drift scale` given the latents and prior.
  
      This is experimental, and compatibility may be broken in the future.
  
      Args:
        latents: A value convertible to a tensor of shape [batch shape, timeseries
          length, latent size].
        sample_shape: Shape of the sample.
        seed: The seed to sample with.
      """
      if not self.allow_drift:
          raise NotImplementedError(
              "experimental_resample_drift_scale is implemented only for seasonal "
              "effects with drift."
          )
      # The drift_scale prior is a scale prior, but this requires the variance
      # prior, thus extract it from the transformed distribution.
      prior = self.get_parameter("drift_scale").prior
      if not (
          isinstance(
              prior, transformed_distribution.TransformedDistribution
          )
          and isinstance(prior.bijector, invert.Invert)
          and isinstance(prior.bijector.bijector, square.Square)
          and isinstance(prior.distribution, inverse_gamma.InverseGamma)
      ):
          raise NotImplementedError(
              "experimental_resample_drift_scale requires drift scale prior be "
              "the square root of an inverse gamma distribution; got {}".format(
                  prior
              )
          )
      drift_scale_variance_prior = prior.distribution
  
      latents = tf.convert_to_tensor(latents, dtype_hint=self.dtype)
      batch_shape = ps.shape(latents)[:-2]
      num_observed_steps = tf.compat.dimension_value(latents.shape[-2])
      drift_scale_shape = ps.concat([batch_shape, [1]], axis=0)
      sssm = self.make_state_space_model(
          num_timesteps=num_observed_steps,
          initial_state_prior=self.initial_state_prior,
          param_vals={},
          # The scale value does not matter, since we are just using this to
          # access the transition matrices.
          drift_scale=tf.ones(
              shape=drift_scale_shape, dtype=self.dtype
          ),
      )
  
      if num_observed_steps < 2:
          # This needs to be special-cased only because:
          #   1) Running with Numpy returns an empty matrix with vectorized_map
          #      (size=(0,)) rather than size=(0, num_latents, num_latents).
          #   2) An issue with Numpy matrix multiplication when sizes are 0, which
          #      occurs if there is a single timestep (and thus no transition noise
          #      to observe).
          #
          # When there are no observed steps, there are no observations of noise.
          transition_noise_observations = tf.zeros(
              shape=(0), dtype=self.dtype
          )
          season_diffs_mask = tf.constant(
              # Value does not matter, not used with empty shape.
              value=True,
              shape=[0],
              dtype=tf.bool,
          )
      else:
          # Transition matrices represent a mapping from t -> t+1 time, hence
          # all but last.
          @tf.function
          def get_all_but_last_transition_matrices():
              return tf.vectorized_map(
                  lambda i: sssm.get_transition_matrix_for_timestep(
                      i
                  ).to_dense(),
                  tf.range(start=0, limit=num_observed_steps - 1),
              )
  
          all_but_last_transition_matrices = (
              get_all_but_last_transition_matrices()
          )
          # Compute the residuals of transition noise by comparing time-adjacent
          # latents, with the transition matrix applied to the previous timestep
          # so that what remains is the addition of transition noise. Since
          # we know that that there is a single transition noise for an entire
          # seasonal component, these residuals can be used to sample the
          # scale of transition noise.
          all_but_first_latents = latents[..., 1:, :]
          all_but_last_latents = latents[..., :-1, :]
  
          # For unconstrained seasonality, the transition noise is added to the
          # latent. For constrained seasonality, it is divided by the number of
          # seasons first, so it that scale needs to be undone.
          latent_to_transition_noise_scale = (
              self.num_seasons
              if self.constrain_mean_effect_to_zero
              else 1.0
          )
          transition_noise_observations = (
              latent_to_transition_noise_scale
              * (
                  all_but_first_latents
                  - tf.linalg.matvec(
                      # Transition matrices represent a mapping from t -> t+1 time,
                      # hence all but last is used to translate the previous latents to
                      # the next latent space.
                      all_but_last_transition_matrices,
                      all_but_last_latents,
                  )
              )[
                  ...,
                  # Using just the last element as:
                  #  1) For unconstrained seasonality, the transition noise is added
                  #     to the last latent.
                  #  2) For constrained seasonality, the transition noise, divided by
                  #     the number of seasons is added to each dimension, so it does
                  #     not matter which is sampled.
                  -1,
              ]
          )
          # Transition noise is only observed when there is a season change, which
          # occurs when there is a non-identity transition matrix. Since both
          # constrained and unconstrained seasonality use a zero in the first-row,
          # first-column value when it is not an identity matrix, this optimization
          # is used instead of checking for identity.
          #
          # Using this as a mask allows us to ignore all the non-changing residual
          # samples, and even handle the case where we, by-chance, observe
          # exactly 0 for transition noise.
          season_diffs_mask = tf.math.not_equal(
              all_but_last_transition_matrices[..., 0, 0], 0
          )
      posterior = sample_parameters.normal_scale_posterior_inverse_gamma_conjugate(
          variance_prior=drift_scale_variance_prior,
          observations=transition_noise_observations,
          is_missing=season_diffs_mask,
      )
      return sample_parameters.sample_with_optional_upper_bound(
          posterior, sample_shape=sample_shape, seed=seed
      )
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import chain
  from tensorflow_probability.python.bijectors import invert
  from tensorflow_probability.python.bijectors import scale
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.bijectors import square
  from tensorflow_probability.python.distributions import inverse_gamma
  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 mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.experimental.sts_gibbs import sample_parameters
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow_probability.python.internal import docstring_util
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.internal import prefer_static
  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 7 for calling the Seasonal obj's underlying member method code with example usage: obj.experimental_resample_drift_scale(latents, sample_shape, seed) and returns: sample_parameters, tf