text: import tensorflow as tf
  from keras.layers import RepeatVector, Lambda
  
  
  def repeat_vector(inputs):
      """
      Temporary solution:
      Use this function within a Lambda layer to get a repeated layer with a variable 1-st dimension (seq_len).
      May be useful to further feed it to a Concatenate layer.
  
      inputs == (layer_for_repeat, layer_for_getting_rep_num):
          layer_for_repeat:           shape == (batch_size, vector_dim)
          layer_for_getting_rep_num:  shape == (batch_size, seq_len, ...)
      :return:
          repeated layer_for_repeat, shape == (batch_size, seq_len, vector_dim)
      """
      layer_for_repeat, layer_for_getting_rep_num = inputs
      repeated_vector = RepeatVector(
          n=tf.keras.backend.shape(layer_for_getting_rep_num)[1], name='custom_repeat_vector')(layer_for_repeat)
      # shape == (batch_size, seq_len, vector_dim)
      return repeated_vector
  
  
  def softmax_with_temperature(logits, temperature):
      """
      :param logits:      shape == (batch_size, seq_len, vocab_size), float32
      :param temperature: shape == (batch_size, 1), float32
      :return:
          transformed tokens probs, shape == (batch_size, seq_len, vocab_size), float32
      """
  
      def softmax_with_temp(args):
          logits, temperature = args
          repeat_num = tf.keras.backend.shape(logits)[1]
          temperature_repeated = RepeatVector(repeat_num)(temperature)
          # shape == (batch_size, seq_len)
          scaled_logits = logits / temperature_repeated
          # shape == (batch_size, seq_len, vocab_size)
  
          # for numerical stability (e.g. for low temperatures):
          scaled_logits = scaled_logits - tf.keras.backend.max(scaled_logits, axis=2, keepdims=True)
          # shape == (batch_size, seq_len, vocab_size)
          transformed_probs = tf.keras.backend.softmax(scaled_logits)
          # shape == (batch_size, seq_len, vocab_size)
          return transformed_probs
  
      # wrap transformation in Lambda to turn the result to Keras layer
      transformed_probs = Lambda(
          function=softmax_with_temp,
          mask=lambda inputs, inputs_masks: inputs_masks[0],  # function to get mask of the first input
          name='softmax_with_temperature')([logits, temperature])
      # output shape == (batch_size, seq_len, vocab_size)
  
      return transformed_probs
  