tf function for rep(x, times) with times as a tensor

tf_expand.R

tf_expand <- function(x, repeats) {
  
  x <- tf$constant(x, dtype = tf$float64)
  repeats <- tf$constant(repeats, dtype = tf$int32)
  
  shape <- tf$reduce_sum(repeats)
  
  idx <- tf$concat(c(tf$constant(0L, dtype = tf$int32)[tf$newaxis],
                     tf$cumsum(repeats[0L:(tf$size(repeats) - 1L)])),
                   axis = 0L)
  
  tmp_index <- tf$concat(c(tf$constant(0, dtype = tf$float64)[tf$newaxis],
                  x[0L:(tf$size(x) - 1L)]), axis = 0L)

  y <- tf$sparse_to_dense(
    sparse_indices = idx,
    output_shape = shape[tf$newaxis],
    sparse_values = x - tf$concat(c(tf$constant(0, dtype = tf$float64)[tf$newaxis],
                                    x[0L:(tf$size(x) - 1L)]), axis = 0L)
  )

  tf$cumsum(y)
    
}