Function TestPicker

tensorflow/core/lib/random/weighted_picker_test.cc:121–181

}

static void TestPicker(SimplePhilox* rnd, int size) {
  VLOG(0) << "======= Testing size " << size;

  // Check that empty picker returns -1
  {
    WeightedPicker picker(size);
    picker.SetAllWeights(0);
    for (int i = 0; i < 100; i++) EXPECT_EQ(picker.Pick(rnd), -1);
  }

  // Create zero weights array
  std::vector<int32> weights(size);
  for (int elem = 0; elem < size; elem++) {
    weights[elem] = 0;
  }

  // Check that singleton picker always returns the same element
  for (int elem = 0; elem < size; elem++) {
    WeightedPicker picker(size);
    picker.SetAllWeights(0);
    picker.set_weight(elem, elem + 1);
    for (int i = 0; i < 100; i++) EXPECT_EQ(picker.Pick(rnd), elem);
    weights[elem] = 10;
    picker.SetWeightsFromArray(size, &weights[0]);
    for (int i = 0; i < 100; i++) EXPECT_EQ(picker.Pick(rnd), elem);
    weights[elem] = 0;
  }

  // Check that uniform picker generates elements roughly uniformly
  {
    WeightedPicker picker(size);
    CheckUniform(rnd, &picker, 100000);
  }

  // Check uniform picker that was grown piecemeal
  if (size / 3 > 0) {
    WeightedPicker picker(size / 3);
    while (picker.num_elements() != size) {
      picker.Append(1);
    }
    CheckUniform(rnd, &picker, 100000);
  }

  // Check that skewed distribution works
  if (size <= 10) {
    // When picker grows one element at a time
    WeightedPicker picker(size);
    int32 weight = 1;
    for (int elem = 0; elem < size; elem++) {
      picker.set_weight(elem, weight);
      weights[elem] = weight;
      weight *= 2;
    }
    CheckSkewed(rnd, &picker, 1000000);

    // When picker is created from an array
    WeightedPicker array_picker(0);
    array_picker.SetWeightsFromArray(size, &weights[0]);