def __init__(self, batch_size, key_depth, val_depth, memory_size,
               sharpen_factor=1., name="neural_memory"):
    """Initialize the memory object.

    Args:
      batch_size: the batch size.
      key_depth: the depth of the memory keys.
      val_depth: the depth of the memory values.
      memory_size: the number of items in the memory.
      sharpen_factor: the sharpen_factor for addressing the memory.
      name: the optional variable scope.
    """
    self.name = name
    self.batch_size = batch_size
    self.key_depth = key_depth
    self.val_depth = val_depth
    self.memory_size = memory_size
    self.sharpen_factor = sharpen_factor
    with tf.variable_scope(name):
      self.segment_number = tf.get_variable(
          "segment_number", [self.batch_size],
          dtype=tf.int32, trainable=False,
          initializer=tf.constant_initializer(100000))
      self.mem_vals = tf.get_variable(
          "memvals", [self.batch_size, self.memory_size, self.val_depth],
          dtype=tf.float32, trainable=False,
          initializer=tf.constant_initializer(.0))
      self.mean_logits = tf.get_variable(
          "meanlogits", [self.batch_size, self.memory_size],
          dtype=tf.float32, trainable=False,
          initializer=tf.constant_initializer(.0))
Initialize the memory object.

    Args:
      batch_size: the batch size.
      key_depth: the depth of the memory keys.
      val_depth: the depth of the memory values.
      memory_size: the number of items in the memory.
      sharpen_factor: the sharpen_factor for addressing the memory.
      name: the optional variable scope.
    
tensor2tensor/layers/transformer_memory.py