sts_seasonal not accepting arrays

Per the TFP documentation sts_seasonal is supposed to be able to accept arrays for num_steps_per_season. It appears the current code forces all values to as.integer which prevents an array from being sent.

See the number of days per month here: https://www.tensorflow.org/probability/api_docs/python/tfp/sts/Seasonal

There is a workaround, I can create the array and pass it directly to tfp$sts$Seasonal.

library(tfprobability)

x = sample(1:100000,1000)
x = tf$cast(x, tf$float32)
first_cycle <- sample(1:100,12)
second_cycle <- sample(1:100,12)
days_array<-array(c(first_cycle,second_cycle),dim = c(2,12))
days_array<-np_array(days_array)

trend = tfprobability::sts_local_linear_trend(observed_time_series = x)

####################
#### THIS WORKS ####
####################
yearly<- tfp$sts$Seasonal(num_seasons = as.integer(12),num_steps_per_season = days_array)
####################

###########################
#### This Doesn't Work ####
###########################
yearly = tfprobability::sts_seasonal(num_seasons=12,num_steps_per_season = days_array,  observed_time_series=x,name = 'yearly')
###########################

model = sts_sum(components = list(trend,yearly), observed_time_series=x)
n_iterations <- 10
n_param_samples <- 50
n_forecast_samples <- 50
n_forecast_steps <- 120

optimizer <- tf$compat$v1$train$AdamOptimizer(0.1)

fit_vi <-
  function(ts,
           ts_train,
           model,
           n_iterations,
           n_param_samples,
           n_forecast_steps,
           n_forecast_samples) {
    loss_and_dists <-
      ts_train %>% sts_build_factored_variational_loss(model = model)
    variational_loss <- loss_and_dists[[1]]
    train_op <- optimizer$minimize(variational_loss)

    with (tf$compat$v1$Session() %as% sess,  {
      sess$run(tf$compat$v1$global_variables_initializer())
      for (step in 1:n_iterations) {
        res <- sess$run(train_op)
        loss <- sess$run(variational_loss)
        if (step %% 25 == 0)
          cat("Loss: ", as.numeric(loss), "\n")
      }
      variational_distributions <- loss_and_dists[[2]]
      posterior_samples <-
        Map(
          function(d)
            d %>% tfd_sample(n_param_samples),
          variational_distributions %>% unname()
        )
      print('postierier samples done')
      forecast_dists <-
        ts_train %>% sts_forecast(model, posterior_samples, n_forecast_steps)
      print('forecast dists done')

      fc_means <- forecast_dists %>% tfd_mean()
      fc_sds <- forecast_dists %>% tfd_stddev()

      c(posterior_samples,
        fc_means,
        fc_sds) %<-%
        sess$run(list(posterior_samples,
                      fc_means,
                      fc_sds))
    })

    list(variational_distributions,
         posterior_samples,
         fc_means[, 1],
         fc_sds[, 1])
  }

c(param_distributions,
  param_samples,
  fc_means,
  fc_sds) %<-% fit_vi(
    x,
    x,
    model,
    n_iterations,
    n_param_samples,
    n_forecast_steps,
    n_forecast_samples
  )

rstudio / tfprobability

sts_seasonal not accepting arrays #71