def sru(x,
        num_layers=2,
        activation=None,
        initial_state=None,
        name=None,
        reuse=None):
  """SRU cell as in https://arxiv.org/abs/1709.02755.

  As defined in the paper:
  (1) x'_t = W x_t
  (2) f_t = sigmoid(Wf x_t + bf)
  (3) r_t = sigmoid(Wr x_t + br)
  (4) c_t = f_t * c_{t-1} + (1 - f_t) * x'_t
  (5) h_t = r_t * activation(c_t) + (1 - r_t) * x_t

  This version uses functional ops to be faster on GPUs with TF-1.9+.

  Args:
    x: A tensor of shape [batch, ..., channels] ; ... is treated as time.
    num_layers: How many SRU layers; default is 2 as results for 1 disappoint.
    activation: Optional activation function, try tf.nn.tanh or tf.nn.relu.
    initial_state: Optional initial c-state, set to zeros if None.
    name: Optional name, "sru" by default.
    reuse: Optional reuse.

  Returns:
    A tensor of the same shape as x.

  Raises:
    ValueError: if num_layers is not positive.
  """
  if num_layers < 1:
    raise ValueError("Number of layers must be positive: %d" % num_layers)
  if is_xla_compiled():  # On TPU the XLA does a good job with while.
    return sru_with_scan(x, num_layers, activation, initial_state, name, reuse)
  try:
    from tensorflow.contrib.recurrent.python.ops import functional_rnn  # pylint: disable=g-import-not-at-top
  except ImportError:
    tf.logging.info("functional_rnn not found, using sru_with_scan instead")
    return sru_with_scan(x, num_layers, activation, initial_state, name, reuse)

  with tf.variable_scope(name, default_name="sru", values=[x], reuse=reuse):
    # We assume x is [batch, ..., channels] and treat all ... as time.
    x_shape = shape_list(x)
    x = tf.reshape(x, [x_shape[0], -1, x_shape[-1]])
    initial_state = initial_state or tf.zeros([x_shape[0], x_shape[-1]])
    cell = CumsumprodCell(initial_state)
    # Calculate SRU on each layer.
    for i in range(num_layers):
      # The parallel part of the SRU.
      x_orig = x
      x, f, r = tf.split(
          layers().Dense(3 * x_shape[-1], name="kernel_%d" % i)(x), 3, axis=-1)
      f, r = tf.sigmoid(f), tf.sigmoid(r)
      x_times_one_minus_f = x * (1.0 - f)  # Compute in parallel for speed.
      # Calculate states.
      concat = tf.concat([x_times_one_minus_f, f], axis=-1)
      c_states, _ = functional_rnn.functional_rnn(
          cell, concat, time_major=False)
      # Final output.
      if activation is not None:
        c_states = activation(c_states)
      h = c_states * r + (1.0 - r) * x_orig
      x = h  # Next layer.
    return tf.reshape(x, x_shape)
def sru ( x , num_layers = 2 , activation = None , initial_state = None , name = None , reuse = None ) : if num_layers < 1 : raise ValueError ( "Number of layers must be positive: %d" % num_layers ) if is_xla_compiled ( ) : return sru_with_scan ( x , num_layers , activation , initial_state , name , reuse ) try : from tensorflow . contrib . recurrent . python . ops import functional_rnn except ImportError : tf . logging . info ( "functional_rnn not found, using sru_with_scan instead" ) return sru_with_scan ( x , num_layers , activation , initial_state , name , reuse ) with tf . variable_scope ( name , default_name = "sru" , values = [ x ] , reuse = reuse ) : x_shape = shape_list ( x ) x = tf . reshape ( x , [ x_shape [ 0 ] , - 1 , x_shape [ - 1 ] ] ) initial_state = initial_state or tf . zeros ( [ x_shape [ 0 ] , x_shape [ - 1 ] ] ) cell = CumsumprodCell ( initial_state ) for i in range ( num_layers ) : x_orig = x x , f , r = tf . split ( layers ( ) . Dense ( 3 * x_shape [ - 1 ] , name = "kernel_%d" % i ) ( x ) , 3 , axis = - 1 ) f , r = tf . sigmoid ( f ) , tf . sigmoid ( r ) x_times_one_minus_f = x * ( 1.0 - f ) concat = tf . concat ( [ x_times_one_minus_f , f ] , axis = - 1 ) c_states , _ = functional_rnn . functional_rnn ( cell , concat , time_major = False ) if activation is not None : c_states = activation ( c_states ) h = c_states * r + ( 1.0 - r ) * x_orig x = h return tf . reshape ( x , x_shape )
SRU cell as in https://arxiv.org/abs/1709.02755.

  As defined in the paper:
  (1) x'_t = W x_t
  (2) f_t = sigmoid(Wf x_t + bf)
  (3) r_t = sigmoid(Wr x_t + br)
  (4) c_t = f_t * c_{t-1} + (1 - f_t) * x'_t
  (5) h_t = r_t * activation(c_t) + (1 - r_t) * x_t

  This version uses functional ops to be faster on GPUs with TF-1.9+.

  Args:
    x: A tensor of shape [batch, ..., channels] ; ... is treated as time.
    num_layers: How many SRU layers; default is 2 as results for 1 disappoint.
    activation: Optional activation function, try tf.nn.tanh or tf.nn.relu.
    initial_state: Optional initial c-state, set to zeros if None.
    name: Optional name, "sru" by default.
    reuse: Optional reuse.

  Returns:
    A tensor of the same shape as x.

  Raises:
    ValueError: if num_layers is not positive.