(n_units, mode='train')
return cb.Serial(
      cb.Branch([], zero_state),
      cb.Scan(GRUCell(n_units=n_units), axis=1, mode=mode),
      cb.Select([0], n_in=2),  # Drop RNN state.
      # Set the name to GRU and don't print sublayers.
      name=f'GRU_{n_units}', sublayers_to_print=[]
  )
GRU running on axis 1.
GRU running on axis 1.
GRU running on axis 1 .
def GRU(n_units, mode='train'):
  """GRU running on axis 1."""
  zero_state = MakeZeroState(depth_multiplier=1)  # pylint: disable=no-value-for-parameter
  return cb.Serial(
      cb.Branch([], zero_state),
      cb.Scan(GRUCell(n_units=n_units), axis=1, mode=mode),
      cb.Select([0], n_in=2),  # Drop RNN state.
      # Set the name to GRU and don't print sublayers.
      name=f'GRU_{n_units}', sublayers_to_print=[]
  )
def GRU ( n_units , mode = 'train' ) : zero_state = MakeZeroState ( depth_multiplier = 1 ) # pylint: disable=no-value-for-parameter return cb . Serial ( cb . Branch ( [ ] , zero_state ) , cb . Scan ( GRUCell ( n_units = n_units ) , axis = 1 , mode = mode ) , cb . Select ( [ 0 ] , n_in = 2 ) , # Drop RNN state. # Set the name to GRU and don't print sublayers. name = f'GRU_{n_units}' , sublayers_to_print = [ ] )