Change dtype of prior parameter of sts.Seasonal from float32 to float64

[CleverEskimo]

trafficstars

Hi all,

I have question based on dtype change in sts.Seasonal, namely I have created multiple of objects like:

Parameter: local_linear_trend/_slope_scale
Prior: tfp.distributions.LogNormal("slope_scale_prior", batch_shape=[], event_shape=[], dtype=float64)
--------------------------------------------------------------------------------------------------------------------------------------------
Parameter: month_of_year/_drift_scale
Prior: tfp.distributions.LogNormal("LogNormal", batch_shape=[], event_shape=[], dtype=float32)
--------------------------------------------------------------------------------------------------------------------------------------------
Parameter: day_of_week/_drift_scale
Prior: tfp.distributions.LogNormal("LogNormal", batch_shape=[], event_shape=[], dtype=float64)

As you can see in the case of month_of_year/_drift_scale Prior dtype is float32. It looks like it is not allowed during training as there is an exception:

ValueError: ConstrainedSeasonalStateSpaceModel, type=<dtype: 'float32'>, must be of the same type (<dtype: 'float64'>) as LocalLinearTrendStateSpaceModel.

Is there any solution to change dtype from float32 to float64?

[jeffpollock9]

Hard to know without the code but I guess you are passing something which is float64 to the local linear trend and day of the week effects, then either nothing or float32 things to the month of the year effect.

import tensorflow as tf
import tensorflow_probability as tfp

seasonal = tfp.sts.Seasonal(
    num_seasons=7,
    num_steps_per_season=24,
    observed_time_series=tf.ones([10], tf.float64),
    name="day_of_week",
)
seasonal.dtype
# tf.float64

seasonal = tfp.sts.Seasonal(
    num_seasons=7,
    num_steps_per_season=24,
    name="day_of_week",
)
seasonal.dtype
# tf.float32