def gru_feedfwd(a_t, h_prev, filters, name=None):
  """position-wise Feed-fwd GRU gates following the MPNN.

  Args:
    a_t: Tensor of shape [batch, length, depth] of current input
    h_prev: Tensor of shape [batch, length, depth] of prev input
    filters: an integer specifying number of dimensions of the filters
    name: A string
  Returns:
    h_t: [batch, length, filters] hidden state
  """

  with tf.variable_scope(name, default_name="GRU", values=[a_t, h_prev]):
    # we use right matrix multiplication to handle batches
    # W_z and W_r have shape 2d, d. U_z U_r have shape d,d
    z_t = (
        tf.sigmoid(
            tpu_conv1d(a_t, filters, 1, padding="SAME", name="W_z") +
            tpu_conv1d(h_prev, filters, 1, padding="SAME", name="U_z")))
    r_t = (
        tf.sigmoid(
            tpu_conv1d(a_t, filters, 1, padding="SAME", name="W_r") +
            tpu_conv1d(h_prev, filters, 1, padding="SAME", name="U_r")))
    h_tilde = (
        tf.tanh(
            tpu_conv1d(a_t, filters, 1, padding="SAME", name="W") +
            tpu_conv1d(r_t * h_prev, filters, 1, padding="SAME", name="U")))
    h_t = (1. - z_t) * h_prev + z_t * h_tilde

  return h_t
def gru_feedfwd ( a_t , h_prev , filters , name = None ) : with tf . variable_scope ( name , default_name = "GRU" , values = [ a_t , h_prev ] ) : z_t = ( tf . sigmoid ( tpu_conv1d ( a_t , filters , 1 , padding = "SAME" , name = "W_z" ) + tpu_conv1d ( h_prev , filters , 1 , padding = "SAME" , name = "U_z" ) ) ) r_t = ( tf . sigmoid ( tpu_conv1d ( a_t , filters , 1 , padding = "SAME" , name = "W_r" ) + tpu_conv1d ( h_prev , filters , 1 , padding = "SAME" , name = "U_r" ) ) ) h_tilde = ( tf . tanh ( tpu_conv1d ( a_t , filters , 1 , padding = "SAME" , name = "W" ) + tpu_conv1d ( r_t * h_prev , filters , 1 , padding = "SAME" , name = "U" ) ) ) h_t = ( 1. - z_t ) * h_prev + z_t * h_tilde return h_t
position-wise Feed-fwd GRU gates following the MPNN.

  Args:
    a_t: Tensor of shape [batch, length, depth] of current input
    h_prev: Tensor of shape [batch, length, depth] of prev input
    filters: an integer specifying number of dimensions of the filters
    name: A string
  Returns:
    h_t: [batch, length, filters] hidden state