def generate_labels()

in recommended-item-search/softmax_model.py [0:0]


def generate_labels(features):
  def _select_random(x):
    """Selectes a random elements from each row of x."""
    def to_float(x):
      return tf.cast(x, tf.float32)
    def to_int(x):
      return tf.cast(x, tf.int64)

    batch_size = tf.shape(x)[0]
    rn = tf.range(batch_size)
    nnz = to_float(tf.count_nonzero(x >= 0, axis=1))
    rnd = tf.random_uniform([batch_size])
    ids = tf.stack([to_int(rn), to_int(nnz * rnd)], axis=1)
    return to_int(tf.gather_nd(x, ids))
  return _select_random(features['movie_ids'])