Python tensorflow.int32() Examples

def __init__(self, resolution=1024, num_channels=3, dtype='uint8', dynamic_range=[0,255], label_size=0, label_dtype='float32'):
        self.resolution         = resolution
        self.resolution_log2    = int(np.log2(resolution))
        self.shape              = [num_channels, resolution, resolution]
        self.dtype              = dtype
        self.dynamic_range      = dynamic_range
        self.label_size         = label_size
        self.label_dtype        = label_dtype
        self._tf_minibatch_var  = None
        self._tf_lod_var        = None
        self._tf_minibatch_np   = None
        self._tf_labels_np      = None

        assert self.resolution == 2 ** self.resolution_log2
        with tf.name_scope('Dataset'):
            self._tf_minibatch_var = tf.Variable(np.int32(0), name='minibatch_var')
            self._tf_lod_var = tf.Variable(np.int32(0), name='lod_var') 

def apply_with_random_selector(x, func, num_cases):
  """Computes func(x, sel), with sel sampled from [0...num_cases-1].

  Args:
    x: input Tensor.
    func: Python function to apply.
    num_cases: Python int32, number of cases to sample sel from.

  Returns:
    The result of func(x, sel), where func receives the value of the
    selector as a python integer, but sel is sampled dynamically.
  """
  sel = tf.random_uniform([], maxval=num_cases, dtype=tf.int32)
  # Pass the real x only to one of the func calls.
  return control_flow_ops.merge([
      func(control_flow_ops.switch(x, tf.equal(sel, case))[1], case)
      for case in range(num_cases)])[0] 

def batch_of_random_bools(batch_size, n):
  """Return a batch of random "boolean" numbers.

  Args:
    batch_size:  Batch size dimension of returned tensor.
    n:  number of entries per batch.

  Returns:
    A [batch_size, n] tensor of "boolean" numbers, where each number is
    preresented as -1 or 1.
  """

  as_int = tf.random_uniform(
      [batch_size, n], minval=0, maxval=2, dtype=tf.int32)
  expanded_range = (as_int * 2) - 1
  return tf.cast(expanded_range, tf.float32) 

def input_fn(partition, training, batch_size):
    """Generate an input_fn for the Estimator."""

    def _input_fn():
        if partition == "train":
            dataset = tf.data.Dataset.from_generator(
                generator(x_train, y_train), (tf.float32, tf.int32), ((28, 28), ()))
        else:
            dataset = tf.data.Dataset.from_generator(
                generator(x_test, y_test), (tf.float32, tf.int32), ((28, 28), ()))

        if training:
            dataset = dataset.shuffle(10 * batch_size, seed=RANDOM_SEED).repeat()

        dataset = dataset.map(preprocess_image).batch(batch_size)
        iterator = dataset.make_one_shot_iterator()
        features, labels = iterator.get_next()
        return features, labels
    return _input_fn 

def _aspect_preserving_resize(image, smallest_side):
  """Resize images preserving the original aspect ratio.

  Args:
    image: A 3-D image `Tensor`.
    smallest_side: A python integer or scalar `Tensor` indicating the size of
      the smallest side after resize.

  Returns:
    resized_image: A 3-D tensor containing the resized image.
  """
  smallest_side = tf.convert_to_tensor(smallest_side, dtype=tf.int32)

  shape = tf.shape(image)
  height = shape[0]
  width = shape[1]
  new_height, new_width = _smallest_size_at_least(height, width, smallest_side)
  image = tf.expand_dims(image, 0)
  resized_image = tf.image.resize_bilinear(image, [new_height, new_width],
                                           align_corners=False)
  resized_image = tf.squeeze(resized_image)
  resized_image.set_shape([None, None, 3])
  return resized_image 

def pitch_shift(
        spectrogram,
        semitone_shift=0.0,
        method=tf.image.ResizeMethod.BILINEAR):
    """ Pitch shift a spectrogram preserving shape in tensorflow. Note that
    this is an approximation in the frequency domain.

    :param spectrogram: Input spectrogram to be pitch shifted as tensor.
    :param semitone_shift: (Optional) Pitch shift in semitone, default to 0.0.
    :param mehtod: (Optional) Interpolation method, default to BILINEAR.
    :returns: Pitch shifted spectrogram (same shape as spectrogram).
    """
    factor = 2 ** (semitone_shift / 12.)
    T = tf.shape(spectrogram)[0]
    F = tf.shape(spectrogram)[1]
    F_ps = tf.cast(tf.cast(F, tf.float32) * factor, tf.int32)[0]
    ps_spec = tf.image.resize_images(
        spectrogram,
        [T, F_ps],
        method=method,
        align_corners=True)
    paddings = [[0, 0], [0, tf.maximum(0, F - F_ps)], [0, 0]]
    return tf.pad(ps_spec[:, :F, :], paddings, 'CONSTANT') 

def time_stretch(
        spectrogram,
        factor=1.0,
        method=tf.image.ResizeMethod.BILINEAR):
    """ Time stretch a spectrogram preserving shape in tensorflow. Note that
    this is an approximation in the frequency domain.

    :param spectrogram: Input spectrogram to be time stretched as tensor.
    :param factor: (Optional) Time stretch factor, must be >0, default to 1.
    :param mehtod: (Optional) Interpolation method, default to BILINEAR.
    :returns: Time stretched spectrogram as tensor with same shape.
    """
    T = tf.shape(spectrogram)[0]
    T_ts = tf.cast(tf.cast(T, tf.float32) * factor, tf.int32)[0]
    F = tf.shape(spectrogram)[1]
    ts_spec = tf.image.resize_images(
        spectrogram,
        [T_ts, F],
        method=method,
        align_corners=True)
    return tf.image.resize_image_with_crop_or_pad(ts_spec, T, F) 

def test_construct_anchor_grid_unnormalized(self):
    base_anchor_size = tf.constant([1, 1], dtype=tf.float32)
    box_specs_list = [[(1.0, 1.0)]]

    exp_anchor_corners = [[0., 0., 320., 320.], [0., 320., 320., 640.]]

    anchor_generator = ag.MultipleGridAnchorGenerator(box_specs_list,
                                                      base_anchor_size)
    anchors = anchor_generator.generate(
        feature_map_shape_list=[(tf.constant(1, dtype=tf.int32), tf.constant(
            2, dtype=tf.int32))],
        im_height=320,
        im_width=640)
    anchor_corners = anchors.get()

    with self.test_session():
      anchor_corners_out = anchor_corners.eval()
      self.assertAllClose(anchor_corners_out, exp_anchor_corners) 

def _reshape_instance_masks(self, keys_to_tensors):
    """Reshape instance segmentation masks.

    The instance segmentation masks are reshaped to [num_instances, height,
    width] and cast to boolean type to save memory.

    Args:
      keys_to_tensors: a dictionary from keys to tensors.

    Returns:
      A 3-D boolean tensor of shape [num_instances, height, width].
    """
    masks = keys_to_tensors['image/segmentation/object']
    if isinstance(masks, tf.SparseTensor):
      masks = tf.sparse_tensor_to_dense(masks)
    height = keys_to_tensors['image/height']
    width = keys_to_tensors['image/width']
    to_shape = tf.cast(tf.stack([-1, height, width]), tf.int32)

    return tf.cast(tf.reshape(masks, to_shape), tf.bool) 

def expanded_shape(orig_shape, start_dim, num_dims):
  """Inserts multiple ones into a shape vector.

  Inserts an all-1 vector of length num_dims at position start_dim into a shape.
  Can be combined with tf.reshape to generalize tf.expand_dims.

  Args:
    orig_shape: the shape into which the all-1 vector is added (int32 vector)
    start_dim: insertion position (int scalar)
    num_dims: length of the inserted all-1 vector (int scalar)
  Returns:
    An int32 vector of length tf.size(orig_shape) + num_dims.
  """
  with tf.name_scope('ExpandedShape'):
    start_dim = tf.expand_dims(start_dim, 0)  # scalar to rank-1
    before = tf.slice(orig_shape, [0], start_dim)
    add_shape = tf.ones(tf.reshape(num_dims, [1]), dtype=tf.int32)
    after = tf.slice(orig_shape, start_dim, [-1])
    new_shape = tf.concat([before, add_shape, after], 0)
    return new_shape 

def structure(self, input_tensor):
        """
        Args:
            input_tensor: NHWC
        """
        rnd = tf.random_uniform((), 135, 160, dtype=tf.int32)
        rescaled = tf.image.resize_images(
            input_tensor, [rnd, rnd], method=tf.image.ResizeMethod.NEAREST_NEIGHBOR)
        h_rem = 160 - rnd
        w_rem = 160 - rnd
        pad_left = tf.random_uniform((), 0, w_rem, dtype=tf.int32)
        pad_right = w_rem - pad_left
        pad_top = tf.random_uniform((), 0, h_rem, dtype=tf.int32)
        pad_bottom = h_rem - pad_top
        padded = tf.pad(rescaled, [[0, 0], [pad_top, pad_bottom], [
                        pad_left, pad_right], [0, 0]])
        padded.set_shape((input_tensor.shape[0], 160, 160, 3))
        output = tf.cond(tf.random_uniform(shape=[1])[0] < tf.constant(0.9),
                         lambda: padded, lambda: input_tensor)
        return output 

def main():
    dataset = tf.data.Dataset.from_generator(gen, (tf.int32, tf.int32),
                                             (tf.TensorShape([BATCH_SIZE]),
                                              tf.TensorShape([BATCH_SIZE, 1])))
    optimizer = tf.compat.v1.train.GradientDescentOptimizer(LEARNING_RATE)
    model = Word2Vec(vocab_size=VOCAB_SIZE, embed_size=EMBED_SIZE)
    grad_fn = tfe.implicit_value_and_gradients(model.compute_loss)
    total_loss = 0.0
    num_train_steps = 0
    while num_train_steps < NUM_TRAIN_STEPS:
        for center_words, target_words in tfe.Iterator(dataset):
            if num_train_steps >= NUM_TRAIN_STEPS:
                break
            loss_batch, grads = grad_fn(center_words, target_words)
            total_loss += loss_batch
            optimizer.apply_gradients(grads)
            if (num_train_steps + 1) % SKIP_STEP == 0:
                print('Average loss at step {}: {:5.1f}'.format(
                    num_train_steps, total_loss / SKIP_STEP
                ))
                total_loss = 0.0
            num_train_steps += 1 

def get_hash_slots(self, query):
    """Gets hashed-to buckets for batch of queries.

    Args:
      query: 2-d Tensor of query vectors.

    Returns:
      A list of hashed-to buckets for each hash function.
    """

    binary_hash = [
        tf.less(tf.matmul(query, self.hash_vecs[i], transpose_b=True), 0)
        for i in xrange(self.num_libraries)]
    hash_slot_idxs = [
        tf.reduce_sum(
            tf.to_int32(binary_hash[i]) *
            tf.constant([[2 ** i for i in xrange(self.num_hashes)]],
                        dtype=tf.int32), 1)
        for i in xrange(self.num_libraries)]
    return hash_slot_idxs 

def read_analogies(self):
    """Reads through the analogy question file.

    Returns:
      questions: a [n, 4] numpy array containing the analogy question's
                 word ids.
      questions_skipped: questions skipped due to unknown words.
    """
    questions = []
    questions_skipped = 0
    with open(self._options.eval_data, "rb") as analogy_f:
      for line in analogy_f:
        if line.startswith(b":"):  # Skip comments.
          continue
        words = line.strip().lower().split(b" ")
        ids = [self._word2id.get(w.strip()) for w in words]
        if None in ids or len(ids) != 4:
          questions_skipped += 1
        else:
          questions.append(np.array(ids))
    print("Eval analogy file: ", self._options.eval_data)
    print("Questions: ", len(questions))
    print("Skipped: ", questions_skipped)
    self._analogy_questions = np.array(questions, dtype=np.int32) 

def test_pad_or_clip_tensor_using_integer_input(self):
    t1 = tf.constant([1], dtype=tf.int32)
    tt1 = shape_utils.pad_or_clip_tensor(t1, 2)
    t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
    tt2 = shape_utils.pad_or_clip_tensor(t2, 2)

    t3 = tf.constant([1, 2, 3], dtype=tf.int32)
    tt3 = shape_utils.clip_tensor(t3, 2)
    t4 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
    tt4 = shape_utils.clip_tensor(t4, 2)

    self.assertEqual(2, tt1.get_shape()[0])
    self.assertEqual(2, tt2.get_shape()[0])
    self.assertEqual(2, tt3.get_shape()[0])
    self.assertEqual(2, tt4.get_shape()[0])

    with self.test_session() as sess:
      tt1_result, tt2_result, tt3_result, tt4_result = sess.run(
          [tt1, tt2, tt3, tt4])
      self.assertAllEqual([1, 0], tt1_result)
      self.assertAllClose([[0.1, 0.2], [0, 0]], tt2_result)
      self.assertAllEqual([1, 2], tt3_result)
      self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], tt4_result) 

def test_pad_or_clip_tensor_using_tensor_input(self):
    t1 = tf.constant([1], dtype=tf.int32)
    tt1 = shape_utils.pad_or_clip_tensor(t1, tf.constant(2))
    t2 = tf.constant([[0.1, 0.2]], dtype=tf.float32)
    tt2 = shape_utils.pad_or_clip_tensor(t2, tf.constant(2))

    t3 = tf.constant([1, 2, 3], dtype=tf.int32)
    tt3 = shape_utils.clip_tensor(t3, tf.constant(2))
    t4 = tf.constant([[0.1, 0.2], [0.2, 0.4], [0.5, 0.8]], dtype=tf.float32)
    tt4 = shape_utils.clip_tensor(t4, tf.constant(2))

    with self.test_session() as sess:
      tt1_result, tt2_result, tt3_result, tt4_result = sess.run(
          [tt1, tt2, tt3, tt4])
      self.assertAllEqual([1, 0], tt1_result)
      self.assertAllClose([[0.1, 0.2], [0, 0]], tt2_result)
      self.assertAllEqual([1, 2], tt3_result)
      self.assertAllClose([[0.1, 0.2], [0.2, 0.4]], tt4_result) 

def test_normalized_to_image_coordinates(self):
    normalized_boxes = tf.placeholder(tf.float32, shape=(None, 1, 4))
    normalized_boxes_np = np.array([[[0.0, 0.0, 1.0, 1.0]],
                                    [[0.5, 0.5, 1.0, 1.0]]])
    image_shape = tf.convert_to_tensor([1, 4, 4, 3], dtype=tf.int32)
    absolute_boxes = ops.normalized_to_image_coordinates(normalized_boxes,
                                                         image_shape,
                                                         parallel_iterations=2)

    expected_boxes = np.array([[[0, 0, 4, 4]],
                               [[2, 2, 4, 4]]])
    with self.test_session() as sess:
      absolute_boxes = sess.run(absolute_boxes,
                                feed_dict={normalized_boxes:
                                           normalized_boxes_np})

    self.assertAllEqual(absolute_boxes, expected_boxes) 

def test_indices_to_dense_vector_size_at_inference(self):
    size = 5000
    num_indices = 250
    all_indices = np.arange(size)
    rand_indices = np.random.permutation(all_indices)[0:num_indices]

    expected_output = np.zeros(size, dtype=np.float32)
    expected_output[rand_indices] = 1.

    tf_all_indices = tf.placeholder(tf.int32)
    tf_rand_indices = tf.constant(rand_indices)
    indicator = ops.indices_to_dense_vector(tf_rand_indices,
                                            tf.shape(tf_all_indices)[0])
    feed_dict = {tf_all_indices: all_indices}

    with self.test_session() as sess:
      output = sess.run(indicator, feed_dict=feed_dict)
      self.assertAllEqual(output, expected_output)
      self.assertEqual(output.dtype, expected_output.dtype) 

def read_analogies(self):
    """Reads through the analogy question file.

    Returns:
      questions: a [n, 4] numpy array containing the analogy question's
                 word ids.
      questions_skipped: questions skipped due to unknown words.
    """
    questions = []
    questions_skipped = 0
    with open(self._options.eval_data, "rb") as analogy_f:
      for line in analogy_f:
        if line.startswith(b":"):  # Skip comments.
          continue
        words = line.strip().lower().split(b" ")
        ids = [self._word2id.get(w.strip()) for w in words]
        if None in ids or len(ids) != 4:
          questions_skipped += 1
        else:
          questions.append(np.array(ids))
    print("Eval analogy file: ", self._options.eval_data)
    print("Questions: ", len(questions))
    print("Skipped: ", questions_skipped)
    self._analogy_questions = np.array(questions, dtype=np.int32) 

def apply_with_random_selector(x, func, num_cases):
  """Computes func(x, sel), with sel sampled from [0...num_cases-1].

  Args:
    x: input Tensor.
    func: Python function to apply.
    num_cases: Python int32, number of cases to sample sel from.

  Returns:
    The result of func(x, sel), where func receives the value of the
    selector as a python integer, but sel is sampled dynamically.
  """
  sel = tf.random_uniform([], maxval=num_cases, dtype=tf.int32)
  # Pass the real x only to one of the func calls.
  return control_flow_ops.merge([
      func(control_flow_ops.switch(x, tf.equal(sel, case))[1], case)
      for case in range(num_cases)
  ])[0] 

def test_position_sensitive_with_single_bin(self):
    num_spatial_bins = [1, 1]
    image_shape = [2, 3, 3, 4]
    crop_size = [2, 2]

    image = tf.random_uniform(image_shape)
    boxes = tf.random_uniform((6, 4))
    box_ind = tf.constant([0, 0, 0, 1, 1, 1], dtype=tf.int32)

    # When a single bin is used, position-sensitive crop and pool should be
    # the same as non-position sensitive crop and pool.
    crop = tf.image.crop_and_resize(image, boxes, box_ind, crop_size)
    crop_and_pool = tf.reduce_mean(crop, [1, 2], keep_dims=True)

    ps_crop_and_pool = ops.position_sensitive_crop_regions(
        image, boxes, box_ind, crop_size, num_spatial_bins, global_pool=True)

    with self.test_session() as sess:
      expected_output, output = sess.run((crop_and_pool, ps_crop_and_pool))
      self.assertAllClose(output, expected_output) 

def test_filter_with_missing_fields(self):
    input_boxes = tf.placeholder(tf.float32, shape=(None, 4))
    input_classes = tf.placeholder(tf.int32, shape=(None,))
    input_tensors = {
        fields.InputDataFields.groundtruth_boxes: input_boxes,
        fields.InputDataFields.groundtruth_classes: input_classes
    }
    valid_indices = tf.placeholder(tf.int32, shape=(None,))

    feed_dict = {
        input_boxes:
        np.array([[0.2, 0.4, 0.1, 0.8], [0.2, 0.4, 1.0, 0.8]], dtype=np.float),
        input_classes:
        np.array([1, 2], dtype=np.int32),
        valid_indices:
        np.array([0], dtype=np.int32)
    }
    expected_tensors = {
        fields.InputDataFields.groundtruth_boxes:
        [[0.2, 0.4, 0.1, 0.8]],
        fields.InputDataFields.groundtruth_classes:
        [1]
    }

    output_tensors = ops.retain_groundtruth(input_tensors, valid_indices)
    with self.test_session() as sess:
      output_tensors = sess.run(output_tensors, feed_dict=feed_dict)
      for key in [fields.InputDataFields.groundtruth_boxes]:
        self.assertAllClose(expected_tensors[key], output_tensors[key])
      for key in [fields.InputDataFields.groundtruth_classes]:
        self.assertAllEqual(expected_tensors[key], output_tensors[key]) 

def GetStep(self):
    def OnesInitializer(shape, dtype=tf.float32, partition_info=None):
      return tf.ones(shape, dtype)
    return self._AddVariable([], tf.int32, 'step', OnesInitializer) 

def test_return_only_valid_boxes_when_input_contains_invalid_boxes(self):
    num_classes = 4
    num_valid_boxes = 3
    num_boxes = 10
    code_size = 4

    dense_location_placeholder = tf.placeholder(tf.float32, shape=(num_boxes,
                                                                   code_size))
    dense_num_boxes_placeholder = tf.placeholder(tf.int32, shape=(num_classes))
    box_locations, box_classes = ops.dense_to_sparse_boxes(
        dense_location_placeholder, dense_num_boxes_placeholder, num_classes)
    feed_dict = {dense_location_placeholder: np.random.uniform(
        size=[num_boxes, code_size]),
                 dense_num_boxes_placeholder: np.array([1, 0, 0, 2],
                                                       dtype=np.int32)}

    expected_box_locations = (feed_dict[dense_location_placeholder]
                              [:num_valid_boxes])
    expected_box_classses = np.array([0, 3, 3])
    with self.test_session() as sess:
      box_locations, box_classes = sess.run([box_locations, box_classes],
                                            feed_dict=feed_dict)

    self.assertAllClose(box_locations, expected_box_locations, rtol=1e-6,
                        atol=1e-6)
    self.assertAllEqual(box_classes, expected_box_classses) 

def test_construct_anchor_grid_non_square(self):
    base_anchor_size = tf.constant([1, 1], dtype=tf.float32)
    box_specs_list = [[(1.0, 1.0)]]

    exp_anchor_corners = [[0., -0.25, 1., 0.75], [0., 0.25, 1., 1.25]]

    anchor_generator = ag.MultipleGridAnchorGenerator(box_specs_list,
                                                      base_anchor_size)
    anchors = anchor_generator.generate(feature_map_shape_list=[(tf.constant(
        1, dtype=tf.int32), tf.constant(2, dtype=tf.int32))])
    anchor_corners = anchors.get()

    with self.test_session():
      anchor_corners_out = anchor_corners.eval()
      self.assertAllClose(anchor_corners_out, exp_anchor_corners) 

def sample_action(rng, action_probs, optimal_action, sample_gt_prob,
                  type='sample', combine_type='one_or_other'):
  optimal_action_ = optimal_action/np.sum(optimal_action+0., 1, keepdims=True)
  action_probs_ = action_probs/np.sum(action_probs+0.001, 1, keepdims=True)
  batch_size = action_probs_.shape[0]

  action = np.zeros((batch_size), dtype=np.int32)
  action_sample_wt = np.zeros((batch_size), dtype=np.float32)
  if combine_type == 'add':
    sample_gt_prob_ = np.minimum(np.maximum(sample_gt_prob, 0.), 1.)

  for i in range(batch_size):
    if combine_type == 'one_or_other':
      sample_gt = rng.rand() < sample_gt_prob
      if sample_gt: distr_ = optimal_action_[i,:]*1.
      else: distr_ = action_probs_[i,:]*1.
    elif combine_type == 'add':
      distr_ = optimal_action_[i,:]*sample_gt_prob_ + \
          (1.-sample_gt_prob_)*action_probs_[i,:]
      distr_ = distr_ / np.sum(distr_)

    if type == 'sample':
      action[i] = np.argmax(rng.multinomial(1, distr_, size=1))
    elif type == 'argmax':
      action[i] = np.argmax(distr_)
    action_sample_wt[i] = action_probs_[i, action[i]] / distr_[action[i]]
  return action, action_sample_wt 

def make_update_op(self, upd_idxs, upd_keys, upd_vals,
                     batch_size, use_recent_idx, intended_output):
    """Function that creates all the update ops."""
    base_update_op = super(LSHMemory, self).make_update_op(
        upd_idxs, upd_keys, upd_vals,
        batch_size, use_recent_idx, intended_output)

    # compute hash slots to be updated
    hash_slot_idxs = self.get_hash_slots(upd_keys)

    # make updates
    update_ops = []
    with tf.control_dependencies([base_update_op]):
      for i, slot_idxs in enumerate(hash_slot_idxs):
        # for each slot, choose which entry to replace
        entry_idx = tf.random_uniform([batch_size],
                                      maxval=self.num_per_hash_slot,
                                      dtype=tf.int32)
        entry_mul = 1 - tf.one_hot(entry_idx, self.num_per_hash_slot,
                                   dtype=tf.int32)
        entry_add = (tf.expand_dims(upd_idxs, 1) *
                     tf.one_hot(entry_idx, self.num_per_hash_slot,
                                dtype=tf.int32))

        mul_op = tf.scatter_mul(self.hash_slots[i], slot_idxs, entry_mul)
        with tf.control_dependencies([mul_op]):
          add_op = tf.scatter_add(self.hash_slots[i], slot_idxs, entry_add)
          update_ops.append(add_op)

    return tf.group(*update_ops) 

def __init__(self, key_dim, memory_size, vocab_size,
               choose_k=256, alpha=0.1, correct_in_top=1, age_noise=8.0,
               var_cache_device='', nn_device=''):
    self.key_dim = key_dim
    self.memory_size = memory_size
    self.vocab_size = vocab_size
    self.choose_k = min(choose_k, memory_size)
    self.alpha = alpha
    self.correct_in_top = correct_in_top
    self.age_noise = age_noise
    self.var_cache_device = var_cache_device  # Variables are cached here.
    self.nn_device = nn_device  # Device to perform nearest neighbour matmul.

    caching_device = var_cache_device if var_cache_device else None
    self.update_memory = tf.constant(True)  # Can be fed "false" if needed.
    self.mem_keys = tf.get_variable(
        'memkeys', [self.memory_size, self.key_dim], trainable=False,
        initializer=tf.random_uniform_initializer(-0.0, 0.0),
        caching_device=caching_device)
    self.mem_vals = tf.get_variable(
        'memvals', [self.memory_size], dtype=tf.int32, trainable=False,
        initializer=tf.constant_initializer(0, tf.int32),
        caching_device=caching_device)
    self.mem_age = tf.get_variable(
        'memage', [self.memory_size], dtype=tf.float32, trainable=False,
        initializer=tf.constant_initializer(0.0), caching_device=caching_device)
    self.recent_idx = tf.get_variable(
        'recent_idx', [self.vocab_size], dtype=tf.int32, trainable=False,
        initializer=tf.constant_initializer(0, tf.int32))

    # variable for projecting query vector into memory key
    self.query_proj = tf.get_variable(
        'memory_query_proj', [self.key_dim, self.key_dim], dtype=tf.float32,
        initializer=tf.truncated_normal_initializer(0, 0.01),
        caching_device=caching_device) 

def train(loss, init_fn, hparams):
  """Wraps slim.learning.train to run a training loop.

  Args:
    loss: a loss tensor
    init_fn: A callable to be executed after all other initialization is done.
    hparams: a model hyper parameters
  """
  optimizer = create_optimizer(hparams)

  if FLAGS.sync_replicas:
    replica_id = tf.constant(FLAGS.task, tf.int32, shape=())
    optimizer = tf.LegacySyncReplicasOptimizer(
        opt=optimizer,
        replicas_to_aggregate=FLAGS.replicas_to_aggregate,
        replica_id=replica_id,
        total_num_replicas=FLAGS.total_num_replicas)
    sync_optimizer = optimizer
    startup_delay_steps = 0
  else:
    startup_delay_steps = 0
    sync_optimizer = None

  train_op = slim.learning.create_train_op(
      loss,
      optimizer,
      summarize_gradients=True,
      clip_gradient_norm=FLAGS.clip_gradient_norm)

  slim.learning.train(
      train_op=train_op,
      logdir=FLAGS.train_log_dir,
      graph=loss.graph,
      master=FLAGS.master,
      is_chief=(FLAGS.task == 0),
      number_of_steps=FLAGS.max_number_of_steps,
      save_summaries_secs=FLAGS.save_summaries_secs,
      save_interval_secs=FLAGS.save_interval_secs,
      startup_delay_steps=startup_delay_steps,
      sync_optimizer=sync_optimizer,
      init_fn=init_fn) 

def sequence_accuracy(predictions, targets, rej_char, streaming=False):
  """Computes sequence level accuracy.

  Both input tensors should have the same shape: [batch_size x seq_length].

  Args:
    predictions: predicted character classes.
    targets: ground truth character classes.
    rej_char: the character id used to mark empty element (end of sequence).
    streaming: if True, uses the streaming mean from the slim.metric module.

  Returns:
    a update_ops for execution and value tensor whose value on evaluation
    returns the total sequence accuracy.
  """

  with tf.variable_scope('SequenceAccuracy'):
    predictions.get_shape().assert_is_compatible_with(targets.get_shape())

    targets = tf.to_int32(targets)
    const_rej_char = tf.constant(
        rej_char, shape=targets.get_shape(), dtype=tf.int32)
    include_mask = tf.not_equal(targets, const_rej_char)
    include_predictions = tf.to_int32(
        tf.where(include_mask, predictions,
                 tf.zeros_like(predictions) + rej_char))
    correct_chars = tf.to_float(tf.equal(include_predictions, targets))
    correct_chars_counts = tf.cast(
        tf.reduce_sum(correct_chars, reduction_indices=[1]), dtype=tf.int32)
    target_length = targets.get_shape().dims[1].value
    target_chars_counts = tf.constant(
        target_length, shape=correct_chars_counts.get_shape())
    accuracy_per_example = tf.to_float(
        tf.equal(correct_chars_counts, target_chars_counts))
    if streaming:
      return tf.contrib.metrics.streaming_mean(accuracy_per_example)
    else:
      return tf.reduce_mean(accuracy_per_example)