code: def fill_memory_slot(memory, value, index):
    
    mask = tf.to_float(
        tf.one_hot(index,
                   tf.shape(memory)[0])[:, None, None, None])
    fill_memory = (1 - mask) * memory + mask * value[None, ...]
    return fill_memory
docs: Fills the memory slot at a particular index with the given value.
  
    Args:
      memory: a 4-d tensor [memory_size, batch, length, channel] containing
        the state of all steps
      value: a 3-d tensor [batch, length, channel] as the sate
      index: integer in [0, memory_size)
  
    Returns:
      filled memory