Apply a stepwise model over a sequence, maintaining state. For RNNs
160598
def Recurrent(step_model):
    """"""
    ops = step_model.ops

    def recurrent_fwd(seqs, drop=0.0):
        lengths = [len(X) for X in seqs]
        X, size_at_t, unpad = ops.square_sequences(seqs)
        Y = ops.allocate((X.shape[0], X.shape[1], step_model.nO))
        cell_drop = ops.get_dropout_mask((len(seqs), step_model.nO), 0.0)
        hidden_drop = ops.get_dropout_mask((len(seqs), step_model.nO), 0.0)
        out_drop = ops.get_dropout_mask((len(seqs), step_model.nO), 0.0)
        backprops = [None] * max(lengths)
        state = step_model.weights.get_initial_state(len(seqs))
        for t in range(max(lengths)):
            state = list(state)
            size = size_at_t[t]
            Xt = X[t, :size]
            state[0] = state[0][:size]
            state[1] = state[1][:size]
            if cell_drop is not None:
                state[0] *= cell_drop
            if hidden_drop is not None:
                state[1] *= hidden_drop
            inputs = (state, Xt)
            (state, Y[t, :size]), backprops[t] = step_model.begin_update(inputs)
            if out_drop is not None:
                Y[t, :size] *= out_drop
        outputs = unpad(Y)

        def recurrent_bwd(d_outputs, sgd=None):
            dY, size_at_t, unpad = step_model.ops.square_sequences(d_outputs)
            d_state = [
                step_model.ops.allocate((dY.shape[1], step_model.nO)),
                step_model.ops.allocate((dY.shape[1], step_model.nO)),
            ]
            updates = {}

            def gather_updates(weights, gradient, key=None):
                updates[key] = (weights, gradient)

            dX = step_model.ops.allocate(
                (dY.shape[0], dY.shape[1], step_model.weights.nI)
            )
            for t in range(max(lengths) - 1, -1, -1):
                if out_drop is not None:
                    dY[t] *= out_drop
                d_state_t, dXt = backprops[t]((d_state, dY[t]), sgd=gather_updates)
                d_state[0][: d_state_t[0].shape[0]] = d_state_t[0]
                d_state[1][: d_state_t[1].shape[0]] = d_state_t[1]
                dX[t, : dXt.shape[0]] = dXt
                if cell_drop is not None:
                    d_state[0] *= cell_drop
                if hidden_drop is not None:
                    d_state[1] *= hidden_drop
            d_cell, d_hidden = d_state
            step_model.weights.d_initial_cells += d_cell.sum(axis=0)
            step_model.weights.d_initial_hiddens += d_hidden.sum(axis=0)
            if sgd is not None:
                for key, (weights, gradient) in updates.items():
                    sgd(weights, gradient, key=key)
            return unpad(dX)

        return outputs, recurrent_bwd

    model = wrap(recurrent_fwd, step_model)
    model.nO = step_model.nO
    return model
424823
def forward(self, X):
        """"""
        # Identify the first non-zero integer from the right (i.e., the length
        # of the sequence before padding starts).
        batch_size, max_seq = X.shape[0], X.shape[1]
        seq_lengths = torch.zeros(batch_size, dtype=torch.long)
        for i in range(batch_size):
            for j in range(max_seq - 1, -1, -1):
                if not torch.all(X[i, j] == 0):
                    seq_lengths[i] = j + 1
                    break

        # Sort by length because pack_padded_sequence requires it
        # Save original order to restore before returning
        seq_lengths, perm_idx = seq_lengths.sort(0, descending=True)
        X = X[perm_idx]
        inv_perm_idx = torch.tensor(
            [i for i, _ in sorted(enumerate(perm_idx), key=lambda idx: idx[1])],
            dtype=torch.long,
        )

        # Encode and pack input sequence
        X_packed = rnn_utils.pack_padded_sequence(
            self.encoder.encode(X), seq_lengths, batch_first=True
        )

        # Run LSTM
        outputs, (h_t, c_t) = self.lstm(X_packed)

        # Unpack and reduce outputs
        outputs_unpacked, _ = rnn_utils.pad_packed_sequence(outputs, batch_first=True)
        reduced = self._reduce_output(outputs_unpacked, seq_lengths)
        return reduced[inv_perm_idx, :]
975116
def hybrid_forward(self, F, x, sampled_values, label, weight, bias):
        """"""
        sampled_candidates, _, _ = sampled_values
        # (batch_size,)
        label = F.reshape(label, shape=(-1,))
        # (num_sampled+batch_size,)
        ids = F.concat(sampled_candidates, label, dim=0)
        # lookup weights and biases
        # (num_sampled+batch_size, dim)
        w_all = F.Embedding(data=ids, weight=weight,
                            input_dim=self._num_classes, output_dim=self._in_unit,
                            sparse_grad=self._sparse_grad)
        # (num_sampled+batch_size, 1)
        b_all = F.take(bias, indices=ids)
        return self._dense(x, sampled_values, label, w_all, b_all)
1171227
def construct_model(images,
                    actions=None,
                    context_frames=2,
                    hparams=None,
                    is_training=True):
  """"""

  pred_depth = 20

  enc_out_all, pred_out_all, van_out_all, van_on_enc_all = [], [], [], []

  lstm_states = [None] * (pred_depth + 2)

  enc_out = encoder_vgg(
      images[0], hparams.enc_size, False, scope_prefix='timestep/',
      hparams=hparams, is_training=is_training)
  enc_out = tf.identity(enc_out, 'enc_out')
  enc_out_all.append(enc_out)

  num_timesteps = len(actions) - 1
  sum_freq = int(num_timesteps / 4 + 1)

  reuse = False
  for timestep, action in zip(range(len(actions) - 1), actions[:-1]):
    done_warm_start = timestep > context_frames - 1

    with tf.variable_scope('timestep', reuse=reuse):
      if done_warm_start:
        pred_input = pred_out_all[-1]
      else:
        pred_input = enc_out_all[-1]
      pred_out = predictor(
          pred_input, action, lstm_states, pred_depth, False, hparams=hparams)
      pred_out = tf.identity(pred_out, 'pred_out')
      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
        tf.summary.histogram('pred_out', pred_out)
      pred_out_all.append(pred_out)

      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
        tf.summary.histogram('lstm_state', lstm_states[0])
      van_out, _, _ = van(
          enc_out_all[0],
          images[0],
          pred_out,
          images[timestep + 1],
          tf.AUTO_REUSE,
          hparams=hparams)
      van_out = tf.identity(van_out, 'van_out')
      van_out_all.append(van_out)

      enc_out = encoder_vgg(
          images[timestep + 1], hparams.enc_size, True, hparams=hparams,
          is_training=is_training)
      enc_out = tf.identity(enc_out, 'enc_out')
      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
        tf.summary.histogram('enc_out', enc_out)
      enc_out_all.append(enc_out)

      van_input = images[0]
      enc_noise = tf.zeros_like(enc_out)
      if timestep % sum_freq == 0:  # and not hparams.use_tpu:
        tf.summary.histogram('enc_noise', enc_noise)
      van_on_enc, _, _ = van(
          enc_out_all[0],
          van_input,
          enc_out + enc_noise,
          images[timestep + 1],
          tf.AUTO_REUSE,
          hparams=hparams)
      van_on_enc = tf.identity(van_on_enc, 'van_on_enc')
      van_on_enc_all.append(van_on_enc)

      reuse = True

  return enc_out_all, pred_out_all, van_out_all, van_on_enc_all
1010714
def update_learning_rate(lr, trainer, epoch, ratio, steps):
    """"""
    new_lr = lr * (ratio ** int(np.sum(np.array(steps) < epoch)))
    trainer.set_learning_rate(new_lr)
    return trainer
148248
def _get_rnn_layer(mode, num_layers, input_size, hidden_size, dropout, weight_dropout):
    """"""
    if mode == 'rnn_relu':
        rnn_block = functools.partial(rnn.RNN, activation='relu')
    elif mode == 'rnn_tanh':
        rnn_block = functools.partial(rnn.RNN, activation='tanh')
    elif mode == 'lstm':
        rnn_block = rnn.LSTM
    elif mode == 'gru':
        rnn_block = rnn.GRU

    block = rnn_block(hidden_size, num_layers, dropout=dropout,
                      input_size=input_size)

    if weight_dropout:
        apply_weight_drop(block, '.*h2h_weight', rate=weight_dropout)

    return block
903880
def ustep(self):
        """"""

        self.U += self.rsdl_r(self.AX, self.Y)
422420
def _after_n_epoch(self, epoch_id: int, **_) -> None:
        
        SaveEvery.save_model(model=self._model, name_suffix=str(epoch_id), on_failure=self._on_save_failure)
1035153
public static void step(long seed, Neurons[] neurons, DeepLearningModel.DeepLearningModelInfo minfo, boolean training, double[] responses) {
    try {
      for (int i=1; i<neurons.length-1; ++i) {
        neurons[i].fprop(seed, training);
      }
      if (minfo.get_params().autoencoder) {
        neurons[neurons.length - 1].fprop(seed, training);
        if (training) {
          for (int i=neurons.length-1; i>0; --i) {
            neurons[i].bprop();
          }
        }
      } else {
        if (minfo.get_params().classification) {
          ((Neurons.Softmax) neurons[neurons.length - 1]).fprop();
          if (training) {
            for (int i = 1; i < neurons.length - 1; i++)
              Arrays.fill(neurons[i]._e.raw(), 0);
            int target_label;
            if (Double.isNaN(responses[0])) { //missing response
              target_label = Neurons.missing_int_value;
            } else {
              assert ((double) (int) responses[0] == responses[0]); //classification -> integer labels expected
              target_label = (int) responses[0];
            }
            ((Neurons.Softmax) neurons[neurons.length - 1]).bprop(target_label);
          }
        } else {
          ((Neurons.Linear) neurons[neurons.length - 1]).fprop();
          if (training) {
            for (int i = 1; i < neurons.length - 1; i++)
              Arrays.fill(neurons[i]._e.raw(), 0);
            float target_value;
            if (Double.isNaN(responses[0])) { //missing response
              target_value = Neurons.missing_real_value;
            } else {
              target_value = (float) responses[0];
            }
            ((Neurons.Linear) neurons[neurons.length - 1]).bprop(target_value);
          }
        }
        if (training) {
          for (int i=neurons.length-2; i>0; --i)
            neurons[i].bprop();
        }
      }
    }
    catch(RuntimeException ex) {
      Log.warn(ex.getMessage());
      minfo.set_unstable();
      throw new Job.JobCancelledException("Canceling job due to numerical instability.");
    }
  }
462997
def forward(self, data_batch, is_train=None, carry_state=True):
        
        # propagate states from the previous iteration
        if carry_state:
            if isinstance(self._next_states, (int, float)):
                self._module.set_states(value=self._next_states)
            else:
                self._module.set_states(states=self._next_states)
        self._module.forward(data_batch, is_train=is_train)
        outputs = self._module.get_outputs(merge_multi_context=False)
        self._next_states = outputs[:-1]
474485
def rlmb_tiny_recurrent():
  """"""
  hparams = rlmb_ppo_tiny()
  hparams.epochs = 1  # Too slow with 2 for regular runs.
  hparams.generative_model = "next_frame_basic_recurrent"
  hparams.generative_model_params = "next_frame_basic_recurrent"
  return hparams
90305
def _euler_step(self, dt):
        
        self.state = self.state + self.ddt() * dt
1005076
def lstm_seq2seq_internal(inputs, targets, hparams, train):
  """"""
  with tf.variable_scope("lstm_seq2seq"):
    if inputs is not None:
      inputs_length = common_layers.length_from_embedding(inputs)
      # Flatten inputs.
      inputs = common_layers.flatten4d3d(inputs)

      # LSTM encoder.
      inputs = tf.reverse_sequence(inputs, inputs_length, seq_axis=1)
      _, final_encoder_state = lstm(inputs, inputs_length, hparams, train,
                                    "encoder")
    else:
      final_encoder_state = None

    # LSTM decoder.
    shifted_targets = common_layers.shift_right(targets)
    # Add 1 to account for the padding added to the left from shift_right
    targets_length = common_layers.length_from_embedding(shifted_targets) + 1
    decoder_outputs, _ = lstm(
        common_layers.flatten4d3d(shifted_targets),
        targets_length,
        hparams,
        train,
        "decoder",
        initial_state=final_encoder_state)
    return tf.expand_dims(decoder_outputs, axis=2)
471004
def BiLSTM(nO, nI):
    """"""
    return Bidirectional(LSTM(nO // 2, nI), LSTM(nO // 2, nI))
997453
def create(input_block: ModelFactory, rnn_type: str, hidden_layers: typing.List[int],
           output_dim: int, dropout=0.0):
    
    def instantiate(**_):
        return MultilayerRnnSequenceModel(
            input_block.instantiate(), rnn_type=rnn_type, hidden_layers=hidden_layers, output_dim=output_dim, dropout=dropout
        )

    return ModelFactory.generic(instantiate)
1279045
def forward(self, inputs, lengths):
        
        x = self.embedder(inputs)

        # bidirectional layer
        x = self.dropout(x)
        x = pack_padded_sequence(x, lengths.cpu().numpy(),
                                 batch_first=self.batch_first)
        x, _ = self.rnn_layers[0](x)
        x, _ = pad_packed_sequence(x, batch_first=self.batch_first)

        # 1st unidirectional layer
        x = self.dropout(x)
        x, _ = self.rnn_layers[1](x)

        # the rest of unidirectional layers,
        # with residual connections starting from 3rd layer
        for i in range(2, len(self.rnn_layers)):
            residual = x
            x = self.dropout(x)
            x, _ = self.rnn_layers[i](x)
            x = x + residual

        return x
293284
def build_sequence(self, xs, masks, init, is_left_to_right):
        '''
        
        '''
        states = []
        last = init
        if is_left_to_right:
            for i, xs_i in enumerate(xs):
                h = self.build(xs_i, last, masks[i])
                states.append(h)
                last = h
        else:
            for i in range(len(xs) - 1, -1, -1):
                h = self.build(xs[i], last, masks[i])
                states.insert(0, h)
                last = h
        return states