text: # coding=utf-8
  # Copyright 2021 The Google Research Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Utility functions for constructing networks."""
  
  import tensorflow as tf
  
  
  def _merge_inputs(shape, name, merged_embedding, inputs):
    rnn_input = tf.keras.layers.Input(shape=shape, name=name)
    merged_embedding = tf.keras.layers.concatenate([merged_embedding, rnn_input],
                                                   axis=-1)
    inputs.append(rnn_input)
    return merged_embedding, inputs
  
  
  def construct_creator_rnn_inputs(document_feature_size=10,
                                   creator_feature_size=None,
                                   num_creators=None,
                                   creator_id_embedding_size=0,
                                   trajectory_length=None):
    """Return creator RNN inputs.
  
    Creator RNN input at each time step is the current time step's
    [#recommendations, #user-clicks, summed user_reward,
    weighted_clicked_document_topics, creator_current_feature], where
    'weighted_clicked_document_topics` has length document_feature_size, and is a
    weighted average of user_clicked document topics weighted by user_reward;
    `creator_current_feature` is creator current context such as the creator
    observed saturated satisfaction if creator_feature_size is not None.
  
    Args:
      document_feature_size: Integer, length of document features, which is the
        number of topics on the platform.
      creator_feature_size: Integer or None, length of creator features. If None,
        no features about creators will be used as input.
      num_creators: Integer, used for embedding creator id.
      creator_id_embedding_size: Integer, if greater than zero, embed creator id.
      trajectory_length: Integer, specify the trajectory length.
  
    Returns:
      rnn_inputs: Keras input layer of shape (episode_length,
      3 + document_feature_size+creator_feature_size)
    """
    number_recommendations = tf.keras.layers.Input(
        shape=(trajectory_length, 1), name='creator_number_recommendations')
    number_clicks = tf.keras.layers.Input(
        shape=(trajectory_length, 1), name='creator_number_clicks')
    creator_rewards = tf.keras.layers.Input(
        shape=(trajectory_length, 1), name='creator_user_rewards')
    weighted_clicked_doc_topics = tf.keras.layers.Input(
        shape=(trajectory_length, document_feature_size),
        name='creator_weighted_clicked_doc_topics')
    merged_embedding = tf.keras.layers.concatenate([
        number_recommendations, number_clicks, creator_rewards,
        weighted_clicked_doc_topics
    ],
                                                   axis=-1)
    inputs = [
        number_recommendations, number_clicks, creator_rewards,
        weighted_clicked_doc_topics
    ]
    if creator_feature_size is not None:
      merged_embedding, inputs = _merge_inputs(
          (trajectory_length, creator_feature_size), 'creator_current_feature',
          merged_embedding, inputs)
    if creator_id_embedding_size > 0:
      creator_id_inputs = tf.keras.layers.Input(
          shape=(trajectory_length), name='creator_id')
      creator_embeddings = tf.keras.layers.Embedding(
          input_dim=num_creators + 1,  # One additional token for padding.
          output_dim=creator_id_embedding_size,
          mask_zero=False,
          name='creator_id_embedding_layer')(
              creator_id_inputs)
      inputs.append(creator_id_inputs)
      merged_embedding = tf.keras.layers.concatenate(
          [merged_embedding, creator_embeddings], axis=-1)
  
    return merged_embedding, inputs
  
  
  def construct_user_rnn_inputs(document_feature_size=10,
                                creator_feature_size=None,
                                user_feature_size=None,
                                input_reward=False):
    """Returns user RNN inputs.
  
    Args:
      document_feature_size: Integer, length of document features.
      creator_feature_size: Integer or None, length of creator features. If None,
        no features about creators will be input.
      user_feature_size: Integer or None, length of user features. If None, no
        features about users will be input.
      input_reward: Boolean, whether to input previous reward to RNN layer.
    """
  
    # Previous consumed document.
    rnn_input_doc_feature = tf.keras.layers.Input(
        shape=(None, document_feature_size), name='user_consumed_doc_feature')
    merged_embedding = rnn_input_doc_feature
    inputs = [rnn_input_doc_feature]
    # Previous consumed document-associated creator.
    if creator_feature_size is not None:
      # This vector includes creator's observable features and/or creator's hidden
      # states inferred by creator model.
      merged_embedding, inputs = _merge_inputs(
          (None, creator_feature_size), 'user_consumed_doc-creator_feature',
          merged_embedding, inputs)
  
    # User current context.
    if user_feature_size is not None:
      merged_embedding, inputs = _merge_inputs(
          (None, user_feature_size), 'user_current_feature', merged_embedding,
          inputs)
    # Previous reward.
    if input_reward:
      merged_embedding, inputs = _merge_inputs((None, 1), 'user_previous_reward',
                                               merged_embedding, inputs)
  
    return merged_embedding, inputs
  
  
  def construct_rnn_layer(rnn_type, rnn_merged_embedding, hidden_size,
                          regularizer_obj):
    """Returns a recurrent layer for the given type or raises ValueError."""
    if rnn_type == 'LSTM':
      rnn_layer = tf.keras.layers.LSTM(
          units=hidden_size,
          return_sequences=True,
          return_state=True,
          name='LSTM',
          kernel_regularizer=regularizer_obj)
      whole_seq_output, final_memory_state, final_carry_state = rnn_layer(
          rnn_merged_embedding)
      final_state = [final_memory_state, final_carry_state]
      return rnn_layer, whole_seq_output, final_state
    elif rnn_type == 'GRU':
      rnn_layer = tf.keras.layers.GRU(
          units=hidden_size,
          return_sequences=True,
          return_state=True,
          name='GRU',
          kernel_regularizer=regularizer_obj)
      whole_seq_output, final_state = rnn_layer(rnn_merged_embedding)
      return rnn_layer, whole_seq_output, final_state
    else:
      raise NotImplementedError(
          "Use recurrent cell as one out of 'LSTM' and 'GRU'.")
  