class_docstr: Basic next-frame model with a tiny discrete latent.
docstr: Update the internal states early in the network in GRU-like way.
docstr_tok: ["Update","the","internal","states","early","in","the","network","in","GRU-like","way","."]
code: def update_internal_states_early(self, internal_states, frames):
      """Update the internal states early in the network in GRU-like way."""
      batch_size = common_layers.shape_list(frames[0])[0]
      internal_state = internal_states[0][0][:batch_size, :, :, :]
      state_activation = tf.concat([internal_state, frames[0]], axis=-1)
      state_gate_candidate = tf.layers.conv2d(
          state_activation,
          2 * self.hparams.recurrent_state_size,
          (3, 3),
          padding="SAME",
          name="state_conv",
      )
      state_gate, state_candidate = tf.split(
          state_gate_candidate, 2, axis=-1
      )
      state_gate = tf.nn.sigmoid(state_gate)
      state_candidate = tf.tanh(state_candidate)
      internal_state = internal_state * state_gate
      internal_state += state_candidate * (1.0 - state_gate)
      max_batch_size = max(_MAX_BATCH, self.hparams.batch_size)
      diff_batch_size = max_batch_size - batch_size
      internal_state = tf.pad(
          internal_state, [[0, diff_batch_size], [0, 0], [0, 0], [0, 0]]
      )
      return [[internal_state]]
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import common_video
  from tensor2tensor.layers import discretization
  from tensor2tensor.models.video import base_vae
  from tensor2tensor.models.video import basic_deterministic
  from tensor2tensor.models.video import basic_deterministic_params
  from tensor2tensor.utils import registry
  import tensorflow.compat.v1
label_desc: Use image node_id 6 for calling the NextFrameBasicStochasticDiscrete obj's underlying member method code with example usage: obj.update_internal_states_early(internal_states, frames) and returns: list