rnn.graph.dev.R

# 
#' Generate a RNN symbolic model - requires CUDA
#' 
#' @param config Either seq-to-one or one-to-one
#' @param cell_type Type of RNN cell: either gru or lstm
#' @param num_rnn_layer int, number of stacked layers
#' @param num_hidden int, size of the state in each RNN layer
#' @param num_embed  int, default = NULL - no embedding. Dimension of the embedding vectors
#' @param num_decode int, number of output variables in the decoding layer
#' @param input_size int, number of levels in the data - only used for embedding
#' @param dropout
#' 
#' @export
rnn.graph <- function (num_rnn_layer, input_size = NULL, num_embed = NULL, 
                       num_hidden, num_decode, dropout = 0, ignore_label = -1, bidirectional = F, 
                       loss_output = NULL, config, cell_type, masking = F, output_last_state = F,
                       rnn.state = NULL, rnn.state.cell = NULL, prefix = "") {
  
  data <- mx.symbol.Variable("data")
  label <- mx.symbol.Variable("label")
  seq.mask <- mx.symbol.Variable("seq.mask")
  if (!is.null(num_embed)) 
    embed.weight <- mx.symbol.Variable("embed.weight")
  rnn.params.weight <- mx.symbol.Variable("rnn.params.weight")
  
  if (is.null(rnn.state)) rnn.state <- mx.symbol.Variable("rnn.state")
  if (cell_type == "lstm" & is.null(rnn.state.cell)) {
    rnn.state.cell <- mx.symbol.Variable("rnn.state.cell")
  }
  
  cls.weight <- mx.symbol.Variable("cls.weight")
  cls.bias <- mx.symbol.Variable("cls.bias")
  if (!is.null(num_embed)) {
    data <- mx.symbol.Embedding(data = data, input_dim = input_size, 
                                weight = embed.weight, output_dim = num_embed, name = "embed")
  }
  
  data = mx.symbol.swapaxes(data = data, dim1 = 0, dim2 = 1, name = paste0(prefix, "swap_pre"))
  
  if (cell_type == "lstm") {
    rnn <- mx.symbol.RNN(data = data, state = rnn.state, 
                         state_cell = rnn.state.cell, parameters = rnn.params.weight, 
                         state.size = num_hidden, num.layers = num_rnn_layer, 
                         bidirectional = bidirectional, mode = cell_type, state.outputs = output_last_state, 
                         p = dropout, name = paste0(prefix, "RNN"))
  } else {
    rnn <- mx.symbol.RNN(data = data, state = rnn.state, 
                         parameters = rnn.params.weight, state.size = num_hidden, 
                         num.layers = num_rnn_layer, bidirectional = bidirectional, mode = cell_type, 
                         state.outputs = output_last_state, p = dropout, 
                         name = paste0(prefix, "RNN"))
  }
  
  if (config == "seq-to-one") {
    if (masking) mask <- mx.symbol.SequenceLast(data = rnn[[1]], use.sequence.length = T, sequence_length = seq.mask, name = "mask") else
      mask <- mx.symbol.SequenceLast(data = rnn[[1]], use.sequence.length = F, name = "mask")
    
    if (!is.null(loss_output)) {
      decode <- mx.symbol.FullyConnected(data = mask, weight = cls.weight, bias = cls.bias, num_hidden = num_decode, name = "decode")
      out <- switch(loss_output, softmax = mx.symbol.SoftmaxOutput(data = decode, label = label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, name = "loss"), 
                    linear = mx.symbol.LinearRegressionOutput(data = decode, label = label, name = "loss"), 
                    logistic = mx.symbol.LogisticRegressionOutput(data = decode, label = label, name = "loss"), 
                    MAE = mx.symbol.MAERegressionOutput(data = decode, label = label, name = "loss"))
    }
    else out <- mask
  }
  
  else if (config == "one-to-one") {
    
    if (masking) mask <- mx.symbol.SequenceMask(data = rnn[[1]], use.sequence.length = T, sequence_length = seq.mask, value = 0, name = "mask") else
      mask <- mx.symbol.identity(data = rnn[[1]], name = "mask")
    
    mask = mx.symbol.swapaxes(data = mask, dim1 = 0, dim2 = 1, name = paste0(prefix, "swap_post"))
    
    if (!is.null(loss_output)) {
      
      mask <- mx.symbol.reshape(data = mask, shape = c(0, -1), reverse = TRUE)
      label <- mx.symbol.reshape(data = label, shape = c(-1))
      
      decode <- mx.symbol.FullyConnected(data = mask, weight = cls.weight, bias = cls.bias, num_hidden = num_decode, 
                                         flatten = TRUE, name = paste0(prefix, "decode"))
      
      out <- switch(loss_output, softmax = mx.symbol.SoftmaxOutput(data = decode, label = label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, name = "loss"), 
                    linear = mx.symbol.LinearRegressionOutput(data = decode, label = label, name = "loss"), 
                    logistic = mx.symbol.LogisticRegressionOutput(data = decode, label = label, name = "loss"), 
                    MAE = mx.symbol.MAERegressionOutput(data = decode, label = label, name = "loss"))
    } else out <- mask
  }
  return(out)
}


# 
#' unroll representation of RNN running on non CUDA device
#' 
#' @export
rnn.graph.unroll <- function(num_rnn_layer, 
                             seq_len, 
                             input_size = NULL,
                             num_embed = NULL, 
                             num_hidden,
                             num_decode,
                             dropout = 0,
                             ignore_label = -1,
                             loss_output = NULL, 
                             init.state = NULL,
                             config,
                             cell_type = "lstm", 
                             masking = F, 
                             output_last_state = F,
                             prefix = "",
                             data_name = "data",
                             label_name = "label") {
  
  if (!is.null(num_embed)) embed.weight <- mx.symbol.Variable(paste0(prefix, "embed.weight"))
  
  # Initial state
  if (is.null(init.state) & output_last_state) {
    init.state <- lapply(1:num_rnn_layer, function(i) {
      if (cell_type=="lstm") {
        state <- list(h = mx.symbol.Variable(paste0("init_", prefix, i, "_h")),
                      c = mx.symbol.Variable(paste0("init_", prefix, i, "_c")))
      } else if (cell_type=="gru") {
        state <- list(h = mx.symbol.Variable(paste0("init_", prefix, i, "_h")))
      }
      return (state)
    })
  }
  
  cls.weight <- mx.symbol.Variable(paste0(prefix, "cls.weight"))
  cls.bias <- mx.symbol.Variable(paste0(prefix, "cls.bias"))
  
  param.cells <- lapply(1:num_rnn_layer, function(i) {
    
    if (cell_type=="lstm") {
      cell <- list(i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".i2h.weight")),
                   i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".i2h.bias")),
                   h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".h2h.weight")),
                   h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".h2h.bias")))
    } else if (cell_type=="gru") {
      cell <- list(gates.i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.i2h.weight")),
                   gates.i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.i2h.bias")),
                   gates.h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.h2h.weight")),
                   gates.h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.h2h.bias")),
                   trans.i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.i2h.weight")),
                   trans.i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.i2h.bias")),
                   trans.h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.h2h.weight")),
                   trans.h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.h2h.bias")))
    } else if (cell_type=="sru") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   write.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.weight")),
                   write.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.weight")),
                   read.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.bias")))
    } else if (cell_type=="straight") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   write.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.weight")),
                   write.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.weight")),
                   read.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.bias")))
    } else if (cell_type=="rich") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   write.in.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.in.weight")),
                   write.in.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.in.bias")),
                   write.c.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.c.weight")),
                   write.c.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.c.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.in.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.in.weight")),
                   read.in.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.in.bias")),
                   read.c.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.c.weight")),
                   read.c.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.c.bias")))
    }
    return (cell)
  })
  
  # embeding layer
  data <- mx.symbol.Variable(data_name)
  label <- mx.symbol.Variable(label_name)
  seq.mask <- mx.symbol.Variable(paste0(prefix, "seq.mask"))
  
  if (!is.null(num_embed)) {
    data <- mx.symbol.Embedding(data = data, input_dim = input_size,
                                weight=embed.weight, output_dim = num_embed, name = paste0(prefix, "embed"))
  }
  
  # swap from [features X seq_len X batch] to [features X batch X seq_len]
  data = mx.symbol.swapaxes(data = data, dim1 = 0, dim2 = 1, name = paste0(prefix, "swap_pre"))
  
  # split into seq_len objects of size: [features X seq_len X batch] to [features X batch X seq_len]
  data <- mx.symbol.split(data = data, axis = 0, num.outputs = seq_len, squeeze_axis = T)
  
  last.hidden <- list()
  last.states <- list()
  
  for (seqidx in 1:seq_len) {
    hidden <- data[[seqidx]]
    
    for (i in 1:num_rnn_layer) {
      
      if (seqidx==1) prev.state <- init.state[[i]] else 
        prev.state <- last.states[[i]]
      
      if (cell_type=="lstm") {
        cell.symbol <- lstm.cell
      } else if (cell_type=="gru"){
        cell.symbol <- gru.cell
      } else if (cell_type=="sru"){
        cell.symbol <- sru.cell
      } else if (cell_type=="straight"){
        cell.symbol <- straight.cell
      } else if (cell_type=="rich"){
        cell.symbol <- rich.cell
      }
      
      next.state <- cell.symbol(num_hidden = num_hidden, 
                                indata = hidden,
                                prev.state = prev.state,
                                param = param.cells[[i]],
                                seqidx = seqidx, 
                                layeridx = i,
                                dropout = dropout,
                                prefix = prefix)
      
      hidden <- next.state$h
      
      # if (dropout > 0)
      #   hidden <- mx.symbol.Dropout(data=hidden, p=dropout)
      
      last.states[[i]] <- next.state
    }
    
    # Aggregate outputs from each timestep
    last.hidden <- c(last.hidden, hidden)
  }
  
  if (output_last_state) {
    out.states = mx.symbol.Group(unlist(last.states))
  }
  
  # concat hidden units - concat seq_len blocks of dimension num_hidden x batch.size
  concat <- mx.symbol.concat(data = last.hidden, num.args = seq_len, dim = 0, name = paste0(prefix, "concat"))
  concat <- mx.symbol.reshape(data = concat, shape = c(num_hidden, -1, seq_len), name = paste0(prefix, "rnn_reshape"))
  
  if (config=="seq-to-one") {
    
    if (masking) mask <- mx.symbol.SequenceLast(data=concat, use.sequence.length = T, sequence_length = seq.mask, name = paste0(prefix, "mask")) else
      mask <- mx.symbol.SequenceLast(data=concat, use.sequence.length = F, name = paste0(prefix, "mask"))
    
    if (!is.null(loss_output)) {
      
      decode <- mx.symbol.FullyConnected(data = mask,
                                         weight = cls.weight,
                                         bias = cls.bias,
                                         num_hidden = num_decode,
                                         name = paste0(prefix, "decode"))
      
      out <- switch(loss_output,
                    softmax = mx.symbol.SoftmaxOutput(data=decode, label=label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, name = paste0(prefix, "loss")),
                    linear = mx.symbol.LinearRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    logistic = mx.symbol.LogisticRegressionOutput(data=decode, label=label, paste0(prefix, name = "loss")),
                    MAE = mx.symbol.MAERegressionOutput(data=decode, label=label, paste0(prefix, name = "loss"))
      )
    } else out <- mask
    
  } else if (config=="one-to-one"){
    
    if (masking) mask <- mx.symbol.SequenceMask(data = concat, use.sequence.length = T, sequence_length = seq.mask, value = 0, name = paste0(prefix, "mask")) else
      mask <- mx.symbol.identity(data = concat, name = paste0(prefix, "mask"))
    
    mask = mx.symbol.swapaxes(data = mask, dim1 = 0, dim2 = 1, name = paste0(prefix, "swap_post"))
    
    if (!is.null(loss_output)) {
      
      mask <- mx.symbol.reshape(data = mask, shape = c(0, -1), reverse = TRUE)
      label <- mx.symbol.reshape(data = label, shape = c(-1))
      
      decode <- mx.symbol.FullyConnected(data = mask, weight = cls.weight, bias = cls.bias, num_hidden = num_decode, 
                                         flatten = T, name = paste0(prefix, "decode"))
      
      out <- switch(loss_output,
                    softmax = mx.symbol.SoftmaxOutput(data=decode, label=label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, 
                                                      name = paste0(prefix, "loss")),
                    linear = mx.symbol.LinearRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    logistic = mx.symbol.LogisticRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    MAE = mx.symbol.MAERegressionOutput(data=decode, label=label, name = paste0(prefix, "loss"))
      )
    } else out <- mask
  }
  
  if (output_last_state) {
    return(mx.symbol.Group(c(out, out.states)))
  } else return(out)
}



# 
#' unroll representation of RNN running on non CUDA device
#' 
#' @export
rnn.graph.unroll.sru <- function(num_rnn_layer, 
                                 seq_len, 
                                 input_size = NULL,
                                 num_embed = NULL, 
                                 num_hidden,
                                 num_decode,
                                 dropout = 0,
                                 ignore_label = -1,
                                 loss_output = NULL, 
                                 init.state = NULL,
                                 config,
                                 cell_type = "lstm", 
                                 masking = F, 
                                 output_last_state = F,
                                 prefix = "",
                                 data_name = "data",
                                 label_name = "label") {
  
  if (!is.null(num_embed)) embed.weight <- mx.symbol.Variable(paste0(prefix, "embed.weight"))
  
  # Initial state
  if (is.null(init.state) & output_last_state) {
    init.state <- lapply(1:num_rnn_layer, function(i) {
      if (cell_type=="lstm") {
        state <- list(h = mx.symbol.Variable(paste0("init_", prefix, i, "_h")),
                      c = mx.symbol.Variable(paste0("init_", prefix, i, "_c")))
      } else if (cell_type=="gru") {
        state <- list(h = mx.symbol.Variable(paste0("init_", prefix, i, "_h")))
      }
      return (state)
    })
  }
  
  cls.weight <- mx.symbol.Variable(paste0(prefix, "cls.weight"))
  cls.bias <- mx.symbol.Variable(paste0(prefix, "cls.bias"))
  
  param.cells <- lapply(1:num_rnn_layer, function(i) {
    
    if (cell_type=="lstm") {
      cell <- list(i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".i2h.weight")),
                   i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".i2h.bias")),
                   h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".h2h.weight")),
                   h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".h2h.bias")))
    } else if (cell_type=="gru") {
      cell <- list(gates.i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.i2h.weight")),
                   gates.i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.i2h.bias")),
                   gates.h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.h2h.weight")),
                   gates.h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.h2h.bias")),
                   trans.i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.i2h.weight")),
                   trans.i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.i2h.bias")),
                   trans.h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.h2h.weight")),
                   trans.h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.h2h.bias")))
    } else if (cell_type=="sru") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   forget.weight = mx.symbol.Variable(paste0(prefix, "l", i, "forget.weight")),
                   forget.bias = mx.symbol.Variable(paste0(prefix, "l", i, "forget.bias")),
                   reset.weight = mx.symbol.Variable(paste0(prefix, "l", i, "reset.weight")),
                   reset.bias = mx.symbol.Variable(paste0(prefix, "l", i, "reset.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.weight")),
                   read.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.bias")))
    } else if (cell_type=="straight") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   write.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.weight")),
                   write.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.weight")),
                   read.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.bias")))
    } else if (cell_type=="rich") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   write.in.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.in.weight")),
                   write.in.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.in.bias")),
                   write.c.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.c.weight")),
                   write.c.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.c.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.in.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.in.weight")),
                   read.in.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.in.bias")),
                   read.c.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.c.weight")),
                   read.c.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.c.bias")))
    }
    return (cell)
  })
  
  # embeding layer
  data <- mx.symbol.Variable(data_name)
  label <- mx.symbol.Variable(label_name)
  seq.mask <- mx.symbol.Variable(paste0(prefix, "seq.mask"))
  
  if (!is.null(num_embed)) {
    data <- mx.symbol.Embedding(data = data, input_dim = input_size,
                                weight=embed.weight, output_dim = num_embed, name = paste0(prefix, "embed"))
  }
  
  # swap from [features X seq_len X batch] to [features X batch X seq_len]
  # concat = mx.symbol.swapaxes(data = data, dim1 = 0, dim2 = 1, name = paste0(prefix, "swap_pre"))
  
  # split into seq_len objects of size: [features X seq_len X batch] to [features X batch X seq_len]
  # data <- mx.symbol.split(data = data, axis = 0, num.outputs = seq_len, squeeze_axis = T)
  
  last.hidden <- list()
  last.states <- list()
  hidden <- list()
  
  for (i in 1:num_rnn_layer) {
    
    param <- param.cells[[i]]
    
    in_proj <- mx.symbol.FullyConnected(data = data, weight = param$input.weight, bias = param$input.bias, flatten = F,
                                        num_hidden = num_hidden, name = paste0(prefix, "input", "_l", i))
    
    forget <- mx.symbol.FullyConnected(data = data, weight = param$forget.weight, bias = param$forget.bias, flatten = F,
                                       num_hidden = num_hidden, name = paste0(prefix, "forget", "_l", i))
    
    reset <- mx.symbol.FullyConnected(data = data, weight = param$reset.weight, bias = param$reset.bias, flatten = F,
                                      num_hidden = num_hidden, name = paste0(prefix, "reset", "_l", i))
    
    # proj_full <- mx.symbol.split(proj_full, num_outputs = 3, axis=2, squeeze.axis = F)
    # in_proj <- proj_full[[1]]
    # forget <- proj_full[[2]]
    # reset <- proj_full[[3]]
    
    forget <- mx.symbol.sigmoid(forget)
    reset <- mx.symbol.sigmoid(reset)
    
    indata <- mx.symbol.split(data, num_outputs = seq_len, axis=1, squeeze.axis = F)
    in_proj <- mx.symbol.split(in_proj, num_outputs = seq_len, axis=1, squeeze.axis = F)
    forget <- mx.symbol.split(forget, num_outputs = seq_len, axis=1, squeeze.axis = F)
    reset <- mx.symbol.split(reset, num_outputs = seq_len, axis=1, squeeze.axis = F)
    
    for (seqidx in 1:seq_len) {
      
      if (seqidx == 1) {
        next.c <- (1 - forget[[seqidx]]) * in_proj[[seqidx]]
        next.h <- (1 - reset[[seqidx]]) * indata[[seqidx]]
      } else {
        next.c <- forget[[seqidx]] * prev.c + (1 - forget[[seqidx]]) * in_proj[[seqidx]]
        next.h <- reset[[seqidx]] * mx.symbol.tanh(next.c) + (1 - reset[[seqidx]]) * indata[[seqidx]]
      }
      
      prev.c <- next.c
      hidden <- c(hidden, next.h)
    }
    
    # concat hidden units - concat seq_len blocks of dimension num_hidden x batch.size -> 
    data <- mx.symbol.concat(hidden, num.args=seq_len, dim=1)
    
  }
  
  # if (output_last_state) {
  #   out.states = mx.symbol.Group(unlist(last.states))
  # }

  if (config=="seq-to-one") {
    
    if (masking) mask <- mx.symbol.SequenceLast(data=data, use.sequence.length = T, sequence_length = seq.mask, axis = 1, name = paste0(prefix, "mask")) else
      mask <- mx.symbol.SequenceLast(data=data, use.sequence.length = F, axis = 1, name = paste0(prefix, "mask"))
    
    if (!is.null(loss_output)) {
      
      decode <- mx.symbol.FullyConnected(data = mask,
                                         weight = cls.weight,
                                         bias = cls.bias,
                                         num_hidden = num_decode,
                                         name = paste0(prefix, "decode"))
      
      out <- switch(loss_output,
                    softmax = mx.symbol.SoftmaxOutput(data=decode, label=label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, name = paste0(prefix, "loss")),
                    linear = mx.symbol.LinearRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    logistic = mx.symbol.LogisticRegressionOutput(data=decode, label=label, paste0(prefix, name = "loss")),
                    MAE = mx.symbol.MAERegressionOutput(data=decode, label=label, paste0(prefix, name = "loss"))
      )
    } else out <- mask
    
  } else if (config=="one-to-one"){
    
    if (masking) mask <- mx.symbol.SequenceMask(data = data, use.sequence.length = T, sequence_length = seq.mask, axis = 1, value = 0, name = paste0(prefix, "mask")) else
      mask <- mx.symbol.identity(data = data, name = paste0(prefix, "mask"))
    
    # mask = mx.symbol.swapaxes(data = mask, dim1 = 0, dim2 = 1, name = paste0(prefix, "swap_post"))
    
    if (!is.null(loss_output)) {
      
      mask <- mx.symbol.reshape(data = mask, shape = c(0, -1), reverse = TRUE)
      label <- mx.symbol.reshape(data = label, shape = c(-1))
      
      decode <- mx.symbol.FullyConnected(data = mask, weight = cls.weight, bias = cls.bias, num_hidden = num_decode, 
                                         flatten = T, name = paste0(prefix, "decode"))
      
      out <- switch(loss_output,
                    softmax = mx.symbol.SoftmaxOutput(data=decode, label=label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, 
                                                      name = paste0(prefix, "loss")),
                    linear = mx.symbol.LinearRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    logistic = mx.symbol.LogisticRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    MAE = mx.symbol.MAERegressionOutput(data=decode, label=label, name = paste0(prefix, "loss"))
      )
    } else out <- mask
  }
  
  if (output_last_state) {
    return(mx.symbol.Group(c(out, out.states)))
  } else return(out)
}


# 
#' unroll representation of RNN running on non CUDA device
#' 
#' @export
rnn.graph.unroll.straight <- function(num_rnn_layer, 
                                      seq_len, 
                                      input_size = NULL,
                                      num_embed = NULL, 
                                      num_hidden,
                                      num_decode,
                                      dropout = 0,
                                      ignore_label = -1,
                                      loss_output = NULL, 
                                      init.state = NULL,
                                      config,
                                      cell_type = "lstm", 
                                      masking = F, 
                                      output_last_state = F,
                                      prefix = "",
                                      data_name = "data",
                                      label_name = "label") {
  
  if (!is.null(num_embed)) embed.weight <- mx.symbol.Variable(paste0(prefix, "embed.weight"))
  
  # Initial state
  if (is.null(init.state) & output_last_state) {
    init.state <- lapply(1:num_rnn_layer, function(i) {
      if (cell_type=="lstm") {
        state <- list(h = mx.symbol.Variable(paste0("init_", prefix, i, "_h")),
                      c = mx.symbol.Variable(paste0("init_", prefix, i, "_c")))
      } else if (cell_type=="gru") {
        state <- list(h = mx.symbol.Variable(paste0("init_", prefix, i, "_h")))
      }
      return (state)
    })
  }
  
  cls.weight <- mx.symbol.Variable(paste0(prefix, "cls.weight"))
  cls.bias <- mx.symbol.Variable(paste0(prefix, "cls.bias"))
  
  param.cells <- lapply(1:num_rnn_layer, function(i) {
    
    if (cell_type=="lstm") {
      cell <- list(i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".i2h.weight")),
                   i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".i2h.bias")),
                   h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".h2h.weight")),
                   h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".h2h.bias")))
    } else if (cell_type=="gru") {
      cell <- list(gates.i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.i2h.weight")),
                   gates.i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.i2h.bias")),
                   gates.h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.h2h.weight")),
                   gates.h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.h2h.bias")),
                   trans.i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.i2h.weight")),
                   trans.i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.i2h.bias")),
                   trans.h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.h2h.weight")),
                   trans.h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.h2h.bias")))
    } else if (cell_type=="sru") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   forget.weight = mx.symbol.Variable(paste0(prefix, "l", i, "forget.weight")),
                   forget.bias = mx.symbol.Variable(paste0(prefix, "l", i, "forget.bias")),
                   reset.weight = mx.symbol.Variable(paste0(prefix, "l", i, "reset.weight")),
                   reset.bias = mx.symbol.Variable(paste0(prefix, "l", i, "reset.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.weight")),
                   read.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.bias")))
    } else if (cell_type=="straight") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   write.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.weight")),
                   write.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.weight")),
                   read.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.bias")))
    } else if (cell_type=="rich") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   write.in.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.in.weight")),
                   write.in.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.in.bias")),
                   write.c.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.c.weight")),
                   write.c.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.c.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.in.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.in.weight")),
                   read.in.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.in.bias")),
                   read.c.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.c.weight")),
                   read.c.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.c.bias")))
    }
    return (cell)
  })
  
  # embeding layer
  data <- mx.symbol.Variable(data_name)
  label <- mx.symbol.Variable(label_name)
  seq.mask <- mx.symbol.Variable(paste0(prefix, "seq.mask"))
  
  if (!is.null(num_embed)) {
    data <- mx.symbol.Embedding(data = data, input_dim = input_size,
                                weight=embed.weight, output_dim = num_embed, name = paste0(prefix, "embed"))
  }
  
  last.hidden <- list()
  last.states <- list()
  
  for (i in 1:num_rnn_layer) {
    
    hidden <- list()
    param <- param.cells[[i]]
    
    in_proj <- mx.symbol.FullyConnected(data = data, weight = param$input.weight, bias = param$input.bias, flatten = F,
                                        num_hidden = num_hidden, name = paste0(prefix, "input", "_l", i))
    
    mem <- mx.symbol.FullyConnected(data = data, weight = param$mem.weight, bias = param$mem.bias, flatten = F,
                                    num_hidden = num_hidden, name = paste0(prefix, "mem", "_l", i))
    
    write <- mx.symbol.FullyConnected(data = data, weight = param$write.weight, bias = param$write.bias, flatten = F,
                                      num_hidden = num_hidden, name = paste0(prefix, "write", "_l", i))
    
    read <- mx.symbol.FullyConnected(data = data, weight = param$read.weight, bias = param$read.bias, flatten = F,
                                     num_hidden = num_hidden, name = paste0(prefix, "read", "_l", i))
    
    highway <- mx.symbol.FullyConnected(data = data, weight = param$highway.weight, bias = param$highway.bias, flatten = F,
                                        num_hidden = num_hidden, name = paste0(prefix, "highway", "_l", i))
    
    mem <- mx.symbol.relu(mem)
    write <- mx.symbol.tanh(write)
    read <- mx.symbol.sigmoid(read)
    highway <- mx.symbol.sigmoid(highway)
    
    in_proj <- mx.symbol.split(in_proj, num_outputs = seq_len, axis=1, squeeze.axis = F)
    mem <- mx.symbol.split(mem, num_outputs = seq_len, axis=1, squeeze.axis = F)
    write <- mx.symbol.split(write, num_outputs = seq_len, axis=1, squeeze.axis = F)
    read <- mx.symbol.split(read, num_outputs = seq_len, axis=1, squeeze.axis = F)
    highway <- mx.symbol.split(highway, num_outputs = seq_len, axis=1, squeeze.axis = F)
    
    for (seqidx in 1:seq_len) {
      
      if (seqidx == 1) {
        next.c <- write[[seqidx]] * mem[[seqidx]]
        next.h <-  highway[[seqidx]] * in_proj[[seqidx]]
      } else {
        next.c <- prev.c + write[[seqidx]] * mem[[seqidx]]
        next.h <- read[[seqidx]] * prev.c + highway[[seqidx]] * in_proj[[seqidx]]
      }
      
      prev.c <- next.c
      hidden <- c(hidden, next.h)
    }
    
    # concat hidden units - concat seq_len blocks of dimension num_hidden x batch.size -> 
    data <- mx.symbol.concat(hidden, num.args=seq_len, dim=1)
    
  }
  
  # if (output_last_state) {
  #   out.states = mx.symbol.Group(unlist(last.states))
  # }
  
  if (config=="seq-to-one") {
    
    if (masking) mask <- mx.symbol.SequenceLast(data = data, use.sequence.length = T, sequence_length = seq.mask, axis = 1, name = paste0(prefix, "mask")) else
      mask <- mx.symbol.SequenceLast(data = data, use.sequence.length = F, axis = 1, name = paste0(prefix, "mask"))
    
    if (!is.null(loss_output)) {
      
      decode <- mx.symbol.FullyConnected(data = mask,
                                         weight = cls.weight,
                                         bias = cls.bias,
                                         num_hidden = num_decode,
                                         name = paste0(prefix, "decode"))
      
      out <- switch(loss_output,
                    softmax = mx.symbol.SoftmaxOutput(data=decode, label=label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, name = paste0(prefix, "loss")),
                    linear = mx.symbol.LinearRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    logistic = mx.symbol.LogisticRegressionOutput(data=decode, label=label, paste0(prefix, name = "loss")),
                    MAE = mx.symbol.MAERegressionOutput(data=decode, label=label, paste0(prefix, name = "loss"))
      )
    } else out <- mask
    
  } else if (config=="one-to-one"){
    
    if (masking) mask <- mx.symbol.SequenceMask(data = data, use.sequence.length = T, sequence_length = seq.mask, axis = 1, value = 0, name = paste0(prefix, "mask")) else
      mask <- mx.symbol.identity(data = data, name = paste0(prefix, "mask"))
    
    # mask = mx.symbol.swapaxes(data = mask, dim1 = 0, dim2 = 1, name = paste0(prefix, "swap_post"))
    
    if (!is.null(loss_output)) {
      
      mask <- mx.symbol.reshape(data = mask, shape = c(0, -1), reverse = TRUE)
      label <- mx.symbol.reshape(data = label, shape = c(-1))
      
      decode <- mx.symbol.FullyConnected(data = mask, weight = cls.weight, bias = cls.bias, num_hidden = num_decode, 
                                         flatten = T, name = paste0(prefix, "decode"))
      
      out <- switch(loss_output,
                    softmax = mx.symbol.SoftmaxOutput(data=decode, label=label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, 
                                                      name = paste0(prefix, "loss")),
                    linear = mx.symbol.LinearRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    logistic = mx.symbol.LogisticRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    MAE = mx.symbol.MAERegressionOutput(data=decode, label=label, name = paste0(prefix, "loss"))
      )
    } else out <- mask
  }
  
  if (output_last_state) {
    return(mx.symbol.Group(c(out, out.states)))
  } else return(out)
}



rnn.graph.unroll.rich <- function(num_rnn_layer, 
                                  seq_len, 
                                  input_size = NULL,
                                  num_embed = NULL, 
                                  num_hidden,
                                  num_decode,
                                  dropout = 0,
                                  ignore_label = -1,
                                  loss_output = NULL, 
                                  init.state = NULL,
                                  config,
                                  cell_type = "lstm", 
                                  masking = F, 
                                  output_last_state = F,
                                  prefix = "",
                                  data_name = "data",
                                  label_name = "label") {
  
  if (!is.null(num_embed)) embed.weight <- mx.symbol.Variable(paste0(prefix, "embed.weight"))
  
  # Initial state
  if (is.null(init.state) & output_last_state) {
    init.state <- lapply(1:num_rnn_layer, function(i) {
      if (cell_type=="lstm") {
        state <- list(h = mx.symbol.Variable(paste0("init_", prefix, i, "_h")),
                      c = mx.symbol.Variable(paste0("init_", prefix, i, "_c")))
      } else if (cell_type=="gru") {
        state <- list(h = mx.symbol.Variable(paste0("init_", prefix, i, "_h")))
      }
      return (state)
    })
  }
  
  cls.weight <- mx.symbol.Variable(paste0(prefix, "cls.weight"))
  cls.bias <- mx.symbol.Variable(paste0(prefix, "cls.bias"))
  
  param.cells <- lapply(1:num_rnn_layer, function(i) {
    
    if (cell_type=="lstm") {
      cell <- list(i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".i2h.weight")),
                   i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".i2h.bias")),
                   h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".h2h.weight")),
                   h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".h2h.bias")))
    } else if (cell_type=="gru") {
      cell <- list(gates.i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.i2h.weight")),
                   gates.i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.i2h.bias")),
                   gates.h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.h2h.weight")),
                   gates.h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".gates.h2h.bias")),
                   trans.i2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.i2h.weight")),
                   trans.i2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.i2h.bias")),
                   trans.h2h.weight = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.h2h.weight")),
                   trans.h2h.bias = mx.symbol.Variable(paste0(prefix, "l", i, ".trans.h2h.bias")))
    } else if (cell_type=="sru") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   forget.weight = mx.symbol.Variable(paste0(prefix, "l", i, "forget.weight")),
                   forget.bias = mx.symbol.Variable(paste0(prefix, "l", i, "forget.bias")),
                   reset.weight = mx.symbol.Variable(paste0(prefix, "l", i, "reset.weight")),
                   reset.bias = mx.symbol.Variable(paste0(prefix, "l", i, "reset.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.weight")),
                   read.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.bias")))
    } else if (cell_type=="straight") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   write.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.weight")),
                   write.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.weight")),
                   read.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.bias")))
    } else if (cell_type=="rich") {
      cell <- list(input.weight = mx.symbol.Variable(paste0(prefix, "l", i, "input.weight")),
                   input.bias = mx.symbol.Variable(paste0(prefix, "l", i, "input.bias")),
                   mem.weight = mx.symbol.Variable(paste0(prefix, "l", i, "mem.weight")),
                   mem.bias = mx.symbol.Variable(paste0(prefix, "l", i, "mem.bias")),
                   write.in.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.in.weight")),
                   write.in.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.in.bias")),
                   write.c.weight = mx.symbol.Variable(paste0(prefix, "l", i, "write.c.weight")),
                   write.c.bias = mx.symbol.Variable(paste0(prefix, "l", i, "write.c.bias")),
                   highway.weight = mx.symbol.Variable(paste0(prefix, "l", i, "highway.weight")),
                   highway.bias = mx.symbol.Variable(paste0(prefix, "l", i, "highway.bias")),
                   read.in.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.in.weight")),
                   read.in.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.in.bias")),
                   read.c.weight = mx.symbol.Variable(paste0(prefix, "l", i, "read.c.weight")),
                   read.c.bias = mx.symbol.Variable(paste0(prefix, "l", i, "read.c.bias")))
    }
    return (cell)
  })
  
  # embeding layer
  data <- mx.symbol.Variable(data_name)
  label <- mx.symbol.Variable(label_name)
  seq.mask <- mx.symbol.Variable(paste0(prefix, "seq.mask"))
  
  if (!is.null(num_embed)) {
    data <- mx.symbol.Embedding(data = data, input_dim = input_size,
                                weight=embed.weight, output_dim = num_embed, name = paste0(prefix, "embed"))
  }
  
  last.hidden <- list()
  last.states <- list()
  
  for (i in 1:num_rnn_layer) {
    
    hidden <- list()
    param <- param.cells[[i]]
    
    in_proj <- mx.symbol.FullyConnected(data = data, weight = param$input.weight, bias = param$input.bias, flatten = F,
                                        num_hidden = num_hidden, name = paste0(prefix, "input", "_l", i))
    
    mem <- mx.symbol.FullyConnected(data = data, weight = param$mem.weight, bias = param$mem.bias, flatten = F,
                                        num_hidden = num_hidden, name = paste0(prefix, "mem", "_l", i))
    
    write.in <- mx.symbol.FullyConnected(data = data, weight = param$write.in.weight, bias = param$write.in.bias, flatten = F,
                                      num_hidden = num_hidden, name = paste0(prefix, "write.in", "_l", i))
    
    read.in <- mx.symbol.FullyConnected(data = data, weight = param$read.in.weight, bias = param$read.in.bias, flatten = F,
                                     num_hidden = num_hidden, name = paste0(prefix, "read.in", "_l", i))
    
    highway <- mx.symbol.FullyConnected(data = data, weight = param$highway.weight, bias = param$highway.bias, flatten = F,
                                        num_hidden = num_hidden, name = paste0(prefix, "highway", "_l", i))
    
    # write.in <- mx.symbol.sigmoid(write.in)
    # read.in <- mx.symbol.tanh(read.in)
    highway <- mx.symbol.sigmoid(highway)
    
    in_proj <- mx.symbol.split(in_proj, num_outputs = seq_len, axis=1, squeeze.axis = F)
    mem <- mx.symbol.split(mem, num_outputs = seq_len, axis=1, squeeze.axis = F)
    write.in <- mx.symbol.split(write.in, num_outputs = seq_len, axis=1, squeeze.axis = F)
    read.in <- mx.symbol.split(read.in, num_outputs = seq_len, axis=1, squeeze.axis = F)
    highway <- mx.symbol.split(highway, num_outputs = seq_len, axis=1, squeeze.axis = F)
    
    for (seqidx in 1:seq_len) {
      
      if (seqidx == 1) {
        write <- mx.symbol.sigmoid(write.in[[seqidx]])
        next.c <- mx.symbol.relu(write.in[[seqidx]] * mem[[seqidx]])
        next.h <- highway[[seqidx]] * in_proj[[seqidx]]
      } else {
        
        read.c <- mx.symbol.FullyConnected(data = prev.c, weight = param$read.c.weight, bias = param$read.c.bias, flatten = F,
                                           num_hidden = num_hidden, name = paste0(prefix, "read.c", "_l", i, "_s", seqidx))
        
        write.c <- mx.symbol.FullyConnected(data = prev.c, weight = param$write.c.weight, bias = param$write.c.bias, flatten = F,
                                            num_hidden = num_hidden, name = paste0(prefix, "write.c", "_l", i, "_s", seqidx))
        
        # read.c <- mx.symbol.tanh(read.c)
        read <- mx.symbol.tanh(read.in[[seqidx]] + read.c)
        write <- mx.symbol.sigmoid(write.in[[seqidx]] + write.c)
        
        next.c <- mx.symbol.relu(prev.c + write * mem[[seqidx]])
        next.h <- read * prev.c + highway[[seqidx]] * in_proj[[seqidx]]
      }
      
      prev.c <- next.c
      hidden <- c(hidden, next.h)
    }
    
    # concat hidden units - concat seq_len blocks of dimension num_hidden x batch.size -> 
    data <- mx.symbol.concat(hidden, num.args=seq_len, dim=1)
    
  }
  
  # if (output_last_state) {
  #   out.states = mx.symbol.Group(unlist(last.states))
  # }
  
  if (config=="seq-to-one") {
    
    if (masking) mask <- mx.symbol.SequenceLast(data=data, use.sequence.length = T, sequence_length = seq.mask, axis = 1, name = paste0(prefix, "mask")) else
      mask <- mx.symbol.SequenceLast(data=data, use.sequence.length = F, axis = 1, name = paste0(prefix, "mask"))
    
    if (!is.null(loss_output)) {
      
      decode <- mx.symbol.FullyConnected(data = mask,
                                         weight = cls.weight,
                                         bias = cls.bias,
                                         num_hidden = num_decode,
                                         name = paste0(prefix, "decode"))
      
      out <- switch(loss_output,
                    softmax = mx.symbol.SoftmaxOutput(data=decode, label=label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, name = paste0(prefix, "loss")),
                    linear = mx.symbol.LinearRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    logistic = mx.symbol.LogisticRegressionOutput(data=decode, label=label, paste0(prefix, name = "loss")),
                    MAE = mx.symbol.MAERegressionOutput(data=decode, label=label, paste0(prefix, name = "loss"))
      )
    } else out <- mask
    
  } else if (config=="one-to-one"){
    
    if (masking) mask <- mx.symbol.SequenceMask(data = data, use.sequence.length = T, sequence_length = seq.mask, axis = 1, value = 0, name = paste0(prefix, "mask")) else
      mask <- mx.symbol.identity(data = data, name = paste0(prefix, "mask"))
    
    if (!is.null(loss_output)) {
      
      mask <- mx.symbol.reshape(data = mask, shape = c(0, -1), reverse = TRUE)
      label <- mx.symbol.reshape(data = label, shape = c(-1))
      
      decode <- mx.symbol.FullyConnected(data = mask, weight = cls.weight, bias = cls.bias, num_hidden = num_decode, 
                                         flatten = T, name = paste0(prefix, "decode"))
      
      out <- switch(loss_output,
                    softmax = mx.symbol.SoftmaxOutput(data=decode, label=label, use_ignore = !ignore_label == -1, ignore_label = ignore_label, 
                                                      name = paste0(prefix, "loss")),
                    linear = mx.symbol.LinearRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    logistic = mx.symbol.LogisticRegressionOutput(data=decode, label=label, name = paste0(prefix, "loss")),
                    MAE = mx.symbol.MAERegressionOutput(data=decode, label=label, name = paste0(prefix, "loss"))
      )
    } else out <- mask
  }
  
  if (output_last_state) {
    return(mx.symbol.Group(c(out, out.states)))
  } else return(out)
}