<|start_filename|>aptamers_mlpd/preprocess/scam_featurize_main.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// A FlumeC++ pipeline that transforms measurements into ScaM inputs.

#include <gflags/gflags.h>

#include "xxx/flume/public/flume.h"
#include "xxx/flume/public/recordio.h"
#include "xxx/flume/public/sstableio.h"
#include "xxx/preprocess/scam_featurize.h"

DEFINE_FLAG(string, input_file, "", "A RecordIO of DNA sequences.");
DEFINE_FLAG(string, output_file, "",
            "An SSTable of N-gram ScaM GFVs keyed by DNA sequence.");

namespace research_biology {
namespace aptamers {

// Run the ScaM featurization on the input file and write to the output file.
void Run(const string& input_file, const string& output_file) {
  // Initialize Flume.
  flume::Flume flume;

  SequenceCollection input = SequenceCollection::Read(
      "ReadSequences",
      flume::RecordIOFile::Source(input_file, flume::Strings()));

  // Transform into ScaM GFV features.
  FeatureVectorTable scam_features = FeaturizeSequences(input);

  // Write to the specified output file.
  flume::PSink<FeatureVectorEntry> output =
      flume::SSTableSink<string, research_scam::GenericFeatureVector>(
          output_file, flume::Strings(),
          flume::Protos<research_scam::GenericFeatureVector>());

  scam_features.Write("WriteFeatures", output);
}

}  // namespace aptamers
}  // namespace research_biology

int main(int argc, char** argv) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  if (base::GetFlag(FLAGS_input_file).empty()) {
    LOG(ERROR) << "--input_file is required";
    return 1;
  }
  if (base::GetFlag(FLAGS_output_file).empty()) {
    LOG(ERROR) << "--output_file is required";
    return 1;
  }

  research_biology::aptamers::Run(base::GetFlag(FLAGS_input_file),
                                  base::GetFlag(FLAGS_output_file));

  return 0;
}


<|start_filename|>scann/scann/scann_ops/cc/python/scann_pybind.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <cstdint>

#include "absl/types/optional.h"
#include "pybind11/pybind11.h"
#include "pybind11/stl.h"
#include "scann/scann_ops/cc/scann_npy.h"

PYBIND11_MODULE(scann_pybind, py_module) {
  py_module.doc() = "pybind11 wrapper for ScaNN";
  pybind11::class_<research_scann::ScannNumpy>(py_module, "ScannNumpy")
      .def(pybind11::init<
           std::optional<const research_scann::np_row_major_arr<float>>,
           std::optional<const research_scann::np_row_major_arr<uint32_t>>,
           std::optional<const research_scann::np_row_major_arr<uint8_t>>,
           std::optional<const research_scann::np_row_major_arr<int8_t>>,
           std::optional<const research_scann::np_row_major_arr<float>>,
           std::optional<const research_scann::np_row_major_arr<float>>,
           const std::string&>())
      .def(pybind11::init<const research_scann::np_row_major_arr<float>&,
                          const std::string&, int>())
      .def("search", &research_scann::ScannNumpy::Search)
      .def("search_batched", &research_scann::ScannNumpy::SearchBatched)
      .def("serialize", &research_scann::ScannNumpy::Serialize);
}


<|start_filename|>scann/scann/distance_measures/distance_measure_factory.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/distance_measures/distance_measure_factory.h"

namespace research_scann {

StatusOr<shared_ptr<DistanceMeasure>> GetDistanceMeasure(
    const DistanceMeasureConfig& config) {
  if (config.distance_measure().empty()) {
    return InvalidArgumentError(
        "Empty DistanceMeasureConfig proto! Must specify distance_measure.");
  }
  return GetDistanceMeasure(config.distance_measure());
}

StatusOr<shared_ptr<DistanceMeasure>> GetDistanceMeasure(string_view name) {
  if (name == "DotProductDistance")
    return shared_ptr<DistanceMeasure>(new DotProductDistance());
  if (name == "BinaryDotProductDistance")
    return shared_ptr<DistanceMeasure>(new BinaryDotProductDistance());
  if (name == "AbsDotProductDistance")
    return shared_ptr<DistanceMeasure>(new AbsDotProductDistance());
  if (name == "L2Distance")
    return shared_ptr<DistanceMeasure>(new L2Distance());
  if (name == "SquaredL2Distance")
    return shared_ptr<DistanceMeasure>(new SquaredL2Distance());
  if (name == "NegatedSquaredL2Distance")
    return shared_ptr<DistanceMeasure>(new NegatedSquaredL2Distance());
  if (name == "L1Distance")
    return shared_ptr<DistanceMeasure>(new L1Distance());
  if (name == "CosineDistance")
    return shared_ptr<DistanceMeasure>(new CosineDistance());
  if (name == "BinaryCosineDistance")
    return shared_ptr<DistanceMeasure>(new BinaryCosineDistance());
  if (name == "GeneralJaccardDistance")
    return shared_ptr<DistanceMeasure>(new GeneralJaccardDistance());
  if (name == "BinaryJaccardDistance")
    return shared_ptr<DistanceMeasure>(new BinaryJaccardDistance());
  if (name == "LimitedInnerProductDistance")
    return shared_ptr<DistanceMeasure>(new LimitedInnerProductDistance());
  if (name == "GeneralHammingDistance")
    return shared_ptr<DistanceMeasure>(new GeneralHammingDistance());
  if (name == "BinaryHammingDistance")
    return shared_ptr<DistanceMeasure>(new BinaryHammingDistance());
  if (name == "NonzeroIntersectDistance")
    return shared_ptr<DistanceMeasure>(new NonzeroIntersectDistance());
  return InvalidArgumentError("Invalid distance_measure: '%s'", name);
}

}  // namespace research_scann


<|start_filename|>zebraix/misc/misc_render.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef ZEBRAIX_MISC_MISC_RENDER_H_
#define ZEBRAIX_MISC_MISC_RENDER_H_

// Functions that apply rendering rules to .

#include <cmath>
#include <limits>
#include <vector>

#include "base/dominance.h"
#include "base/zebraix_graph.proto.h"

namespace zebraix {
namespace misc {

static constexpr char kLayoutDefaultPbTxt[] =
    "base_node {"
    "   label_text: \"Missing label\","
    "   prime_rank: -1,"
    "   node_radius: 10.5,"
    "   label_radius: 18.0,"
    "}"
    "layout {"
    "  direction: E,"
    "  sep_points: 72.0,"
    "}";
static constexpr char kImputedSource[] =
    "label_text: \"Extrapolated source\",";  // Ranks edited in.
static constexpr char kImputedSink[] =
    "label_text: \"Extrapolated sink\",";  // Ranks edited in.

inline void NormalizeVector(double* x, double* y) {
  double norm = 1.0 / std::sqrt(*x * *x + *y * *y +
                                std::numeric_limits<double>::epsilon());
  *x = *x * norm;
  *y = *y * norm;
}

inline zebraix_proto::ShowHide LeastShowy(zebraix_proto::ShowHide v1,
                                          zebraix_proto::ShowHide v2) {
  if (v1 == v2) {
    return v1;
  } else if (v1 == zebraix_proto::HIDE || v2 == zebraix_proto::HIDE) {
    return zebraix_proto::HIDE;
  } else if (v1 == zebraix_proto::GHOST || v2 == zebraix_proto::GHOST) {
    return zebraix_proto::GHOST;
  } else if (v1 == zebraix_proto::WAYPOINT || v2 == zebraix_proto::WAYPOINT) {
    return zebraix_proto::WAYPOINT;
  } else {
    return v1;
  }
}

void ApplyNodeLogic(const std::vector<zebraix::base::DominanceNode>& d_nodes,
                    const zebraix_proto::Layout& proto_layout,
                    std::vector<zebraix_proto::Node>* target_nodes);

}  // namespace misc
}  // namespace zebraix

#endif  // ZEBRAIX_MISC_MISC_RENDER_H_


<|start_filename|>cube_unfoldings/unfoldings.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <assert.h>
#include <stdio.h>

#include <array>
#include <fstream>
#include <iostream>
#include <string>
#include <unordered_set>
#include <vector>

#include "absl/container/flat_hash_set.h"
#include "nauty_wrapper.h"
#include "perfect_matchings.h"

namespace cube_unfoldings {
namespace {
// Writes a matching in "canonical form": edges are sorted so that the first
// node is < the second one, and all edges are in ascending order.
void Canonicalize(Matching& pm) {
  for (size_t i = 0; i < pm.size(); i++) {
    int a = pm[i][0];
    int b = pm[i][1];
    if (a > b) std::swap(a, b);
    pm[i][0] = a;
    pm[i][1] = b;
  }
  std::sort(pm.begin(), pm.end());
}

// Find all the matchings that are equivalent to the given one under
// automorphism of the tree. All matchings are added to `all_found_matchings`.
void FindAllEquivalentMatchings(
    const Matching& pm, const std::vector<Generator>& tree_autom_gen,
    absl::flat_hash_set<Matching>* all_found_matchings) {
  absl::flat_hash_set<Matching> matchings;
  all_found_matchings->insert(pm);
  matchings.insert(std::move(pm));
  Matching new_m;
  std::vector<Matching> new_matchings;
  // Finds the closure of the current set of equivalent matchings under the
  // tree automorphism generators `tree_autom_gen`, by iteratively applying
  // the generators until no new matching is produced.
  while (true) {
    size_t n = matchings.size();
    for (size_t i = 0; i < tree_autom_gen.size(); i++) {
      new_matchings.clear();
      for (const Matching& m : matchings) {
        new_m = m;
        for (size_t j = 0; j < new_m.size(); j++) {
          new_m[j][0] = tree_autom_gen[i][new_m[j][0]];
          new_m[j][1] = tree_autom_gen[i][new_m[j][1]];
        }
        Canonicalize(new_m);
        if (!matchings.contains(new_m)) new_matchings.emplace_back(new_m);
      }
      for (Matching& m : new_matchings) {
        all_found_matchings->insert(m);
        matchings.insert(std::move(m));
      }
    }
    if (n == matchings.size()) break;
  }
}
}  // namespace

size_t NumberOfUniqueMatchings(const char* tree) {
  auto tree_data = TreeData::FromString(tree);
  absl::flat_hash_set<Matching> all_found_matchings;
  size_t num = 0;
  auto cb = [&](const Matching& matching) {
    if (all_found_matchings.contains(matching)) return;
    num++;
    FindAllEquivalentMatchings(matching, tree_data.generators,
                               &all_found_matchings);
  };
  AllPerfectMatchings(tree_data.complement_edges, tree_data.n, cb);

  return num;
}
}  // namespace cube_unfoldings


<|start_filename|>fairness_and_bias_in_online_selection/synthetic_data.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "synthetic_data.h"

#include <iostream>

#include "random_handler.h"

namespace fair_secretary {

using std::vector;

vector<SecretaryInstance> SyntheticData::GetSecretaryInput(
    const vector<int>& sizes, const std::vector<double>& prob) {
  num_colors = sizes.size();
  vector<SecretaryInstance> instance;
  double rand = 2.0;
  if (!prob.empty()) {
    rand = static_cast<double>(RandomHandler::eng_()) /
           std::numeric_limits<uint64_t>::max();
  }
  for (int i = 0; i < sizes.size(); i++) {
    for (int j = 0; j < sizes[i]; j++) {
      instance.push_back(SecretaryInstance(
          static_cast<double>(RandomHandler::eng_()) / 10, i));
    }
    if (!prob.empty()) {
      if (prob[i] > rand && rand >= 0) {
        instance[instance.size() - 1].value =
            std::numeric_limits<uint64_t>::max();
      }
      rand -= prob[i];
    }
  }
  return instance;
}

vector<SecretaryInstance> SyntheticData::GetProphetInput(
    const int size,
    const vector<std::reference_wrapper<RandomDistribution>>& dist) {
  num_colors = size;
  vector<SecretaryInstance> instance;
  for (int i = 0; i < size; i++) {
    if (i < size / 2) {
      instance.push_back(SecretaryInstance(dist[0].get().Sample(), i, 0));

    } else {
      instance.push_back(SecretaryInstance(dist[1].get().Sample(), i, 1));
    }
  }
  return instance;
}

}  // namespace fair_secretary


<|start_filename|>scann/scann/utils/top_n_amortized_constant.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/utils/top_n_amortized_constant.h"

#include "absl/flags/flag.h"
#include "scann/utils/zip_sort.h"

ABSL_RETIRED_FLAG(bool, use_branch_optimized_top_n, , );

namespace research_scann {

template <typename Distance>
void TopNeighbors<Distance>::PartitionElements(vector<Neighbor>* elements,
                                               const DistanceComparator& cmp) {
  ZipNthElementBranchOptimized(DistanceComparatorBranchOptimized(),
                               this->limit() - 1, elements->begin(),
                               elements->end());
}

SCANN_INSTANTIATE_TYPED_CLASS(, TopNeighbors);

}  // namespace research_scann


<|start_filename|>fairness_and_bias_in_online_selection/fair-secretary.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "fair-secretary.h"

#include <iostream>
#include <ostream>

namespace fair_secretary {

using std::vector;

FairSecretaryAlgorithm::FairSecretaryAlgorithm(const vector<int>& thre,
                                               int num_colors) {
  thre_ = thre;
  num_colors_ = num_colors;
}

SecretaryInstance FairSecretaryAlgorithm::ComputeSolution(
    const vector<SecretaryInstance>& elements) {
  vector<double> max_color(num_colors_, 0);
  // Compute max values.
  for (int i = 0; i < elements.size(); i++) {
    if (i < thre_[elements[i].color]) {
      max_color[elements[i].color] =
          std::max(max_color[elements[i].color], elements[i].value);
    }
  }
  for (int i = 0; i < elements.size(); i++) {
    if (i >= thre_[elements[i].color]) {
      if (elements[i].value >= max_color[elements[i].color]) {
        return elements[i];
      }
    }
  }
  return SecretaryInstance(-1, -1);
}

}  // namespace fair_secretary


<|start_filename|>fairness_and_bias_in_online_selection/distributions.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "distributions.h"

namespace fair_secretary {

using std::vector;

void UniformDistribution::Init(int n, double range) {
  range_ = range;
  p_th_ = vector<double>(n, 0);
  vector<double> V(n, 0.0);
  V[n - 1] = range / 2;
  for (int i = n - 2; i >= 0; i--) {
    p_th_[i] = V[i + 1];
    V[i] = (range + p_th_[i]) / 2;
  }
}

double UniformDistribution::PThreshold(int index) { return p_th_[index]; }

double UniformDistribution::Middle(int n) {
  return range_ * pow(1.0 / 2, 1.0 / n);
}

double UniformDistribution::Reverse(double x) { return x * range_; }

double UniformDistribution::Sample() {
  return (static_cast<double>(RandomHandler::eng_()) /
          std::numeric_limits<uint64_t>::max()) *
         range_;
}

void BinomialDistribution::Init(int n, double p) {
  choose_ = vector<vector<double>>(n + 1, vector<double>(n + 1, 0));
  probability_ = vector<double>(n + 1, 1.0);
  r_probability_ = vector<double>(n + 1, 1.0);
  for (int i = 0; i <= n; i++) {
    choose_[i][0] = 1;
  }
  for (int i = 1; i <= n; i++) {
    for (int j = 1; j <= n; j++) {
      choose_[i][j] = choose_[i - 1][j - 1] + choose_[i - 1][j];
    }
  }
  for (int i = 1; i <= n; i++) {
    probability_[i] = probability_[i - 1] * p;
    r_probability_[i] = r_probability_[i - 1] * (1 - p);
  }
  distribution_ =
      std::binomial_distribution<int>(probability_.size() - 1, probability_[1]);
}

double BinomialDistribution::Expected(double lower_bound) {
  double ans = 0;
  double range = 0;
  int n = probability_.size() - 1;
  for (int i = ceil(lower_bound * n); i <= n; i++) {
    ans += probability_[i] * r_probability_[n - i] * choose_[n][i] * i / n;
    range += probability_[i] * r_probability_[n - i] * choose_[n][i];
  }
  return ans / range;
}

double BinomialDistribution::Reverse(double x) {
  int n = probability_.size() - 1;
  for (int i = 0; i <= n; i++) {
    x -= choose_[n][i] * probability_[i] * r_probability_[n - i];
    if (x <= 0) {
      return static_cast<double>(i) / n;
    }
  }
  return 1.0;
}

double BinomialDistribution::Sample() {
  double rand = (static_cast<double>(RandomHandler::eng_()) /
                 std::numeric_limits<uint64_t>::max());
  return (static_cast<double>(distribution_(RandomHandler::generator_)) +
          rand) /
         (probability_.size() - 1);
}

double BinomialDistribution::Middle(int n) {
  for (int i = max_dist_.size() - 1; i >= 0; i--) {
    if (max_dist_[i] >= 0.5) {
      return static_cast<double>(i) / (max_dist_.size() - 1);
    }
  }
  return 0.0;
}

void BinomialDistribution::ComputeMaxDist(double num_dists) {
  max_dist_ = vector<double>(probability_.size(), 0.0);
  int n = probability_.size() - 1;
  double x = 0;
  for (int i = n; i >= 0; i--) {
    x += choose_[n][i] * probability_[i] * r_probability_[n - i];
    max_dist_[i] = (1 - pow(1 - x, num_dists));
  }
  // Computing PThreshold.
  p_th_ = vector<double>(num_dists, 0);
  vector<double> V(num_dists, 0.0);
  V[num_dists - 1] = Expected(0);
  for (int i = num_dists - 2; i >= 0; i--) {
    p_th_[i] = V[i + 1];
    V[i] = Expected(p_th_[i]);
  }
}

double BinomialDistribution::PThreshold(int index) { return p_th_[index]; }

}  // namespace fair_secretary


<|start_filename|>parallel_clustering/clustering/util/dynamic_weight_threshold.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef PARALLEL_CLUSTERING_CLUSTERING_UTIL_DYNAMIC_WEIGHT_THRESHOLD_H_
#define PARALLEL_CLUSTERING_CLUSTERING_UTIL_DYNAMIC_WEIGHT_THRESHOLD_H_

#include "clustering/util/dynamic_weight_threshold.pb.h"
#include "absl/status/statusor.h"

namespace research_graph {

// Computes weight threshold for the given iteration of affinity clustering,
// based on the provided DynamicWeightThresholdConfig.
absl::StatusOr<double> DynamicWeightThreshold(
    const DynamicWeightThresholdConfig& config, int num_iteration,
    int iteration);

}  // namespace research_graph

#endif  // PARALLEL_CLUSTERING_CLUSTERING_UTIL_DYNAMIC_WEIGHT_THRESHOLD_H_


<|start_filename|>fairness_and_bias_in_online_selection/random_handler.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Simple utility function for handling random varialbes.

#ifndef EXPERIMENTAL_USERS_ASHKANNOROUZI_FAST_K_MEANS_2019_RANDOM_HANDLER_H_
#define EXPERIMENTAL_USERS_ASHKANNOROUZI_FAST_K_MEANS_2019_RANDOM_HANDLER_H_

#include <cstdint>
#include <random>

namespace fair_secretary {

class RandomHandler {
 public:
  // Get a random seed from the OS entropy device, or whatever
  // Use the 64-bit Mersenne Twister 19937
  // generator and seed it with entropy.
  // In case of using changing random values, use std::mt19937_64 eng_(rd_);
  static std::random_device rd_;

  // Define the distribution, by default it goes from 0 to MAX(unsigned long
  // long) or the type supported by your architecture.
  static std::mt19937_64 eng_;

  static std::default_random_engine generator_;
};

}  // namespace fair_secretary

#endif  // EXPERIMENTAL_USERS_ASHKANNOROUZI_FAST_K_MEANS_2019_RANDOM_HANDLER_H_


<|start_filename|>scann/scann/partitioning/kmeans_tree_partitioner_helper.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <cstdint>

#include "scann/base/single_machine_base.h"
#include "scann/brute_force/brute_force.h"
#include "scann/hashes/asymmetric_hashing2/indexing.h"
#include "scann/hashes/asymmetric_hashing2/querying.h"
#include "scann/hashes/asymmetric_hashing2/searcher.h"
#include "scann/hashes/asymmetric_hashing2/training.h"
#include "scann/hashes/asymmetric_hashing2/training_options.h"
#include "scann/partitioning/kmeans_tree_partitioner.h"
#include "scann/proto/distance_measure.pb.h"
#include "scann/proto/hash.pb.h"
#include "scann/utils/datapoint_utils.h"
#include "scann/utils/parallel_for.h"
#include "scann/utils/types.h"
#include "tensorflow/core/lib/core/errors.h"

namespace research_scann {

namespace internal {

StatusOr<unique_ptr<SingleMachineSearcherBase<float>>>
CreateRecommendedAsymmetricSearcher(
    shared_ptr<DenseDataset<float>> dataset,
    shared_ptr<const DistanceMeasure> quantization_distance,
    int32_t num_neighbors, float epsilon = numeric_limits<float>::infinity(),
    bool with_exact_reordering = true, shared_ptr<ThreadPool> pool = nullptr,
    int num_clusters_per_block = 16, int num_dimension_per_block = 2) {
  DCHECK(dataset);

  if (dataset->size() < num_clusters_per_block ||
      dataset->dimensionality() < num_dimension_per_block) {
    unique_ptr<SingleMachineSearcherBase<float>> bf_searcher(
        new BruteForceSearcher<float>(quantization_distance, dataset,
                                      num_neighbors, epsilon));
    return std::move(bf_searcher);
  }

  AsymmetricHasherConfig hasher_config;
  hasher_config.set_num_clusters_per_block(num_clusters_per_block);

  DimensionIndex dim = dataset->dimensionality();
  hasher_config.mutable_projection()->set_input_dim(dim);
  hasher_config.mutable_projection()->set_projection_type(
      ProjectionConfig::CHUNK);
  hasher_config.mutable_projection()->set_num_blocks(
      (dim + num_dimension_per_block - 1) / num_dimension_per_block);
  hasher_config.mutable_projection()->set_num_dims_per_block(
      num_dimension_per_block);
  if (num_clusters_per_block == 16) {
    hasher_config.set_lookup_type(AsymmetricHasherConfig::INT8_LUT16);
  } else {
    hasher_config.set_lookup_type(AsymmetricHasherConfig::FLOAT);
  }
  hasher_config.mutable_quantization_distance()->set_distance_measure(
      std::string(quantization_distance->name()));
  asymmetric_hashing2::TrainingOptions<float> training_opts(
      hasher_config, quantization_distance, *dataset);
  TF_ASSIGN_OR_RETURN(shared_ptr<asymmetric_hashing2::Model<float>> model,
                      asymmetric_hashing2::TrainSingleMachine<float>(
                          *dataset, training_opts, pool));
  auto indexer = make_unique<asymmetric_hashing2::Indexer<float>>(
      training_opts.projector(), quantization_distance, model);

  auto hashed_dataset = std::make_shared<DenseDataset<uint8_t>>();
  hashed_dataset->Reserve(dataset->size());
  for (int i = 0; i < dataset->size(); ++i) {
    Datapoint<uint8_t> dp;
    SCANN_RETURN_IF_ERROR(indexer->Hash((*dataset)[i], &dp));
    hashed_dataset->AppendOrDie(dp.ToPtr());
  }

  auto queryer = make_unique<asymmetric_hashing2::AsymmetricQueryer<float>>(
      training_opts.projector(), quantization_distance, model);
  asymmetric_hashing2::SearcherOptions<float> searcher_opts(std::move(queryer));
  searcher_opts.set_asymmetric_lookup_type(hasher_config.lookup_type());
  unique_ptr<SingleMachineSearcherBase<float>> ah_searcher(
      new asymmetric_hashing2::Searcher<float>(
          std::move(dataset), std::move(hashed_dataset), searcher_opts,
          num_neighbors, numeric_limits<float>::infinity()));

  if (with_exact_reordering) {
    auto helper = std::make_shared<ExactReorderingHelper<float>>(
        quantization_distance, ah_searcher->shared_dataset());
    ah_searcher->EnableReordering(
        std::make_shared<ExactReorderingHelper<float>>(
            quantization_distance, ah_searcher->shared_dataset()),
        num_neighbors, epsilon);
  }

  return std::move(ah_searcher);
}

}  // namespace internal
}  // namespace research_scann


<|start_filename|>fairness_and_bias_in_online_selection/synthetic_data.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef FAIRNESS_AND_BIAS_SELECTION_SYNTHETIC_DATA_H_
#define FAIRNESS_AND_BIAS_SELECTION_SYNTHETIC_DATA_H_

#include "distributions.h"
#include "utils.h"

namespace fair_secretary {

// Creates synthetic dataset based on the sizes and probabilities given.
class SyntheticData {
 public:
  std::vector<SecretaryInstance> GetSecretaryInput(
      const std::vector<int>& sizes, const std::vector<double>& prob = {});

  std::vector<SecretaryInstance> GetProphetInput(
      const int size,
      const std::vector<std::reference_wrapper<RandomDistribution>>& dist);

  // Indicates the number of colors.
  int num_colors;
};

}  // namespace fair_secretary

#endif  // FAIRNESS_AND_BIAS_SELECTION_SYNTHETIC_DATA_H_


<|start_filename|>scann/scann/distance_measures/one_to_one/dot_product.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_DISTANCE_MEASURES_ONE_TO_ONE_DOT_PRODUCT_H_
#define SCANN_DISTANCE_MEASURES_ONE_TO_ONE_DOT_PRODUCT_H_

#include <cstdint>

#include "scann/distance_measures/distance_measure_base.h"
#include "scann/distance_measures/one_to_one/binary_distance_measure_base.h"
#include "scann/distance_measures/one_to_one/common.h"
#include "scann/distance_measures/one_to_one/dot_product_avx1.h"
#include "scann/distance_measures/one_to_one/dot_product_avx2.h"
#include "scann/distance_measures/one_to_one/dot_product_sse4.h"
#include "scann/utils/intrinsics/flags.h"
#include "scann/utils/reduction.h"
#include "scann/utils/types.h"

namespace research_scann {

template <typename T, typename U>
double SparseDotProduct(const DatapointPtr<T>& a, const DatapointPtr<U>& b);

template <typename T, typename U>
double DenseDotProduct(const DatapointPtr<T>& a, const DatapointPtr<U>& b);

template <typename T, typename U>
double HybridDotProduct(const DatapointPtr<T>& a, const DatapointPtr<U>& b);

template <typename T, typename U>
double DotProduct(const DatapointPtr<T>& a, const DatapointPtr<U>& b) {
  return SparseDenseDispatch(a, b, &SparseDotProduct<T>, &DenseDotProduct<T>,
                             &HybridDotProduct<T>);
}

inline DimensionIndex BinaryDotProduct(const DatapointPtr<uint8_t>& a,
                                       const DatapointPtr<uint8_t>& b);
inline DimensionIndex DenseBinaryDotProduct(const DatapointPtr<uint8_t>& a,
                                            const DatapointPtr<uint8_t>& b);
inline DimensionIndex SparseBinaryDotProduct(const DatapointPtr<uint8_t>& a,
                                             const DatapointPtr<uint8_t>& b);
inline DimensionIndex HybridBinaryDotProduct(const DatapointPtr<uint8_t>& a,
                                             const DatapointPtr<uint8_t>& b);

class DotProductDistance final : public DistanceMeasure {
 public:
  SCANN_DECLARE_DISTANCE_MEASURE_VIRTUAL_METHODS(DOT_PRODUCT);

 private:
  template <typename T>
  SCANN_INLINE double GetDistanceDenseImpl(const DatapointPtr<T>& a,
                                           const DatapointPtr<T>& b) const {
    return -DenseDotProduct(a, b);
  }

  template <typename T>
  SCANN_INLINE double GetDistanceSparseImpl(const DatapointPtr<T>& a,
                                            const DatapointPtr<T>& b) const {
    return -SparseDotProduct(a, b);
  }

  template <typename T>
  SCANN_INLINE double GetDistanceHybridImpl(const DatapointPtr<T>& a,
                                            const DatapointPtr<T>& b) const {
    return -HybridDotProduct(a, b);
  }
};

class AbsDotProductDistance final : public DistanceMeasure {
 public:
  SCANN_DECLARE_DISTANCE_MEASURE_VIRTUAL_METHODS(ABS_DOT_PRODUCT);

 private:
  template <typename T>
  SCANN_INLINE double GetDistanceDenseImpl(const DatapointPtr<T>& a,
                                           const DatapointPtr<T>& b) const {
    return -std::abs(DenseDotProduct(a, b));
  }

  template <typename T>
  SCANN_INLINE double GetDistanceSparseImpl(const DatapointPtr<T>& a,
                                            const DatapointPtr<T>& b) const {
    return -std::abs(SparseDotProduct(a, b));
  }

  template <typename T>
  SCANN_INLINE double GetDistanceHybridImpl(const DatapointPtr<T>& a,
                                            const DatapointPtr<T>& b) const {
    return -std::abs(HybridDotProduct(a, b));
  }
};

class BinaryDotProductDistance final : public BinaryDistanceMeasureBase {
 public:
  string_view name() const final;

  using BinaryDistanceMeasureBase::GetDistanceDense;
  using BinaryDistanceMeasureBase::GetDistanceHybrid;
  using BinaryDistanceMeasureBase::GetDistanceSparse;

  double GetDistanceDense(const DatapointPtr<uint8_t>& a,
                          const DatapointPtr<uint8_t>& b) const final;
  double GetDistanceSparse(const DatapointPtr<uint8_t>& a,
                           const DatapointPtr<uint8_t>& b) const final;
  double GetDistanceHybrid(const DatapointPtr<uint8_t>& a,
                           const DatapointPtr<uint8_t>& b) const final {
    return GetDistanceSparse(a, b);
  }
};

struct DotProductReduce {
  template <typename Accumulator, typename T, typename U>
  void operator()(Accumulator* acc, const T a, const U b) {
    *acc += static_cast<Accumulator>(a) * static_cast<Accumulator>(b);
  }
};

template <typename T, typename U>
double SparseDotProduct(const DatapointPtr<T>& a, const DatapointPtr<U>& b) {
  return SparsePairAccumulate(a, b, DotProductReduce(), DoNothingReduce());
}

template <typename T, typename U>
double DenseDotProductFallback(const DatapointPtr<T>& a,
                               const DatapointPtr<U>& b) {
  return DensePairAccumulate(a.values(), b.values(), a.nonzero_entries(),
                             DotProductReduce());
}

template <typename T, typename U>
double DenseDotProduct(const DatapointPtr<T>& a, const DatapointPtr<U>& b) {
  return DenseDotProductFallback(a, b);
}

template <typename T, typename U, typename V>
double DenseDotProductFallback(const DatapointPtr<T>& a,
                               const DatapointPtr<U>& b,
                               const DatapointPtr<V>& c) {
  using AT = AccumulatorTypeFor<T, U, V>;
  AT acc = 0;
  DCHECK_EQ(a.nonzero_entries(), b.nonzero_entries());
  DCHECK_EQ(a.nonzero_entries(), c.nonzero_entries());
  const T* aptr = a.values();
  const U* bptr = b.values();
  const V* cptr = c.values();

  for (size_t i : Seq(a.nonzero_entries())) {
    acc += static_cast<AT>(aptr[i]) * static_cast<AT>(bptr[i]) *
           static_cast<AT>(cptr[i]);
  }
  return static_cast<double>(acc);
}

template <typename T, typename U, typename V>
double DenseDotProduct(const DatapointPtr<T>& a, const DatapointPtr<U>& b,
                       const DatapointPtr<V>& c) {
  return DenseDotProductFallback(a, b, c);
}

#ifdef __x86_64__

template <>
inline double DenseDotProduct<uint8_t, uint8_t>(
    const DatapointPtr<uint8_t>& a, const DatapointPtr<uint8_t>& b) {
  if (RuntimeSupportsSse4()) {
    return dp_internal::DenseDotProductSse4(a, b);
  } else {
    return DenseDotProductFallback(a, b);
  }
}

template <>
inline double DenseDotProduct<int8_t, int8_t>(const DatapointPtr<int8_t>& a,
                                              const DatapointPtr<int8_t>& b) {
  if (RuntimeSupportsSse4()) {
    return dp_internal::DenseDotProductSse4(a, b);
  } else {
    return DenseDotProductFallback(a, b);
  }
}

template <>
inline double DenseDotProduct<float, float>(const DatapointPtr<float>& a,
                                            const DatapointPtr<float>& b) {
  if (RuntimeSupportsSse4()) {
    return dp_internal::DenseDotProductSse4(a, b);
  } else {
    return DenseDotProductFallback(a, b);
  }
}

template <>
inline double DenseDotProduct<double, double>(const DatapointPtr<double>& a,
                                              const DatapointPtr<double>& b) {
  if (RuntimeSupportsSse4()) {
    return dp_internal::DenseDotProductSse4(a, b);
  } else {
    return DenseDotProductFallback(a, b);
  }
}

template <>
inline double DenseDotProduct<int8_t, float>(const DatapointPtr<int8_t>& a,
                                             const DatapointPtr<float>& b) {
  if (RuntimeSupportsAvx2()) {
    return dp_internal::DenseDotProductAvx2(a, b);
  } else if (RuntimeSupportsAvx1()) {
    return dp_internal::DenseDotProductAvx1(a, b);
  } else if (RuntimeSupportsSse4()) {
    return dp_internal::DenseDotProductSse4(a, b);
  } else {
    return DenseDotProductFallback(a, b);
  }
}

template <>
inline double DenseDotProduct<float, int8_t>(const DatapointPtr<float>& a,
                                             const DatapointPtr<int8_t>& b) {
  return DenseDotProduct(b, a);
}

template <>
inline double DenseDotProduct<int8_t, float, float>(
    const DatapointPtr<int8_t>& a, const DatapointPtr<float>& b,
    const DatapointPtr<float>& c) {
  if (RuntimeSupportsAvx2()) {
    return dp_internal::DenseDotProductAvx2(a, b, c);
  } else if (RuntimeSupportsAvx1()) {
    return dp_internal::DenseDotProductAvx1(a, b, c);
  } else if (RuntimeSupportsSse4()) {
    return dp_internal::DenseDotProductSse4(a, b, c);
  } else {
    return DenseDotProductFallback(a, b, c);
  }
}

template <>
inline double DenseDotProduct<float, int8_t, float>(
    const DatapointPtr<float>& a, const DatapointPtr<int8_t>& b,
    const DatapointPtr<float>& c) {
  return DenseDotProduct(b, a, c);
}

template <>
inline double DenseDotProduct<float, float, int8_t>(
    const DatapointPtr<float>& a, const DatapointPtr<float>& b,
    const DatapointPtr<int8_t>& c) {
  return DenseDotProduct(c, a, b);
}

template <>
inline double DenseDotProduct<int8_t, int8_t, float>(
    const DatapointPtr<int8_t>& a, const DatapointPtr<int8_t>& b,
    const DatapointPtr<float>& c) {
  if (RuntimeSupportsAvx2()) {
    return dp_internal::DenseDotProductAvx2(a, b, c);
  } else if (RuntimeSupportsAvx1()) {
    return dp_internal::DenseDotProductAvx1(a, b, c);
  } else if (RuntimeSupportsSse4()) {
    return dp_internal::DenseDotProductSse4(a, b, c);
  } else {
    return DenseDotProductFallback(a, b, c);
  }
}

template <>
inline double DenseDotProduct<int8_t, float, int8_t>(
    const DatapointPtr<int8_t>& a, const DatapointPtr<float>& b,
    const DatapointPtr<int8_t>& c) {
  return DenseDotProduct(a, c, b);
}

template <>
inline double DenseDotProduct<float, int8_t, int8_t>(
    const DatapointPtr<float>& a, const DatapointPtr<int8_t>& b,
    const DatapointPtr<int8_t>& c) {
  return DenseDotProduct(c, b, a);
}

#endif

template <typename T, typename U>
double HybridDotProduct(const DatapointPtr<T>& a, const DatapointPtr<U>& b) {
  return HybridPairAccumulate(a, b, DotProductReduce(), DoNothingReduce());
}

inline DimensionIndex BinaryDotProduct(const DatapointPtr<uint8_t>& a,
                                       const DatapointPtr<uint8_t>& b) {
  return SparseDenseDispatch(a, b, &SparseBinaryDotProduct,
                             &DenseBinaryDotProduct, &HybridBinaryDotProduct);
}

class DenseBinaryDotProductAnd {
 public:
  template <typename T>
  T operator()(T a, T b) {
    return a & b;
  }
};

inline DimensionIndex DenseBinaryDotProduct(const DatapointPtr<uint8_t>& a,
                                            const DatapointPtr<uint8_t>& b) {
  return DenseBinaryMergeAndPopcnt(a, b, DenseBinaryDotProductAnd());
}

inline DimensionIndex SparseBinaryDotProduct(const DatapointPtr<uint8_t>& a,
                                             const DatapointPtr<uint8_t>& b) {
  if (a.nonzero_entries() == 0 || b.nonzero_entries() == 0) {
    return 0;
  }
  DimensionIndex intersection = 0;
  size_t a_front = 0, b_front = 0;
  size_t a_back = a.nonzero_entries() - 1, b_back = b.nonzero_entries() - 1;
  while (a_front < a_back && b_front < b_back) {
    const size_t to_add_front1 = a.indices()[a_front] <= b.indices()[b_front];
    const size_t to_add_front2 = a.indices()[a_front] >= b.indices()[b_front];
    const size_t to_sub_back2 = a.indices()[a_back] <= b.indices()[b_back];
    const size_t to_sub_back1 = a.indices()[a_back] >= b.indices()[b_back];
    intersection += a.indices()[a_front] == b.indices()[b_front];
    intersection += a.indices()[a_back] == b.indices()[b_back];
    a_front += to_add_front1;
    b_front += to_add_front2;
    a_back -= to_sub_back1;
    b_back -= to_sub_back2;
  }
  if (a_front == a_back) {
    for (; b_front <= b_back; ++b_front) {
      if (ABSL_PREDICT_FALSE(a.indices()[a_front] == b.indices()[b_front])) {
        intersection++;
        break;
      }
    }
  } else if (b_front == b_back) {
    for (; a_front <= a_back; ++a_front) {
      if (ABSL_PREDICT_FALSE(a.indices()[a_front] == b.indices()[b_front])) {
        intersection++;
        break;
      }
    }
  }
  return intersection;
}

inline DimensionIndex HybridBinaryDotProduct(const DatapointPtr<uint8_t>& a,
                                             const DatapointPtr<uint8_t>& b) {
  LOG(FATAL) << "Hybrid binary dot product not yet implemented.";
}

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/utils/sampled_index_list.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_UTILS_SAMPLED_INDEX_LIST_H_
#define SCANN_UTILS_SAMPLED_INDEX_LIST_H_

#include "absl/types/variant.h"
#include "scann/utils/types.h"
#include "tensorflow/core/platform/macros.h"

namespace research_scann {
namespace internal {

template <typename Index>
class SampledIndexList {
 public:
  SampledIndexList();

  SampledIndexList(Index start, Index stop);

  explicit SampledIndexList(vector<Index> sparse);

  bool GetNextIndex(Index* i);

  void RestartIndex();

 private:
  using DenseIndices = pair<Index, Index>;

  using SparseIndices = vector<Index>;

  const DenseIndices* dense() const {
    return absl::get_if<DenseIndices>(&indices_);
  }

  const SparseIndices* sparse() const {
    return absl::get_if<SparseIndices>(&indices_);
  }

  absl::variant<DenseIndices, SparseIndices> indices_;
  Index current_;
};

template <typename Index>
SampledIndexList<Index>::SampledIndexList()
    : indices_(DenseIndices(0, 0)), current_(0) {}

template <typename Index>
SampledIndexList<Index>::SampledIndexList(Index start, Index stop)
    : indices_(DenseIndices(start, stop)), current_(start) {}

template <typename Index>
SampledIndexList<Index>::SampledIndexList(vector<Index> sparse)
    : indices_(std::move(sparse)), current_(0) {}

template <typename Index>
bool SampledIndexList<Index>::GetNextIndex(Index* i) {
  if (dense()) {
    if (current_ == dense()->second) return false;
    *i = current_;
    ++current_;
  } else {
    if (current_ == sparse()->size()) return false;
    *i = (*sparse())[current_];
    ++current_;
  }
  return true;
}

template <typename Index>
void SampledIndexList<Index>::RestartIndex() {
  if (dense()) {
    current_ = dense()->first;
  } else {
    current_ = 0;
  }
}

}  // namespace internal
}  // namespace research_scann

#endif


<|start_filename|>scann/scann/hashes/asymmetric_hashing2/training_model.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_HASHES_ASYMMETRIC_HASHING2_TRAINING_MODEL_H_
#define SCANN_HASHES_ASYMMETRIC_HASHING2_TRAINING_MODEL_H_

#include <cstdint>

#include "scann/data_format/dataset.h"
#include "scann/proto/centers.pb.h"
#include "scann/proto/hash.pb.h"
#include "scann/utils/types.h"

namespace research_scann {
namespace asymmetric_hashing2 {

template <typename T>
class Model {
 public:
  using FloatT = FloatingTypeFor<T>;

  static StatusOr<unique_ptr<Model<T>>> FromCenters(
      std::vector<DenseDataset<FloatT>> centers,
      AsymmetricHasherConfig::QuantizationScheme quantization_scheme =
          AsymmetricHasherConfig::PRODUCT);

  static StatusOr<unique_ptr<Model<T>>> FromProto(
      const CentersForAllSubspaces& proto);

  CentersForAllSubspaces ToProto() const;

  ConstSpan<DenseDataset<FloatT>> centers() const { return centers_; }

  uint32_t num_clusters_per_block() const { return num_clusters_per_block_; }

  size_t num_blocks() const { return centers_.size(); }

  AsymmetricHasherConfig::QuantizationScheme quantization_scheme() const {
    return quantization_scheme_;
  }

  bool CentersEqual(const Model& rhs) const;

 private:
  explicit Model(
      std::vector<DenseDataset<FloatT>> centers,
      AsymmetricHasherConfig::QuantizationScheme quantization_scheme);

  std::vector<DenseDataset<FloatT>> centers_ = {};

  uint32_t num_clusters_per_block_ = numeric_limits<uint32_t>::max();

  AsymmetricHasherConfig::QuantizationScheme quantization_scheme_ =
      AsymmetricHasherConfig::PRODUCT;

  TF_DISALLOW_COPY_AND_ASSIGN(Model);
};

SCANN_INSTANTIATE_TYPED_CLASS(extern, Model);

}  // namespace asymmetric_hashing2
}  // namespace research_scann

#endif


<|start_filename|>scann/scann/utils/memory_logging.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_UTILS_MEMORY_LOGGING_H_
#define SCANN_UTILS_MEMORY_LOGGING_H_

#include <hash_set>
#include <type_traits>
#include <unordered_map>
#include <unordered_set>

#include "scann/data_format/datapoint.h"
#include "scann/utils/types.h"

namespace research_scann {

static const double bytes_in_mb = 1024.0 * 1024.0;

constexpr double kBytesInMb = 1024.0 * 1024.0;

template <typename T>
size_t VectorStorage(const vector<T>& vec) {
  return sizeof(vec) + vec.capacity() * sizeof(vec[0]);
}

template <typename T>
size_t HashSetStorage(const absl::node_hash_set<T>& set) {
  return sizeof(set) + set.bucket_count() * sizeof(void*) +
         set.size() * sizeof(pair<T, void*>);
}

template <typename T, typename U, typename Hash>
size_t HashMapStorage(const std::unordered_map<T, U, Hash>& map) {
  return sizeof(map) + map.bucket_count() * sizeof(void*) +
         map.size() * sizeof(std::tuple<T, U, void*>);
}

template <typename K, typename V, typename... Args>
size_t HashMapStorage(const flat_hash_map<K, V, Args...>& map) {
  return sizeof(map) + map.bucket_count() * (sizeof(pair<K, V>) + 1);
}

template <typename T>
size_t HashSetStorage(const flat_hash_set<T>& set) {
  return sizeof(set) + set.bucket_count() * (sizeof(T) + 1);
}

template <typename Vector>
size_t TwoDimensionalVectorStorage(const Vector& vec) {
  size_t result = sizeof(vec) + vec.capacity() * sizeof(vec[0]);
  for (const auto& elem : vec) {
    result += elem.capacity() * sizeof(elem[0]);
  }

  return result;
}

size_t GfvStorage(const GenericFeatureVector& gfv);

size_t GfvStorage(const vector<GenericFeatureVector>& gfvs);

template <typename T>
size_t DatapointStorage(const Datapoint<T>& dp) {
  return sizeof(dp) + VectorStorage(dp.indices()) + VectorStorage(dp.values());
}

template <typename T>
size_t DatapointStorage(const vector<Datapoint<T> >& dps) {
  size_t result = sizeof(dps);
  for (const auto& dp : dps) {
    result += DatapointStorage(dp);
  }

  return result;
}

std::string GetTcMallocLogString();

#define SCANN_LOG_TCMALLOC() LOG(INFO) << GetTcMallocLogString();
#define SCANN_VLOG_TCMALLOC(verbosity_level) \
  VLOG(verbosity_level) << GetTcMallocLogString();

inline void LogTcMalloc(int verbosity_level = 0) {
  VLOG(verbosity_level) << GetTcMallocLogString();
}

}  // namespace research_scann

#endif


<|start_filename|>cube_unfoldings/unfoldings.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef CUBE_UNFOLDINGS_UNFOLDINGS_H_
#define CUBE_UNFOLDINGS_UNFOLDINGS_H_
#include <cstdlib>

namespace cube_unfoldings {
// Counts the number of unique perfect matchings (under automorphism) in the
// given tree. The tree is assumed to be in a format recognized by nauty
// (sparse6, graph6).
size_t NumberOfUniqueMatchings(const char* tree);
}  // namespace cube_unfoldings

#endif  // CUBE_UNFOLDINGS_UNFOLDINGS_H_


<|start_filename|>scann/scann/projection/identity_projection.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_PROJECTION_IDENTITY_PROJECTION_H_
#define SCANN_PROJECTION_IDENTITY_PROJECTION_H_

#include <cstdint>

#include "scann/data_format/datapoint.h"
#include "scann/projection/projection_base.h"
#include "scann/utils/types.h"

namespace research_scann {

template <typename T>
class IdentityProjection : public Projection<T> {
 public:
  IdentityProjection() {}

  Status ProjectInput(const DatapointPtr<T>& input,
                      Datapoint<float>* projected) const override;
  Status ProjectInput(const DatapointPtr<T>& input,
                      Datapoint<double>* projected) const override;

  int32_t projected_dimensionality() const override { return -1; }

 private:
  template <typename FloatT>
  Status ProjectInputImpl(const DatapointPtr<T>& input,
                          Datapoint<FloatT>* projected) const;
};

SCANN_INSTANTIATE_TYPED_CLASS(extern, IdentityProjection);

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/hashes/internal/stacked_quantizers.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_HASHES_INTERNAL_STACKED_QUANTIZERS_H_
#define SCANN_HASHES_INTERNAL_STACKED_QUANTIZERS_H_

#include <cstdint>

#include "scann/data_format/datapoint.h"
#include "scann/data_format/dataset.h"
#include "scann/distance_measures/distance_measure_base.h"
#include "scann/hashes/asymmetric_hashing2/training_options_base.h"
#include "scann/oss_wrappers/scann_threadpool.h"
#include "scann/utils/common.h"
#include "scann/utils/types.h"

namespace research_scann {
namespace asymmetric_hashing_internal {

class CodesList;

template <typename T>
class StackedQuantizers {
 public:
  using FloatT = FloatingTypeFor<T>;
  using TrainingOptions = asymmetric_hashing2::TrainingOptionsTyped<T>;

  template <typename U>
  using CodebookList = vector<DenseDataset<U>>;
  template <typename U>
  using CodebookListView = ConstSpan<DenseDataset<U>>;

  static Status Hash(const DatapointPtr<T>& input,
                     const ChunkingProjection<T>& projector,
                     const DistanceMeasure& quantization_distance,
                     CodebookListView<FloatT> codebook_list,
                     Datapoint<uint8_t>* output);

  static Status Hash(const DatapointPtr<T>& input,
                     const ChunkingProjection<T>& projector,
                     const DistanceMeasure& quantization_distance,
                     CodebookListView<FloatT> codebook_list,
                     MutableSpan<uint8_t> output);

  static Datapoint<FloatT> Reconstruct(const DatapointPtr<uint8_t>& input,
                                       CodebookListView<FloatT> codebook_list);

  static void Reconstruct(ConstSpan<uint8_t> input,
                          CodebookListView<FloatT> codebook_list,
                          MutableSpan<FloatingTypeFor<T>> output);

  static StatusOr<CodebookList<FloatT>> Train(const DenseDataset<T>& dataset,
                                              const TrainingOptions& opts,
                                              shared_ptr<ThreadPool> pool);

 private:
  static StatusOr<const DenseDataset<double>*> SampledDataset(
      const DenseDataset<T>& dataset, const TrainingOptions& opts,
      DenseDataset<double>* buffer);

  static StatusOr<CodebookList<double>> HierarchicalKMeans(
      const DenseDataset<double>& dataset, const TrainingOptions& opts,
      int num_codebooks, shared_ptr<ThreadPool> pool);

  template <typename U>
  static void GreedilyAssignCodes(const DatapointPtr<U>& input,
                                  const DistanceMeasure& quantization_distance,
                                  CodebookListView<U> codebook_list,
                                  MutableSpan<uint8_t> mutable_codes,
                                  Datapoint<U>* output_residual = nullptr);

  static vector<Datapoint<double>> ComputeUpdatesToCodebook(
      int codebook_i, DimensionIndex dim, int num_centers,
      const DistanceMeasure& quantization_distance, const CodesList& codes_list,
      const DenseDataset<double>& residual);

  static Status InitializeCodes(const DenseDataset<double>& dataset,
                                const DistanceMeasure& quantization_distance,
                                CodebookListView<double> codebook_list,
                                CodesList* codes_list,
                                DenseDataset<double>* residual,
                                ThreadPool* pool);
};

SCANN_INSTANTIATE_TYPED_CLASS(extern, StackedQuantizers);

}  // namespace asymmetric_hashing_internal
}  // namespace research_scann

#endif


<|start_filename|>fairness_and_bias_in_online_selection/unfair-prophet.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "unfair-prophet.h"

namespace fair_secretary {

using std::vector;

SecretaryInstance UnfairProphetAlgorithm::ComputeSolutionOneMinusOneE(
    const vector<SecretaryInstance>& elements,
    const std::vector<std::reference_wrapper<RandomDistribution>> distributions,
    const vector<double>& q) {
  for (int i = 0; i < elements.size(); i++) {
    if (elements[i].value >= distributions[elements[i].type].get().Reverse(
                                 1.0 - (1.0 / elements.size()))) {
      return elements[i];
    }
  }
  return SecretaryInstance(-1, -1);
}

SecretaryInstance UnfairProphetAlgorithm::ComputeSolutionOneHalf(
    const vector<SecretaryInstance>& elements,
    const std::vector<std::reference_wrapper<RandomDistribution>> distributions,
    const vector<double>& q) {
  for (int i = 0; i < elements.size(); i++) {
    if (elements[i].value >=
        distributions[elements[i].type].get().Middle(elements.size())) {
      return elements[i];
    }
  }
  return SecretaryInstance(-1, -1);
}

SecretaryInstance UnfairProphetAlgorithm::ComputeSolutionThreeForth(
    const vector<SecretaryInstance>& elements,
    const std::vector<std::reference_wrapper<RandomDistribution>> distributions,
    const vector<double>& q) {
  for (int i = 0; i < elements.size(); i++) {
    if (elements[i].value >=
        distributions[elements[i].type].get().PThreshold(i)) {
      return elements[i];
    }
  }
  return SecretaryInstance(-1, -1);
}

SecretaryInstance UnfairProphetAlgorithm::ComputeSolutionDiffEq(
    const vector<SecretaryInstance>& elements,
    const std::vector<std::reference_wrapper<RandomDistribution>> distributions,
    const vector<double>& q) {
  // Precomputed Threshold.
  vector<double> diff_solution_50 = {
      1.00E+00, 9.73E-01, 9.45E-01, 9.18E-01, 8.91E-01, 8.63E-01, 8.36E-01,
      8.09E-01, 7.82E-01, 7.56E-01, 7.29E-01, 7.03E-01, 6.76E-01, 6.50E-01,
      6.24E-01, 5.99E-01, 5.74E-01, 5.49E-01, 5.24E-01, 4.99E-01, 4.75E-01,
      4.52E-01, 4.28E-01, 4.05E-01, 3.83E-01, 3.61E-01, 3.39E-01, 3.18E-01,
      2.98E-01, 2.77E-01, 2.58E-01, 2.39E-01, 2.20E-01, 2.02E-01, 1.85E-01,
      1.68E-01, 1.52E-01, 1.36E-01, 1.21E-01, 1.07E-01, 9.32E-02, 8.02E-02,
      6.78E-02, 5.60E-02, 4.50E-02, 3.46E-02, 2.49E-02, 1.59E-02, 7.65E-03,
      1.95E-04};
  vector<double> diff_solution_1000 = {
      1.00E+00, 9.99E-01, 9.97E-01, 9.96E-01, 9.95E-01, 9.93E-01, 9.92E-01,
      9.91E-01, 9.89E-01, 9.88E-01, 9.87E-01, 9.85E-01, 9.84E-01, 9.83E-01,
      9.81E-01, 9.80E-01, 9.79E-01, 9.77E-01, 9.76E-01, 9.74E-01, 9.73E-01,
      9.72E-01, 9.70E-01, 9.69E-01, 9.68E-01, 9.66E-01, 9.65E-01, 9.64E-01,
      9.62E-01, 9.61E-01, 9.60E-01, 9.58E-01, 9.57E-01, 9.56E-01, 9.54E-01,
      9.53E-01, 9.52E-01, 9.50E-01, 9.49E-01, 9.48E-01, 9.46E-01, 9.45E-01,
      9.44E-01, 9.42E-01, 9.41E-01, 9.40E-01, 9.38E-01, 9.37E-01, 9.36E-01,
      9.34E-01, 9.33E-01, 9.32E-01, 9.30E-01, 9.29E-01, 9.28E-01, 9.26E-01,
      9.25E-01, 9.24E-01, 9.22E-01, 9.21E-01, 9.20E-01, 9.18E-01, 9.17E-01,
      9.16E-01, 9.14E-01, 9.13E-01, 9.11E-01, 9.10E-01, 9.09E-01, 9.07E-01,
      9.06E-01, 9.05E-01, 9.03E-01, 9.02E-01, 9.01E-01, 8.99E-01, 8.98E-01,
      8.97E-01, 8.95E-01, 8.94E-01, 8.93E-01, 8.91E-01, 8.90E-01, 8.89E-01,
      8.87E-01, 8.86E-01, 8.85E-01, 8.83E-01, 8.82E-01, 8.81E-01, 8.79E-01,
      8.78E-01, 8.77E-01, 8.75E-01, 8.74E-01, 8.73E-01, 8.71E-01, 8.70E-01,
      8.69E-01, 8.67E-01, 8.66E-01, 8.65E-01, 8.63E-01, 8.62E-01, 8.61E-01,
      8.59E-01, 8.58E-01, 8.57E-01, 8.55E-01, 8.54E-01, 8.53E-01, 8.51E-01,
      8.50E-01, 8.49E-01, 8.47E-01, 8.46E-01, 8.45E-01, 8.43E-01, 8.42E-01,
      8.41E-01, 8.39E-01, 8.38E-01, 8.37E-01, 8.35E-01, 8.34E-01, 8.33E-01,
      8.31E-01, 8.30E-01, 8.29E-01, 8.28E-01, 8.26E-01, 8.25E-01, 8.24E-01,
      8.22E-01, 8.21E-01, 8.20E-01, 8.18E-01, 8.17E-01, 8.16E-01, 8.14E-01,
      8.13E-01, 8.12E-01, 8.10E-01, 8.09E-01, 8.08E-01, 8.06E-01, 8.05E-01,
      8.04E-01, 8.02E-01, 8.01E-01, 8.00E-01, 7.98E-01, 7.97E-01, 7.96E-01,
      7.94E-01, 7.93E-01, 7.92E-01, 7.90E-01, 7.89E-01, 7.88E-01, 7.87E-01,
      7.85E-01, 7.84E-01, 7.83E-01, 7.81E-01, 7.80E-01, 7.79E-01, 7.77E-01,
      7.76E-01, 7.75E-01, 7.73E-01, 7.72E-01, 7.71E-01, 7.69E-01, 7.68E-01,
      7.67E-01, 7.65E-01, 7.64E-01, 7.63E-01, 7.62E-01, 7.60E-01, 7.59E-01,
      7.58E-01, 7.56E-01, 7.55E-01, 7.54E-01, 7.52E-01, 7.51E-01, 7.50E-01,
      7.48E-01, 7.47E-01, 7.46E-01, 7.45E-01, 7.43E-01, 7.42E-01, 7.41E-01,
      7.39E-01, 7.38E-01, 7.37E-01, 7.35E-01, 7.34E-01, 7.33E-01, 7.31E-01,
      7.30E-01, 7.29E-01, 7.28E-01, 7.26E-01, 7.25E-01, 7.24E-01, 7.22E-01,
      7.21E-01, 7.20E-01, 7.18E-01, 7.17E-01, 7.16E-01, 7.15E-01, 7.13E-01,
      7.12E-01, 7.11E-01, 7.09E-01, 7.08E-01, 7.07E-01, 7.06E-01, 7.04E-01,
      7.03E-01, 7.02E-01, 7.00E-01, 6.99E-01, 6.98E-01, 6.96E-01, 6.95E-01,
      6.94E-01, 6.93E-01, 6.91E-01, 6.90E-01, 6.89E-01, 6.87E-01, 6.86E-01,
      6.85E-01, 6.84E-01, 6.82E-01, 6.81E-01, 6.80E-01, 6.78E-01, 6.77E-01,
      6.76E-01, 6.75E-01, 6.73E-01, 6.72E-01, 6.71E-01, 6.69E-01, 6.68E-01,
      6.67E-01, 6.66E-01, 6.64E-01, 6.63E-01, 6.62E-01, 6.61E-01, 6.59E-01,
      6.58E-01, 6.57E-01, 6.55E-01, 6.54E-01, 6.53E-01, 6.52E-01, 6.50E-01,
      6.49E-01, 6.48E-01, 6.47E-01, 6.45E-01, 6.44E-01, 6.43E-01, 6.41E-01,
      6.40E-01, 6.39E-01, 6.38E-01, 6.36E-01, 6.35E-01, 6.34E-01, 6.33E-01,
      6.31E-01, 6.30E-01, 6.29E-01, 6.28E-01, 6.26E-01, 6.25E-01, 6.24E-01,
      6.22E-01, 6.21E-01, 6.20E-01, 6.19E-01, 6.17E-01, 6.16E-01, 6.15E-01,
      6.14E-01, 6.12E-01, 6.11E-01, 6.10E-01, 6.09E-01, 6.07E-01, 6.06E-01,
      6.05E-01, 6.04E-01, 6.02E-01, 6.01E-01, 6.00E-01, 5.99E-01, 5.97E-01,
      5.96E-01, 5.95E-01, 5.94E-01, 5.92E-01, 5.91E-01, 5.90E-01, 5.89E-01,
      5.87E-01, 5.86E-01, 5.85E-01, 5.84E-01, 5.82E-01, 5.81E-01, 5.80E-01,
      5.79E-01, 5.78E-01, 5.76E-01, 5.75E-01, 5.74E-01, 5.73E-01, 5.71E-01,
      5.70E-01, 5.69E-01, 5.68E-01, 5.66E-01, 5.65E-01, 5.64E-01, 5.63E-01,
      5.61E-01, 5.60E-01, 5.59E-01, 5.58E-01, 5.57E-01, 5.55E-01, 5.54E-01,
      5.53E-01, 5.52E-01, 5.50E-01, 5.49E-01, 5.48E-01, 5.47E-01, 5.46E-01,
      5.44E-01, 5.43E-01, 5.42E-01, 5.41E-01, 5.40E-01, 5.38E-01, 5.37E-01,
      5.36E-01, 5.35E-01, 5.33E-01, 5.32E-01, 5.31E-01, 5.30E-01, 5.29E-01,
      5.27E-01, 5.26E-01, 5.25E-01, 5.24E-01, 5.23E-01, 5.21E-01, 5.20E-01,
      5.19E-01, 5.18E-01, 5.17E-01, 5.15E-01, 5.14E-01, 5.13E-01, 5.12E-01,
      5.11E-01, 5.09E-01, 5.08E-01, 5.07E-01, 5.06E-01, 5.05E-01, 5.03E-01,
      5.02E-01, 5.01E-01, 5.00E-01, 4.99E-01, 4.97E-01, 4.96E-01, 4.95E-01,
      4.94E-01, 4.93E-01, 4.92E-01, 4.90E-01, 4.89E-01, 4.88E-01, 4.87E-01,
      4.86E-01, 4.84E-01, 4.83E-01, 4.82E-01, 4.81E-01, 4.80E-01, 4.79E-01,
      4.77E-01, 4.76E-01, 4.75E-01, 4.74E-01, 4.73E-01, 4.72E-01, 4.70E-01,
      4.69E-01, 4.68E-01, 4.67E-01, 4.66E-01, 4.65E-01, 4.63E-01, 4.62E-01,
      4.61E-01, 4.60E-01, 4.59E-01, 4.58E-01, 4.56E-01, 4.55E-01, 4.54E-01,
      4.53E-01, 4.52E-01, 4.51E-01, 4.50E-01, 4.48E-01, 4.47E-01, 4.46E-01,
      4.45E-01, 4.44E-01, 4.43E-01, 4.41E-01, 4.40E-01, 4.39E-01, 4.38E-01,
      4.37E-01, 4.36E-01, 4.35E-01, 4.34E-01, 4.32E-01, 4.31E-01, 4.30E-01,
      4.29E-01, 4.28E-01, 4.27E-01, 4.26E-01, 4.24E-01, 4.23E-01, 4.22E-01,
      4.21E-01, 4.20E-01, 4.19E-01, 4.18E-01, 4.17E-01, 4.15E-01, 4.14E-01,
      4.13E-01, 4.12E-01, 4.11E-01, 4.10E-01, 4.09E-01, 4.08E-01, 4.06E-01,
      4.05E-01, 4.04E-01, 4.03E-01, 4.02E-01, 4.01E-01, 4.00E-01, 3.99E-01,
      3.98E-01, 3.96E-01, 3.95E-01, 3.94E-01, 3.93E-01, 3.92E-01, 3.91E-01,
      3.90E-01, 3.89E-01, 3.88E-01, 3.87E-01, 3.85E-01, 3.84E-01, 3.83E-01,
      3.82E-01, 3.81E-01, 3.80E-01, 3.79E-01, 3.78E-01, 3.77E-01, 3.76E-01,
      3.75E-01, 3.74E-01, 3.72E-01, 3.71E-01, 3.70E-01, 3.69E-01, 3.68E-01,
      3.67E-01, 3.66E-01, 3.65E-01, 3.64E-01, 3.63E-01, 3.62E-01, 3.61E-01,
      3.60E-01, 3.58E-01, 3.57E-01, 3.56E-01, 3.55E-01, 3.54E-01, 3.53E-01,
      3.52E-01, 3.51E-01, 3.50E-01, 3.49E-01, 3.48E-01, 3.47E-01, 3.46E-01,
      3.45E-01, 3.44E-01, 3.43E-01, 3.41E-01, 3.40E-01, 3.39E-01, 3.38E-01,
      3.37E-01, 3.36E-01, 3.35E-01, 3.34E-01, 3.33E-01, 3.32E-01, 3.31E-01,
      3.30E-01, 3.29E-01, 3.28E-01, 3.27E-01, 3.26E-01, 3.25E-01, 3.24E-01,
      3.23E-01, 3.22E-01, 3.21E-01, 3.20E-01, 3.19E-01, 3.18E-01, 3.17E-01,
      3.16E-01, 3.15E-01, 3.14E-01, 3.13E-01, 3.12E-01, 3.10E-01, 3.09E-01,
      3.08E-01, 3.07E-01, 3.06E-01, 3.05E-01, 3.04E-01, 3.03E-01, 3.02E-01,
      3.01E-01, 3.00E-01, 2.99E-01, 2.98E-01, 2.97E-01, 2.96E-01, 2.95E-01,
      2.94E-01, 2.93E-01, 2.92E-01, 2.91E-01, 2.90E-01, 2.89E-01, 2.88E-01,
      2.87E-01, 2.86E-01, 2.85E-01, 2.84E-01, 2.84E-01, 2.83E-01, 2.82E-01,
      2.81E-01, 2.80E-01, 2.79E-01, 2.78E-01, 2.77E-01, 2.76E-01, 2.75E-01,
      2.74E-01, 2.73E-01, 2.72E-01, 2.71E-01, 2.70E-01, 2.69E-01, 2.68E-01,
      2.67E-01, 2.66E-01, 2.65E-01, 2.64E-01, 2.63E-01, 2.62E-01, 2.61E-01,
      2.60E-01, 2.59E-01, 2.58E-01, 2.57E-01, 2.56E-01, 2.56E-01, 2.55E-01,
      2.54E-01, 2.53E-01, 2.52E-01, 2.51E-01, 2.50E-01, 2.49E-01, 2.48E-01,
      2.47E-01, 2.46E-01, 2.45E-01, 2.44E-01, 2.43E-01, 2.42E-01, 2.41E-01,
      2.41E-01, 2.40E-01, 2.39E-01, 2.38E-01, 2.37E-01, 2.36E-01, 2.35E-01,
      2.34E-01, 2.33E-01, 2.32E-01, 2.31E-01, 2.30E-01, 2.30E-01, 2.29E-01,
      2.28E-01, 2.27E-01, 2.26E-01, 2.25E-01, 2.24E-01, 2.23E-01, 2.22E-01,
      2.21E-01, 2.21E-01, 2.20E-01, 2.19E-01, 2.18E-01, 2.17E-01, 2.16E-01,
      2.15E-01, 2.14E-01, 2.13E-01, 2.13E-01, 2.12E-01, 2.11E-01, 2.10E-01,
      2.09E-01, 2.08E-01, 2.07E-01, 2.06E-01, 2.06E-01, 2.05E-01, 2.04E-01,
      2.03E-01, 2.02E-01, 2.01E-01, 2.00E-01, 1.99E-01, 1.99E-01, 1.98E-01,
      1.97E-01, 1.96E-01, 1.95E-01, 1.94E-01, 1.93E-01, 1.93E-01, 1.92E-01,
      1.91E-01, 1.90E-01, 1.89E-01, 1.88E-01, 1.87E-01, 1.87E-01, 1.86E-01,
      1.85E-01, 1.84E-01, 1.83E-01, 1.82E-01, 1.82E-01, 1.81E-01, 1.80E-01,
      1.79E-01, 1.78E-01, 1.77E-01, 1.77E-01, 1.76E-01, 1.75E-01, 1.74E-01,
      1.73E-01, 1.73E-01, 1.72E-01, 1.71E-01, 1.70E-01, 1.69E-01, 1.68E-01,
      1.68E-01, 1.67E-01, 1.66E-01, 1.65E-01, 1.64E-01, 1.64E-01, 1.63E-01,
      1.62E-01, 1.61E-01, 1.60E-01, 1.60E-01, 1.59E-01, 1.58E-01, 1.57E-01,
      1.56E-01, 1.56E-01, 1.55E-01, 1.54E-01, 1.53E-01, 1.53E-01, 1.52E-01,
      1.51E-01, 1.50E-01, 1.49E-01, 1.49E-01, 1.48E-01, 1.47E-01, 1.46E-01,
      1.46E-01, 1.45E-01, 1.44E-01, 1.43E-01, 1.43E-01, 1.42E-01, 1.41E-01,
      1.40E-01, 1.39E-01, 1.39E-01, 1.38E-01, 1.37E-01, 1.36E-01, 1.36E-01,
      1.35E-01, 1.34E-01, 1.33E-01, 1.33E-01, 1.32E-01, 1.31E-01, 1.31E-01,
      1.30E-01, 1.29E-01, 1.28E-01, 1.28E-01, 1.27E-01, 1.26E-01, 1.25E-01,
      1.25E-01, 1.24E-01, 1.23E-01, 1.22E-01, 1.22E-01, 1.21E-01, 1.20E-01,
      1.20E-01, 1.19E-01, 1.18E-01, 1.17E-01, 1.17E-01, 1.16E-01, 1.15E-01,
      1.15E-01, 1.14E-01, 1.13E-01, 1.13E-01, 1.12E-01, 1.11E-01, 1.10E-01,
      1.10E-01, 1.09E-01, 1.08E-01, 1.08E-01, 1.07E-01, 1.06E-01, 1.06E-01,
      1.05E-01, 1.04E-01, 1.04E-01, 1.03E-01, 1.02E-01, 1.02E-01, 1.01E-01,
      1.00E-01, 9.95E-02, 9.88E-02, 9.82E-02, 9.75E-02, 9.68E-02, 9.62E-02,
      9.55E-02, 9.49E-02, 9.42E-02, 9.35E-02, 9.29E-02, 9.22E-02, 9.16E-02,
      9.09E-02, 9.03E-02, 8.96E-02, 8.90E-02, 8.83E-02, 8.77E-02, 8.71E-02,
      8.64E-02, 8.58E-02, 8.51E-02, 8.45E-02, 8.39E-02, 8.32E-02, 8.26E-02,
      8.20E-02, 8.13E-02, 8.07E-02, 8.01E-02, 7.95E-02, 7.89E-02, 7.82E-02,
      7.76E-02, 7.70E-02, 7.64E-02, 7.58E-02, 7.52E-02, 7.45E-02, 7.39E-02,
      7.33E-02, 7.27E-02, 7.21E-02, 7.15E-02, 7.09E-02, 7.03E-02, 6.97E-02,
      6.91E-02, 6.85E-02, 6.79E-02, 6.73E-02, 6.68E-02, 6.62E-02, 6.56E-02,
      6.50E-02, 6.44E-02, 6.38E-02, 6.32E-02, 6.27E-02, 6.21E-02, 6.15E-02,
      6.09E-02, 6.04E-02, 5.98E-02, 5.92E-02, 5.87E-02, 5.81E-02, 5.75E-02,
      5.70E-02, 5.64E-02, 5.59E-02, 5.53E-02, 5.47E-02, 5.42E-02, 5.36E-02,
      5.31E-02, 5.25E-02, 5.20E-02, 5.14E-02, 5.09E-02, 5.03E-02, 4.98E-02,
      4.93E-02, 4.87E-02, 4.82E-02, 4.77E-02, 4.71E-02, 4.66E-02, 4.61E-02,
      4.55E-02, 4.50E-02, 4.45E-02, 4.40E-02, 4.34E-02, 4.29E-02, 4.24E-02,
      4.19E-02, 4.14E-02, 4.08E-02, 4.03E-02, 3.98E-02, 3.93E-02, 3.88E-02,
      3.83E-02, 3.78E-02, 3.73E-02, 3.68E-02, 3.63E-02, 3.58E-02, 3.53E-02,
      3.48E-02, 3.43E-02, 3.38E-02, 3.33E-02, 3.29E-02, 3.24E-02, 3.19E-02,
      3.14E-02, 3.09E-02, 3.04E-02, 3.00E-02, 2.95E-02, 2.90E-02, 2.85E-02,
      2.81E-02, 2.76E-02, 2.71E-02, 2.67E-02, 2.62E-02, 2.57E-02, 2.53E-02,
      2.48E-02, 2.44E-02, 2.39E-02, 2.35E-02, 2.30E-02, 2.26E-02, 2.21E-02,
      2.17E-02, 2.12E-02, 2.08E-02, 2.03E-02, 1.99E-02, 1.94E-02, 1.90E-02,
      1.86E-02, 1.81E-02, 1.77E-02, 1.73E-02, 1.69E-02, 1.64E-02, 1.60E-02,
      1.56E-02, 1.52E-02, 1.47E-02, 1.43E-02, 1.39E-02, 1.35E-02, 1.31E-02,
      1.27E-02, 1.23E-02, 1.19E-02, 1.15E-02, 1.10E-02, 1.06E-02, 1.02E-02,
      9.85E-03, 9.45E-03, 9.06E-03, 8.67E-03, 8.28E-03, 7.89E-03, 7.50E-03,
      7.12E-03, 6.73E-03, 6.35E-03, 5.97E-03, 5.60E-03, 5.22E-03, 4.85E-03,
      4.48E-03, 4.11E-03, 3.74E-03, 3.38E-03, 3.01E-03, 2.65E-03, 2.30E-03,
      1.94E-03, 1.59E-03, 1.24E-03, 8.87E-04, 5.40E-04, 1.95E-04};
  auto diff_solution =
      elements.size() == 50 ? diff_solution_50 : diff_solution_1000;
  for (int i = 0; i < elements.size(); i++) {
    if (elements[i].value >=
        distributions[elements[i].type].get().Reverse(
            std::pow(diff_solution[i], 1.0 / (elements.size() - 1)))) {
      return elements[i];
    }
  }
  return SecretaryInstance(-1, -1);
}

}  // namespace fair_secretary


<|start_filename|>aptamers_mlpd/preprocess/merge/merge_measurements.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "xxx/merge/merge_measurements.h"

#include "glog/logging.h"
#include "pipeline/flume/public/flume.h"
#include "absl/strings/string_view.h"

namespace research_biology {
namespace aptamers {

namespace {

// MergeFn is the function that merges multiple MeasurementEntry values
// into a single one by taking the union of the counts.
//
// This is a stateful operation because we need to the join tags
// for accessing the fields in a JoinResult.
class MergeFn : public flume::DoFn<JoinEntry, MeasurementEntry> {
 public:
  explicit MergeFn(const std::vector<JoinTag> &tags)
      : tags_(tags) {}

  // Do performs the merge operation.
  void Do(const JoinEntry &in,
          const flume::EmitFn<MeasurementEntry> &emit_fn) override {
    const flume::JoinResult &value = in.value();

    // Combine the counts values.  We assume that the columns are disjoint,
    // meaning that the union never encounters the same map key in multiple
    // input tables.
    Measurement result;
    proto2::Map<int32, int32> *counts = result.mutable_counts();
    for (const JoinTag &tag : tags_) {
      flume::Stream<Measurement> vs = value.Get(tag);
      for (; !vs.at_end(); ++vs) {
        const Measurement m = *vs;
        for (const auto &value : m.counts()) {
          const std::pair<proto2::Map<int32, int32>::iterator, bool> p =
              counts->insert(value);
          if (!p.second) {
            LOG(ERROR) << "Duplicate key: " << value.first;
          }
        }
      }
    }
    emit_fn.Emit(MeasurementEntry(in.key(), result));
  }

  // Encode the state, for saving an instance of this class.
  void EncodeState(string *encoded_state) const override {
    std::vector<string> keys;
    for (const auto &tag : tags_) {
      keys.push_back(tag.id());
    }
    flume::VectorsOf(flume::Strings()).Encode(keys, encoded_state);
  }

  // Recreate an instance of this class from a copy of the state.
  static MergeFn *DecodeState(absl::string_view encoded_state) {
    std::vector<string> keys;
    CHECK(flume::VectorsOf(flume::Strings()).Decode(encoded_state, &keys));
    std::vector<JoinTag> tags;
    for (const string &key : keys) {
      tags.push_back(JoinTag(key));
    }
    return new MergeFn(tags);
  }

 private:
  const std::vector<JoinTag> tags_;

  REGISTER_AS_STATEFUL_FN(MergeFn);
};

}  // namespace

// Perform the join.
MeasurementTable MergeMeasurements(
    const std::vector<JoinTag> &keys, const JoinTable &in) {
  return in.ParDo("merge", new MergeFn(keys));
}

}  // namespace aptamers
}  // namespace research_biology


<|start_filename|>simulation_research/traffic/libsumo_demo_main.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Simple main to demonstrate use of libsumo interface.
//
// The inputs to the job can be generated by following the steps at:
// https://sumo.dlr.de/wiki/Tutorials/Hello_Sumo. Alternatively, these files
// have been pregenerated and are hecked in at
// //third_party/sumo/sumo-1.1.0/docs/tutorial/hello/data.
//
// Example invocation:
//     SUMO_HOME=${SUMO_PATH}/sumo/sumo-1.1.0 \
//         blaze-bin/research/simulation/traffic/libsumo_demo_main \
//         --config_file=third_party/sumo/sumo-1.1.0/docs/tutorial/hello/data\
//         /hello.sumocfg --num_steps=100 --alsologtostderr
// NB: The way that SUMO imports headers is problematic. Vehicle.h requires
// PositionVector.h, so we need to include the header before we include
// Vehicle.h.
// TODO(yusef): Devise a better solution to deal with such imports.
#include <utils/geom/PositionVector.h>

#include <gflags/gflags.h>

#include "glog/logging.h"
#include "absl/flags/flag.h"
#include "third_party/sumo/sumo_1_6_0/src/libsumo/Simulation.h"
#include "third_party/sumo/sumo_1_6_0/src/libsumo/Vehicle.h"

DEFINE_string(
    config_file,
    "third_party/sumo/sumo_1_6_0/docs/tutorial/hello/data/hello.sumocfg",
    "Simulation configuration.");
DEFINE_int64(num_steps, 0, "Number of steps.");

namespace research {
namespace simulation {
int Main(int argc, char** argv) {
  std::vector<std::string> args = {absl::GetFlag(FLAGS_config_file)};
  // The following line will crash (throw an exception) if the args are invalid.
  libsumo::Simulation::load(args);

  // We demonstrate two simple use cases of the libsumo API.
  // First, we use the subscription API. 0x7a corresponds to a list of arrived
  // vehicles (https://sumo.dlr.de/wiki/TraCI/Simulation_Value_Retrieval).
  // Second, we call the Vehicle API directly without a subscription to retrieve
  // current vehicles and their speeds.
  libsumo::Simulation::subscribe({0x7a}, 0, 10000);
  for (int i = 0; i < absl::GetFlag(FLAGS_num_steps); ++i) {
    libsumo::Simulation::step();
    const auto results = libsumo::Simulation::getSubscriptionResults();
    for (const auto& pair : results) {
      LOG(INFO) << pair.first;
      LOG(INFO) << pair.second->getString();
    }
    for (const auto& vehicle : libsumo::Vehicle::getIDList()) {
      LOG(INFO) << vehicle << " speed: " << libsumo::Vehicle::getSpeed(vehicle);
    }
  }
  return 0;
}

}  // namespace simulation
}  // namespace research

int main(int argc, char** argv) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  return research::simulation::Main(argc, argv);
}


<|start_filename|>scann/scann/distance_measures/one_to_one/l2_distance_sse4.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/distance_measures/one_to_one/l2_distance_sse4.h"

#include <cstdint>
#ifdef __x86_64__

#include "scann/utils/intrinsics/sse4.h"

namespace research_scann {
namespace l2_internal {

template <typename Byte, typename SseFuncs>
SCANN_SSE4_INLINE double DenseSquaredL2DistanceByteImpl(const Byte* aptr,
                                                        const Byte* bptr,
                                                        size_t length) {
  const Byte* aend = aptr + length;

  auto as_m128i = [](const Byte* x) SCANN_SSE4_INLINE_LAMBDA -> __m128i* {
    return reinterpret_cast<__m128i*>(const_cast<Byte*>(x));
  };

  auto get_terms = [&as_m128i](const Byte* aptr, const Byte* bptr)
                       SCANN_SSE4_INLINE_LAMBDA {
                         __m128i avals = _mm_loadu_si128(as_m128i(aptr));
                         __m128i bvals = _mm_loadu_si128(as_m128i(bptr));
                         __m128i diff = SseFuncs::AbsDiff(avals, bvals);

                         __m128i lower = SseFuncs::ZeroExtendLower8To16(diff);
                         __m128i upper = SseFuncs::ZeroExtendUpper8To16(diff);
                         lower = _mm_mullo_epi16(lower, lower);
                         upper = _mm_mullo_epi16(upper, upper);
                         return std::make_pair(lower, upper);
                       };

  uint32_t scalar_accumulator = 0;
  if (aptr + 4 <= aend) {
    __m128i accumulator0 = _mm_setzero_si128();
    __m128i accumulator1 = _mm_setzero_si128();

    auto do_accumulations = [&accumulator0, &accumulator1](
                                __m128i term) SCANN_SSE4_INLINE_LAMBDA {
      accumulator0 =
          _mm_add_epi32(accumulator0, SseFuncs::ZeroExtendLower16To32(term));
      accumulator1 =
          _mm_add_epi32(accumulator1, SseFuncs::ZeroExtendUpper16To32(term));
    };

    for (; aptr + 16 <= aend; aptr += 16, bptr += 16) {
      const pair<__m128i, __m128i> terms = get_terms(aptr, bptr);
      do_accumulations(terms.first);
      do_accumulations(terms.second);
    }

    if (aptr + 8 <= aend) {
      __m128i avals = _mm_loadl_epi64(as_m128i(aptr));
      __m128i bvals = _mm_loadl_epi64(as_m128i(bptr));
      __m128i diff = SseFuncs::AbsDiff(avals, bvals);
      __m128i lower = SseFuncs::ZeroExtendLower8To16(diff);
      lower = _mm_mullo_epi16(lower, lower);
      do_accumulations(lower);
      aptr += 8;
      bptr += 8;
    }

    if (aptr + 4 <= aend) {
      __m128i avals = _mm_cvtsi32_si128(*reinterpret_cast<const int*>(aptr));
      __m128i bvals = _mm_cvtsi32_si128(*reinterpret_cast<const int*>(bptr));
      __m128i diff = SseFuncs::AbsDiff(avals, bvals);
      __m128i lower = _mm_unpacklo_epi8(diff, _mm_setzero_si128());
      lower = _mm_mullo_epi16(lower, lower);
      do_accumulations(lower);
      aptr += 4;
      bptr += 4;
    }

    scalar_accumulator =
        SseFuncs::HorizontalSum32(_mm_add_epi32(accumulator0, accumulator1));
  }

  DCHECK_LT(aend - aptr, 4);
  for (; aptr < aend; ++aptr, ++bptr) {
    int32_t diff = static_cast<int32_t>(*aptr) - static_cast<int32_t>(*bptr);
    scalar_accumulator += diff * diff;
  }

  return static_cast<double>(scalar_accumulator);
}

class SseFunctionsSse4 {
 public:
  SCANN_SSE4_INLINE static __m128i ZeroExtendLower8To16(__m128i v) {
    return _mm_cvtepu8_epi16(v);
  }

  SCANN_SSE4_INLINE static __m128i ZeroExtendLower16To32(__m128i v) {
    return _mm_cvtepu16_epi32(v);
  }

  SCANN_SSE4_INLINE static __m128i ZeroExtendUpper8To16(__m128i v) {
    return _mm_unpackhi_epi8(v, _mm_setzero_si128());
  }

  SCANN_SSE4_INLINE static __m128i ZeroExtendUpper16To32(__m128i v) {
    return _mm_unpackhi_epi16(v, _mm_setzero_si128());
  }

  SCANN_SSE4_INLINE static uint32_t HorizontalSum32(__m128i v) {
    v = _mm_add_epi32(v, _mm_srli_si128(v, 8));
    v = _mm_add_epi32(v, _mm_srli_si128(v, 4));
    return _mm_cvtsi128_si32(v);
  }
};

class SignedSquaredL2SseFunctionsSse4 : public SseFunctionsSse4 {
 public:
  SCANN_SSE4_INLINE static __m128i AbsDiff(__m128i a, __m128i b) {
    return _mm_sub_epi8(_mm_max_epi8(a, b), _mm_min_epi8(a, b));
  }
};

class UnsignedSquaredL2SseFunctionsSse4 : public SseFunctionsSse4 {
 public:
  SCANN_SSE4_INLINE static __m128i AbsDiff(__m128i a, __m128i b) {
    return _mm_sub_epi8(_mm_max_epu8(a, b), _mm_min_epu8(a, b));
  }
};

SCANN_SSE4_OUTLINE double DenseSquaredL2DistanceSse4(
    const DatapointPtr<uint8_t>& a, const DatapointPtr<uint8_t>& b) {
  DCHECK_EQ(a.nonzero_entries(), b.nonzero_entries());
  DCHECK(a.IsDense());
  DCHECK(b.IsDense());
  return DenseSquaredL2DistanceByteImpl<uint8_t,
                                        UnsignedSquaredL2SseFunctionsSse4>(
      a.values(), b.values(), a.nonzero_entries());
}

SCANN_SSE4_OUTLINE double DenseSquaredL2DistanceSse4(
    const DatapointPtr<int8_t>& a, const DatapointPtr<int8_t>& b) {
  DCHECK_EQ(a.nonzero_entries(), b.nonzero_entries());
  DCHECK(a.IsDense());
  DCHECK(b.IsDense());
  return DenseSquaredL2DistanceByteImpl<int8_t,
                                        SignedSquaredL2SseFunctionsSse4>(
      a.values(), b.values(), a.nonzero_entries());
}

SCANN_SSE4_OUTLINE double DenseSquaredL2DistanceSse4(
    const DatapointPtr<float>& a, const DatapointPtr<float>& b) {
  DCHECK_EQ(a.nonzero_entries(), b.nonzero_entries());
  DCHECK(a.IsDense());
  DCHECK(b.IsDense());

  auto get_terms = [](const float* aptr, const float* bptr)
                       SCANN_SSE4_INLINE_LAMBDA {
                         __m128 avals = _mm_loadu_ps(aptr);
                         __m128 bvals = _mm_loadu_ps(bptr);
                         __m128 diff = _mm_sub_ps(avals, bvals);
                         return _mm_mul_ps(diff, diff);
                       };

  const float* aptr = a.values();
  const float* bptr = b.values();
  const float* aend = aptr + a.nonzero_entries();
  __m128 accumulator = _mm_setzero_ps();
  if (aptr + 8 <= aend) {
    __m128 accumulator0 = get_terms(aptr, bptr);
    __m128 accumulator1 = get_terms(aptr + 4, bptr + 4);
    aptr += 8;
    bptr += 8;
    for (; aptr + 8 <= aend; aptr += 8, bptr += 8) {
      accumulator0 = _mm_add_ps(accumulator0, get_terms(aptr, bptr));
      accumulator1 = _mm_add_ps(accumulator1, get_terms(aptr + 4, bptr + 4));
    }

    accumulator = _mm_add_ps(accumulator0, accumulator1);
  }

  if (aptr + 4 <= aend) {
    accumulator = _mm_add_ps(accumulator, get_terms(aptr, bptr));
    aptr += 4;
    bptr += 4;
  }

  if (aptr + 2 <= aend) {
    __m128 avals = _mm_setzero_ps();
    __m128 bvals = _mm_setzero_ps();
    avals = _mm_loadh_pi(avals, reinterpret_cast<const __m64*>(aptr));
    bvals = _mm_loadh_pi(bvals, reinterpret_cast<const __m64*>(bptr));
    __m128 diff = _mm_sub_ps(avals, bvals);
    __m128 squared = _mm_mul_ps(diff, diff);
    accumulator = _mm_add_ps(accumulator, squared);
    aptr += 2;
    bptr += 2;
  }

  if (aptr < aend) {
    accumulator[0] += (aptr[0] - bptr[0]) * (aptr[0] - bptr[0]);
  }

  accumulator = _mm_hadd_ps(accumulator, accumulator);
  accumulator = _mm_hadd_ps(accumulator, accumulator);
  return accumulator[0];
}

SCANN_SSE4_OUTLINE double DenseSquaredL2DistanceSse4(
    const DatapointPtr<double>& a, const DatapointPtr<double>& b) {
  DCHECK_EQ(a.nonzero_entries(), b.nonzero_entries());
  DCHECK(a.IsDense());
  DCHECK(b.IsDense());

  auto get_terms = [](const double* aptr, const double* bptr)
                       SCANN_SSE4_INLINE_LAMBDA {
                         __m128d avals = _mm_loadu_pd(aptr);
                         __m128d bvals = _mm_loadu_pd(bptr);
                         __m128d diff = _mm_sub_pd(avals, bvals);
                         return _mm_mul_pd(diff, diff);
                       };

  const double* aptr = a.values();
  const double* bptr = b.values();
  const double* aend = aptr + a.nonzero_entries();
  __m128d accumulator = _mm_setzero_pd();
  if (aptr + 4 <= aend) {
    __m128d accumulator0 = get_terms(aptr, bptr);
    __m128d accumulator1 = get_terms(aptr + 2, bptr + 2);
    aptr += 4;
    bptr += 4;
    for (; aptr + 4 <= aend; aptr += 4, bptr += 4) {
      accumulator0 = _mm_add_pd(accumulator0, get_terms(aptr, bptr));
      accumulator1 = _mm_add_pd(accumulator1, get_terms(aptr + 2, bptr + 2));
    }

    accumulator = _mm_add_pd(accumulator0, accumulator1);
  }

  if (aptr + 2 <= aend) {
    accumulator = _mm_add_pd(accumulator, get_terms(aptr, bptr));
    aptr += 2;
    bptr += 2;
  }

  accumulator = _mm_hadd_pd(accumulator, accumulator);
  double result = accumulator[0];

  if (aptr < aend) {
    const double diff = *aptr - *bptr;
    result += diff * diff;
  }

  return result;
}

}  // namespace l2_internal
}  // namespace research_scann

#endif


<|start_filename|>npy_array/npy_array/npy_array.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef NPY_ARRAY_NPY_ARRAY_NPY_ARRAY_H_
#define NPY_ARRAY_NPY_ARRAY_NPY_ARRAY_H_

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <cstdio>
#include <string>
#include <utility>
#include <vector>

#include "glog/logging.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "array/array.h"

namespace npy_array {

// This header exposes a single function, `SerializeToNpyString`, whose behavior
// can be configured with `NpySerializeOptions`.

struct NpySerializeOptions {
  // # In Numpy, shapes are indexed outermost to innermost. #
  // Let npy_shape = (480, 640, 3), with axis names (i, j, k).
  // i is the outermost axis with extent 480.
  // k is the innermost axis with extent 3.
  //
  // In the default layout (compact, C-contiguous), i changes least frequently
  // and k changes most frequently. (This corresponds to an interleaved image,
  // where k names a color channel.)
  //
  // In the alternative layout (compact, F-contiguous aka fortran_order), this
  // is reversed: i changes *most* frequently and k changes *least* frequently.
  //
  // # In nda, shapes are indexed innermost to outermost. #
  // Let nda_shape = (1920, 1080, 3) with axis names (x, y, c).
  // x is the innermost axis with extent 1920.
  // c is the outermost axis with extent 3.
  //
  // In the default layout (compact), x changes most frequently and c changes
  // least frequently. (This corresponds to an planar image, where c names a
  // color channel.)

  // If `reverse_axes` is true, then the ordering of the axes in the serialized
  // NPY file will be reverse that of nda. I.e., if the nda::shape is {1920,
  // 1080, 3}, the NPY shape will be (3, 1080, 1920). This leads to natural
  // interpretations if, in nda, we index:
  // - Planar images as (x, y, c):
  //   This yields an npy indexing of (c, y, x) aka (k, i, j).
  // - Interleaved images as (c, x, y):
  //   This yields an npy indexing of (y, x, c) aka (i, j, c).
  //   - Unfortunately, this does *not* lead to a natural ordering if in nda, we
  //     index interleaved images as (x, y, c). To reorder (x, y, c) as (c, x,
  //     y), call nda::reorder<2, 0, 1>(src).
  //
  // If `reverse_axes` is false, then the ordering of axes is preserved.
  bool reverse_axes = true;
};

// Serializes `src` to a std::string in the NPY file format
// (https://numpy.org/devdocs/reference/generated/numpy.lib.format.html).
//
// DataType may be any type specialized by NpyDataTypeString(). This includes
// all signed and unsigned fixed width types in C++ (8, 16, 32, and 64 bits), as
// we as float (32 bits), double (64 bits), and float16 (see NpyFloat16 below).
//
// Caveats:
// - Non-zero mins are not preserved since the format does not support it.
// - If `src.empty()`, returns an empty string since the NPY format does not
//   support empty arrays (the shape "()" refers to a scalar).
//
// Future work:
// - Structured data.
// - Version 3.0, which uses optional utf8 strings to name structured data axes.
template <typename DataType, typename ShapeType>
std::string SerializeToNpyString(
    nda::array_ref<DataType, ShapeType> src,
    const NpySerializeOptions& options = NpySerializeOptions());

// A stand-in type for the purposes of template specialization that allows the
// client to use any float16 library. To serialize a float16 buffer, first
// reinterpret_cast the client pointer to const NpyFloat16* and wrap it in an
// nda::array_ref<const NpyFloat16>.
//
// If your source buffer is already an nda::array or nda::array_ref, use
// nda:reinterpret<const NpyFloat16>(src), which reinterprets the pointer type
// and preserves the shape.
struct NpyFloat16 {
  uint16_t u16;
};
static_assert(sizeof(NpyFloat16) == 2, "sizeof(NpyFloat16) should be 2 bytes.");

namespace internal {
// Returns "<" if this machine is little-endian and ">" if this machine is
// big-endian.
std::string NpyEndiannessString();

// Returns the single-character type string for the given DataType. See:
// https://numpy.org/devdocs/reference/arrays.interface.html#arrays-interface.
//
// This function template is explicitly specialized only for supported types.
template <typename DataType>
std::string NpyDataTypeString();

// Returns the NPY "descr string" describing DataType (it is a string that can
// be passed to the constructor of np.dtype).
//
// TODO(jiawen): This can be explicitly specialized too.
template <typename DataType>
std::string NpyDescrString() {
  return absl::StrCat(NpyEndiannessString(), NpyDataTypeString<DataType>(),
                      sizeof(DataType));
}

// Converts array shape to a vector.
template <typename ShapeType>
std::vector<size_t> NpyShapeVector(ShapeType shape) {
  std::vector<size_t> dims;
  for (size_t i = 0; i < ShapeType::rank(); ++i) {
    dims.push_back(shape.dim(i).extent());
  }
  return dims;
}

// Returns the NPY shape string for the given shape's extents.
// - Rank 0 (scalar) --> "()".
// - Rank 1 (vector) --> "(len,)". Trailing comma is intentional.
// - Rank >= 2 --> "(rows,cols,...)".
std::string NpyShapeString(const std::vector<size_t>& shape);

// Encodes the length of `header` in NPY format (four bytes, little endian).
std::string NpyHeaderLengthString(absl::string_view header);

// Returns the "full" NPY file header for the given DataType and ShapeType.
// The "full header" consists of the magic 6 bytes, version number, and the
// "NPY header" that describes the data.
//
// This returns a header for version 2.0.
template <typename DataType, typename ShapeType>
std::string NpyFullHeaderString(ShapeType shape, bool reverse_axes) {
  constexpr absl::string_view kMagic("\x93NUMPY");

  // The explicit count is required since the \x00 in the literal would be
  // interpreted to construct a string of length 1.
  constexpr absl::string_view kVersion("\x02\x00", /*count=*/2);

  // The NPY format says that:
  // - If fortran_order = False (the default NPY ordering):
  //   Then the data is stored with the innermost axis changing most frequently.
  // - If fortran_order = True:
  //   Then the data is stored with the outermost axis changing most frequently.
  // Since `NpyDataString` always serializes data with the innermost *nda* axis
  // changing most frequently, if we *do* reverse the axes so that in npy, the
  // last axis is the innermost, then fortran_order is False.
  const bool fortran_order = !reverse_axes;

  std::vector<size_t> shape_vector = NpyShapeVector(shape);
  if (reverse_axes) {
    std::reverse(shape_vector.begin(), shape_vector.end());
  }

  const std::string header =
      absl::StrCat("{'descr': ", "'", NpyDescrString<DataType>(), "'", ", ",
                   "'fortran_order': ", fortran_order ? "True, " : "False, ",
                   "'shape': ", NpyShapeString(shape_vector), "}");

  return absl::StrCat(kMagic, kVersion, NpyHeaderLengthString(header), header);
}

// Returns a copy of `src` serialized compactly. Pedentically:
// - The output buffer has the exact size needed to represent all elements of
//   `src`, no more, no less.
// - The output buffer stores data with the innermost axis of `src` changing
//   most frequently.
// - In Numpy parlance, this means the data is C-contiguous: the innermost
//   axis (the first axis in nda convention, the last axis in numpy convention)
//   changes most frequently.
template <typename DataType, typename ShapeType>
std::string NpyDataString(nda::array_ref<const DataType, ShapeType> src) {
  // Allocate a buffer with exactly the amount of space needed to compactly
  // store `src`.
  const size_t dst_buffer_size_bytes = src.size() * sizeof(DataType);
  std::string dst_buffer(dst_buffer_size_bytes, '\0');

  // Wrap `dst_buffer` in an nda::array_ref with the same shape as `src` but
  // compact. nda::make_compact won't remove known-at-compile-time padding, only
  // dynamic padding. So we first promote it to a shape_of_rank<R>.
  DataType* dst_ptr = reinterpret_cast<DataType*>(dst_buffer.data());
  const nda::shape_of_rank<ShapeType::rank()> dynamic_shape = src.shape();
  auto dst_ref = nda::make_array_ref(dst_ptr, nda::make_compact(dynamic_shape));

  // Now just do a copy.
  nda::copy(src, dst_ref);

  return dst_buffer;
}

template <typename DataType, typename ShapeType>
std::string SerializeToNpyString(nda::array_ref<const DataType, ShapeType> src,
                                 const NpySerializeOptions& options) {
  // TODO(jiawen): Write a test for this case.
  if (src.empty()) {
    return "";
  }

  return absl::StrCat(internal::NpyFullHeaderString<DataType>(
                          src.shape(), options.reverse_axes),
                      internal::NpyDataString(src));
}

struct NpyHeader {
  // Array shape and total element count derived from it. total_element_count is
  // product of all sizes if it's a non-empty vector (representing an array of
  // rank > 0), or 1 if it's empty.
  std::vector<size_t> shape;
  size_t total_element_count = 0;

  // A single character describing interpretation of the type, e.g. 'f' for
  // float etc.
  char type_char = 'x';

  // Word size, with interpretation forming full used type. e.g. 'f' and 16
  // means half float, 32 full float etc.
  size_t word_size = 0;

  // See NpyFullHeaderString above.
  bool fortran_order = false;

  // Offset in the byte stream that describes where the actual array data
  // starts.
  size_t data_start_offset = 0;

  // If this is false, parsing of the header has failed.
  bool valid = false;
};

NpyHeader ReadHeader(absl::string_view src);

template <class Shape, size_t... Is>
Shape ToShapeImpl(const std::vector<size_t>& sizes,
                  std::index_sequence<Is...>) {
  return Shape({sizes[Is]...});
}

template <class Shape>
Shape ToShape(const std::vector<size_t>& sizes) {
  return ToShapeImpl<Shape>(sizes, std::make_index_sequence<Shape::rank()>());
}

}  // namespace internal

template <typename DataType, typename ShapeType>
std::string SerializeToNpyString(nda::array_ref<DataType, ShapeType> src,
                                 const NpySerializeOptions& options) {
  // Ensure that we only work on const (read-only) arrays.
  return internal::SerializeToNpyString(src.cref(), options);
}

template <typename DataType, typename ShapeType>
nda::array<DataType, ShapeType> DeserializeFromNpyString(
    absl::string_view src) {
  if (src.empty()) {
    LOG(ERROR) << "DeserializeFromNpyString: unable to deserialize, got an "
                  "empty string.";
    return nda::array<DataType, ShapeType>();
  }

  internal::NpyHeader header = internal::ReadHeader(src);
  if (!header.valid) {
    LOG(ERROR) << "DeserializeFromNpyString: unable to deserialize, got an "
                  "invalid npy header.";
    return nda::array<DataType, ShapeType>();
  }

  // After parsing the header, perform expected data verification.
  const size_t expected_data_size =
      header.total_element_count * header.word_size;
  if (header.data_start_offset + expected_data_size != src.size()) {
    LOG(ERROR) << "DeserializeFromNpyString: unable to deserialize, expected "
                  "string of length "
               << header.data_start_offset + expected_data_size << ", got "
               << src.size() << ".";
    return nda::array<DataType, ShapeType>();
  }
  if (internal::NpyDataTypeString<DataType>()[0] != header.type_char ||
      sizeof(DataType) != header.word_size) {
    LOG(ERROR) << "DeserializeFromNpyString: unable to deserialize, npy "
                  "contains data type "
               << header.type_char << header.word_size << ", while requested "
               << internal::NpyDataTypeString<DataType>()[0]
               << sizeof(DataType) << ".";
    return nda::array<DataType, ShapeType>();
  }

  if (ShapeType::rank() != header.shape.size()) {
    LOG(ERROR) << "DeserializeFromNpyString: unable to deserialize, got rank "
               << header.shape.size() << ", while requested "
               << ShapeType::rank() << ".";
    return nda::array<DataType, ShapeType>();
  }

  // See the rationale for flipping in NpySerializeOptions.
  if (!header.fortran_order) {
    std::reverse(header.shape.begin(), header.shape.end());
  }

  nda::array<DataType, ShapeType> array = nda::array<DataType, ShapeType>(
      internal::ToShape<ShapeType>(header.shape));
  src.copy(reinterpret_cast<char*>(array.data()), expected_data_size,
           header.data_start_offset);
  return array;
}

}  // namespace npy_array

#endif  // NPY_ARRAY_NPY_ARRAY_NPY_ARRAY_H_


<|start_filename|>scann/scann/projection/random_orthogonal_projection.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_PROJECTION_RANDOM_ORTHOGONAL_PROJECTION_H_
#define SCANN_PROJECTION_RANDOM_ORTHOGONAL_PROJECTION_H_

#include <cstdint>

#include "scann/data_format/datapoint.h"
#include "scann/oss_wrappers/scann_random.h"
#include "scann/projection/projection_base.h"
#include "scann/utils/types.h"

namespace research_scann {

template <typename T>
class RandomOrthogonalProjection : public Projection<T> {
 public:
  RandomOrthogonalProjection(const int32_t input_dims,
                             const int32_t projected_dims, const int32_t seed);

  void Create();

  StatusOr<shared_ptr<const TypedDataset<float>>> GetDirections() const final {
    return std::dynamic_pointer_cast<const TypedDataset<float>>(
        random_rotation_matrix_);
  }

  Status ProjectInput(const DatapointPtr<T>& input,
                      Datapoint<float>* projected) const override;
  Status ProjectInput(const DatapointPtr<T>& input,
                      Datapoint<double>* projected) const override;

  int32_t projected_dimensionality() const override { return projected_dims_; }

 private:
  template <typename FloatT>
  Status ProjectInputImpl(const DatapointPtr<T>& input,
                          Datapoint<FloatT>* projected) const;

  int32_t input_dims_;
  int32_t projected_dims_;
  unique_ptr<MTRandom> random_;
  int32_t seed_;
  shared_ptr<const DenseDataset<float>> random_rotation_matrix_;
};

SCANN_INSTANTIATE_TYPED_CLASS(extern, RandomOrthogonalProjection);

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/utils/intrinsics/flags.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_UTILS_INTRINSICS_FLAGS_H_
#define SCANN_UTILS_INTRINSICS_FLAGS_H_

#include "absl/flags/declare.h"
#include "absl/flags/flag.h"
#include "scann/utils/types.h"

ABSL_DECLARE_FLAG(bool, ignore_avx512);

ABSL_DECLARE_FLAG(bool, ignore_avx2);

ABSL_DECLARE_FLAG(bool, ignore_avx);

namespace research_scann {
namespace flags_internal {

extern bool should_use_avx1;
extern bool should_use_avx2;
extern bool should_use_avx512;
extern bool should_use_sse4;

}  // namespace flags_internal

inline bool RuntimeSupportsSse4() { return flags_internal::should_use_sse4; }
inline bool RuntimeSupportsAvx1() { return flags_internal::should_use_avx1; }
inline bool RuntimeSupportsAvx2() { return flags_internal::should_use_avx2; }
inline bool RuntimeSupportsAvx512() {
  return flags_internal::should_use_avx512;
}

enum PlatformGeneration {
  kFallbackForNonX86 = 99,

  kBaselineSse4 = 0,

  kSandyBridgeAvx1 = 1,

  kHaswellAvx2 = 2,

  kSkylakeAvx512 = 3,
};

inline string_view PlatformName(PlatformGeneration x86_arch) {
  switch (x86_arch) {
    case kFallbackForNonX86:
      return "FallbackForNonX86";
    case kBaselineSse4:
      return "SSE4";
    case kSandyBridgeAvx1:
      return "AVX1";
    case kHaswellAvx2:
      return "AVX2";
    case kSkylakeAvx512:
      return "AVX512";
    default:
      return "INVALID_X86_ARCH";
  }
}

class ScopedPlatformOverride {
 public:
  SCANN_DECLARE_MOVE_ONLY_CLASS(ScopedPlatformOverride);

  explicit ScopedPlatformOverride(PlatformGeneration generation);

  ~ScopedPlatformOverride();

  bool IsSupported();

 private:
  bool original_avx1_;
  bool original_avx2_;
  bool original_avx512_;
  bool original_sse4_;
};

ScopedPlatformOverride TestHookOverridePlatform(PlatformGeneration generation);

}  // namespace research_scann

#endif


<|start_filename|>fairness_and_bias_in_online_selection/unfair-secretary.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "unfair-secretary.h"

#include <algorithm>
#include <cmath>
#include <iostream>

namespace fair_secretary {

using std::vector;

SecretaryInstance UnfairSecretaryAlgorithm::ComputeSolution(
    const vector<SecretaryInstance>& elements) {
  double max_value = 0;
  int th = elements.size() * (1.0 / std::exp(1.0));
  // Compute max values.
  for (int i = 0; i < th; i++) {
    max_value = std::max(max_value, elements[i].value);
  }
  for (int i = th; i < elements.size(); i++) {
    if (elements[i].value >= max_value) {
      return elements[i];
    }
  }
  return SecretaryInstance(-1, -1);
}

SecretaryInstance UnfairSecretaryAlgorithm::ComputeSolutionSingleColor(
    const vector<SecretaryInstance>& elements, const vector<double>& prob) {
  double max_value = 0;
  int rand_color = 0;
  int rand = RandomHandler::eng_() % 1000000;
  double rand_balanced = static_cast<double>(rand) / 1000000;
  for (int i = 0; i < prob.size(); i++) {
    if (rand_balanced <= prob[i]) {
      rand_color = i;
      break;
    }
    rand_balanced -= prob[i];
  }
  int th = elements.size() * (1.0 / std::exp(1.0));
  // Compute max values.
  for (int i = 0; i < th; i++) {
    if (rand_color == elements[i].color) {
      max_value = std::max(max_value, elements[i].value);
    }
  }
  for (int i = th; i < elements.size(); i++) {
    if (elements[i].value >= max_value && elements[i].color == rand_color) {
      return elements[i];
    }
  }
  return SecretaryInstance(-1, -1);
}

}  // namespace fair_secretary


<|start_filename|>fairness_and_bias_in_online_selection/fair-prophet.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef FAIRNESS_AND_BIAS_SELECTION_FAIR_PROPHET_H_
#define FAIRNESS_AND_BIAS_SELECTION_FAIR_PROPHET_H_

#include <vector>

#include "distributions.h"
#include "utils.h"

namespace fair_secretary {

// Fair Prophet algorithm introduced in this work. It support both IID and
// non-IID algorithms.
class FairProphetAlgorithm {
 public:
  SecretaryInstance ComputeSolution(
      const std::vector<SecretaryInstance>& elements,
      const std::vector<std::reference_wrapper<RandomDistribution>>
          distributions,
      const std::vector<double>& q);
  SecretaryInstance ComputeSolutionIID(
      const std::vector<SecretaryInstance>& elements,
      const std::vector<std::reference_wrapper<RandomDistribution>>
          distributions,
      const std::vector<double>& q);
};

}  // namespace fair_secretary

#endif  // FAIRNESS_AND_BIAS_SELECTION_FAIR_PROPHET_H_


<|start_filename|>scann/scann/utils/intrinsics/fallback.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_UTILS_INTRINSICS_FALLBACK_H_
#define SCANN_UTILS_INTRINSICS_FALLBACK_H_

#include <cstdint>

#include "scann/utils/index_sequence.h"
#include "scann/utils/intrinsics/attributes.h"
#include "scann/utils/intrinsics/flags.h"
#include "scann/utils/types.h"

namespace research_scann {
namespace fallback {

static constexpr PlatformGeneration kPlatformGeneration = kFallbackForNonX86;

template <typename T = void, size_t kTensorNumElements0 = 1,
          size_t... kTensorNumElements>
class Simd;

template <typename T, size_t kTensorNumElements0, size_t... kTensorNumElements>
using SimdFor = Simd<T, kTensorNumElements0, kTensorNumElements...>;

struct Zeros {};
struct Uninitialized {};

SCANN_INLINE string_view SimdName() { return "<Non-x86 Fallback>"; }
SCANN_INLINE bool RuntimeSupportsSimd() { return true; }

template <typename T, size_t kNumElementsRequired>
SCANN_INLINE constexpr size_t InferNumRegisters() {
  return kNumElementsRequired;
}

template <typename T>
struct FakeIntelType {
  T val;
  SCANN_INLINE FakeIntelType() {}
  SCANN_INLINE explicit FakeIntelType(T val) : val(val) {}
  SCANN_INLINE T operator[](size_t idx) const {
    DCHECK_EQ(idx, 0);
    return val;
  }
};

template <typename T, size_t kNumElementsArg>
class Simd<T, kNumElementsArg> {
 public:
  SCANN_DECLARE_COPYABLE_CLASS(Simd);
  static_assert(IsSameAny<T, float, double, int8_t, int16_t, int32_t, int64_t,
                          uint8_t, uint16_t, uint32_t, uint64_t>());

  static constexpr size_t kRegisterBits = sizeof(T) * 8;
  static constexpr size_t kRegisterBytes = sizeof(T);
  static constexpr size_t kNumRegisters = kNumElementsArg;
  static constexpr size_t kNumElements = kNumElementsArg;
  static constexpr size_t kElementsPerRegister = 1;

  using IntelType = FakeIntelType<T>;
  static_assert(sizeof(IntelType) == kRegisterBytes);

  Simd(Uninitialized) {}
  Simd() : Simd(Uninitialized()) {}

  SCANN_INLINE Simd(Zeros) { Clear(); }

  SCANN_INLINE Simd(IntelType val) {
    static_assert(kNumElements == 1);
    *this = val;
  }

  SCANN_INLINE Simd(T val) { *this = Broadcast(val); }

  template <typename U, size_t kOther>
  SCANN_INLINE explicit Simd(const Simd<U, kOther>& other) {
    Simd& me = *this;
    for (size_t j : Seq(kNumElements)) {
      if constexpr (kOther == kNumElements) {
        me[j] = *other[j];
      } else if constexpr (kOther == 1) {
        me[j] = *other[0];
      } else {
        static_assert(kOther == kNumElements || kOther == 1);
      }
    }
  }

  SCANN_INLINE Simd& operator=(Zeros val) {
    Clear();
    return *this;
  }

  SCANN_INLINE Simd& operator=(IntelType val) {
    static_assert(kNumElements == 1,
                  "To intentionally perform register-wise broadcast, "
                  "explicitly cast to an Simd<T>");
    registers_[0] = val;
    return *this;
  }

  SCANN_INLINE Simd& operator=(T val) {
    *this = Broadcast(val);
    return *this;
  }

  SCANN_INLINE IntelType operator*() const {
    static_assert(kNumElements == 1);
    return registers_[0];
  }

  SCANN_INLINE Simd<T, 1>& operator[](size_t idx) {
    if constexpr (kNumElements == 1) {
      DCHECK_EQ(idx, 0);
      return *this;
    } else {
      DCHECK_LT(idx, kNumElements);
      return registers_[idx];
    }
  }

  SCANN_INLINE const Simd<T, 1>& operator[](size_t idx) const {
    if constexpr (kNumElements == 1) {
      DCHECK_EQ(idx, 0);
      return *this;
    } else {
      DCHECK_LT(idx, kNumElements);
      return registers_[idx];
    }
  }

  static SCANN_INLINE Simd Zeros() {
    Simd<T, kNumElements> ret;
    for (size_t j : Seq(kNumElements)) {
      ret[j] = IntelType(0);
    }
    return ret;
  }

  SCANN_INLINE Simd& Clear() {
    for (size_t j : Seq(kNumElements)) {
      registers_[j] = IntelType(0);
    }
    return *this;
  }

  SCANN_INLINE static Simd Broadcast(T x) {
    Simd ret;
    for (size_t j : Seq(kNumElements)) {
      ret[j] = IntelType(x);
    }
    return ret;
  }

  template <bool kAligned = false>
  SCANN_INLINE static Simd Load(const T* address) {
    Simd ret;
    for (size_t j : Seq(kNumElements)) {
      ret[j] = address[j];
    }
    return ret;
  }

  SCANN_INLINE void Store(T* address) const {
    const Simd& me = *this;
    for (size_t j : Seq(kNumElements)) {
      address[j] = me[j].Unwrap();
    }
  }

  SCANN_INLINE array<T, kNumElements> Store() const {
    array<T, kNumElements> ret;
    Store(ret.data());
    return ret;
  }

  template <size_t kOther, typename Op,
            size_t kOutput = std::max(kNumElements, kOther)>
  static SCANN_INLINE Simd<T, kOutput> BinaryOperatorImpl(
      const Simd& me, const Simd<T, kOther>& other, Op fn) {
    Simd<T, kOutput> ret;
    for (size_t j : Seq(Simd<T, kOutput>::kNumElements)) {
      if constexpr (kOther == kNumElements) {
        ret[j] = fn(*me[j], *other[j]);
      } else if constexpr (kNumElements == 1) {
        ret[j] = fn(*me[0], *other[j]);
      } else if constexpr (kOther == 1) {
        ret[j] = fn(*me[j], *other[0]);
      } else {
        static_assert(kOther == kNumElements || kNumElements == 1 ||
                      kOther == 1);
      }
    }
    return ret;
  }

  static SCANN_INLINE IntelType Add(IntelType a, IntelType b) {
    return IntelType(a.val + b.val);
  }

  template <size_t kOther>
  SCANN_INLINE auto operator+(const Simd<T, kOther>& other) const {
    return BinaryOperatorImpl(*this, other, &Add);
  }

  static SCANN_INLINE IntelType Subtract(IntelType a, IntelType b) {
    return IntelType(a.val - b.val);
  }

  template <size_t kOther>
  SCANN_INLINE auto operator-(const Simd<T, kOther>& other) const {
    return BinaryOperatorImpl(*this, other, &Subtract);
  }

  static SCANN_INLINE IntelType Multiply(IntelType a, IntelType b) {
    return IntelType(a.val * b.val);
  }

  template <size_t kOther>
  SCANN_INLINE auto operator*(const Simd<T, kOther>& other) const {
    return BinaryOperatorImpl(*this, other, &Multiply);
  }

  static SCANN_INLINE IntelType Divide(IntelType a, IntelType b) {
    return IntelType(a.val / b.val);
  }

  template <size_t kOther>
  SCANN_INLINE auto operator/(const Simd<T, kOther>& other) const {
    return BinaryOperatorImpl(*this, other, &Divide);
  }

  static SCANN_INLINE IntelType BitwiseAnd(IntelType a, IntelType b) {
    if constexpr (IsSame<T, float>()) {
      return IntelType(a.val & b.val);
    }
    if constexpr (IsSame<T, double>()) {
      return IntelType(a.val & b.val);
    }
    if constexpr (IsSameAny<T, int8_t, uint8_t, int16_t, uint16_t, int32_t,
                            uint32_t, int64_t, uint64_t>()) {
      return IntelType(a.val & b.val);
    }
  }

  template <size_t kOther>
  SCANN_INLINE auto operator&(const Simd<T, kOther>& other) const {
    return BinaryOperatorImpl(*this, other, &BitwiseAnd);
  }

  static SCANN_INLINE IntelType BitwiseOr(IntelType a, IntelType b) {
    if constexpr (IsSame<T, float>()) {
      return IntelType(a.val | b.val);
    }
    if constexpr (IsSame<T, double>()) {
      return IntelType(a.val | b.val);
    }
    if constexpr (IsSameAny<T, int8_t, uint8_t, int16_t, uint16_t, int32_t,
                            uint32_t, int64_t, uint64_t>()) {
      return IntelType(a.val | b.val);
    }
  }

  template <size_t kOther>
  SCANN_INLINE auto operator|(const Simd<T, kOther>& other) const {
    return BinaryOperatorImpl(*this, other, &BitwiseOr);
  }

  static SCANN_INLINE IntelType ShiftRight(IntelType x, int count) {
    static_assert(!IsSameAny<T, float, double>(),
                  "Bit shifting isn't defined for floating-point types.");
    return IntelType(x.val << count);
  }

  SCANN_INLINE Simd operator>>(int count) const {
    const Simd& me = *this;
    Simd ret;
    for (size_t j : Seq(kNumElements)) {
      ret[j] = ShiftRight(*me[j], count);
    }
    return ret;
  }

  static SCANN_INLINE IntelType ShiftLeft(IntelType x, int count) {
    static_assert(!IsSameAny<T, float, double>(),
                  "Bit shifting isn't defined for floating-point types.");
    return IntelType(x.val << count);
  }

  SCANN_INLINE Simd operator<<(int count) const {
    const Simd& me = *this;
    Simd ret;
    for (size_t j : Seq(kNumElements)) {
      ret[j] = ShiftLeft(*me[j], count);
    }
    return ret;
  }

  template <size_t kOther, typename Op>
  SCANN_INLINE Simd& AccumulateOperatorImpl(const Simd<T, kOther>& other,
                                            Op fn) {
    Simd& me = *this;
    for (size_t j : Seq(kNumElements)) {
      if constexpr (kOther == kNumElements) {
        me[j] = fn(*me[j], *other[j]);
      } else if constexpr (kOther == 1) {
        me[j] = fn(*me[j], *other[0]);
      } else {
        static_assert(kOther == kNumElements || kOther == 1);
      }
    }
    return *this;
  }

  template <size_t kOther>
  SCANN_INLINE Simd& operator+=(const Simd<T, kOther>& other) {
    return AccumulateOperatorImpl(other, &Add);
  }

  template <size_t kOther>
  SCANN_INLINE Simd& operator-=(const Simd<T, kOther>& other) {
    return AccumulateOperatorImpl(other, &Subtract);
  }

  template <size_t kOther>
  SCANN_INLINE Simd& operator*=(const Simd<T, kOther>& other) {
    return AccumulateOperatorImpl(other, &Multiply);
  }

  template <size_t kOther>
  SCANN_INLINE Simd& operator/=(const Simd<T, kOther>& other) {
    return AccumulateOperatorImpl(other, &Divide);
  }

  template <size_t kOther>
  SCANN_INLINE Simd& operator&=(const Simd<T, kOther>& other) {
    return AccumulateOperatorImpl(other, &BitwiseAnd);
  }

  template <size_t kOther>
  SCANN_INLINE Simd& operator|=(const Simd<T, kOther>& other) {
    return AccumulateOperatorImpl(other, &BitwiseOr);
  }

  SCANN_INLINE Simd& operator<<=(int count) {
    Simd& me = *this;
    for (size_t j : Seq(kNumElements)) {
      me[j] = ShiftLeft(*me[j], count);
    }
    return *this;
  }

  SCANN_INLINE Simd& operator>>=(int count) {
    Simd& me = *this;
    for (size_t j : Seq(kNumElements)) {
      me[j] = ShiftRight(*me[j], count);
    }
    return *this;
  }

  template <size_t kOther = kNumElements, typename Op>
  SCANN_INLINE auto ComparisonOperatorImpl(const Simd& me,
                                           const Simd<T, kOther>& other,
                                           Op fn) const {
    array<bool, std::max(kNumElements, kOther)> masks;
    for (size_t j : Seq(std::max(kNumElements, kOther))) {
      if constexpr (kOther == kNumElements) {
        masks[j] = fn(*me[j], *other[j]);
      } else if constexpr (kNumElements == 1) {
        masks[j] = fn(*me[0], *other[j]);
      } else if constexpr (kOther == 1) {
        masks[j] = fn(*me[j], *other[0]);
      } else {
        static_assert(kOther == kNumElements || kNumElements == 1 ||
                      kOther == 1);
      }
    }
    return masks;
  }

  static SCANN_INLINE bool LessThan(IntelType a, IntelType b) {
    LOG(INFO) << StrFormat("Called: %s < %s", StrCat(a.val), StrCat(b.val));
    return a.val < b.val;
  }

  template <size_t kOther = kNumElements>
  SCANN_INLINE auto operator<(const Simd<T, kOther>& other) const {
    return ComparisonOperatorImpl(*this, other, &LessThan);
  }

  static SCANN_INLINE bool LessOrEquals(IntelType a, IntelType b) {
    return a.val <= b.val;
  }

  template <size_t kOther = kNumElements>
  SCANN_INLINE auto operator<=(const Simd<T, kOther>& other) const {
    return ComparisonOperatorImpl(*this, other, &LessOrEquals);
  }

  static SCANN_INLINE bool Equals(IntelType a, IntelType b) {
    return a.val == b.val;
  }

  template <size_t kOther>
  SCANN_INLINE auto operator==(const Simd<T, kOther>& other) const {
    return ComparisonOperatorImpl(*this, other, &Equals);
  }

  static SCANN_INLINE bool GreaterOrEquals(IntelType a, IntelType b) {
    return a.val >= b.val;
  }

  template <size_t kOther = kNumElements>
  SCANN_INLINE auto operator>=(const Simd<T, kOther>& other) const {
    return ComparisonOperatorImpl(*this, other, &GreaterOrEquals);
  }

  static SCANN_INLINE bool GreaterThan(IntelType a, IntelType b) {
    return a.val > b.val;
  }

  template <size_t kOther = kNumElements>
  SCANN_INLINE auto operator>(const Simd<T, kOther>& other) const {
    return ComparisonOperatorImpl(*this, other, &GreaterThan);
  }

  SCANN_INLINE int MaskFromHighBits() const {
    static_assert(0 != 1, "Not Implemented. TODO: implement if/when needed.");
  }

  SCANN_INLINE T GetLowElement() const {
    const Simd& me = *this;
    return (*me[0]).val;
  }

  SCANN_INLINE T Unwrap() const { return GetLowElement(); }

  template <typename U>
  SCANN_INLINE Simd<U, kNumElements> ConvertTo() const {
    static_assert(IsDivisibleBy(kNumElements, 2));
    const Simd& me = *this;
    Simd<U, kNumElements> ret;
    for (size_t j : Seq(kNumElements)) {
      ret[j] = static_cast<U>((*me[j]).val);
    }
    return ret;
  }

  static SCANN_INLINE auto InferExpansionType() {
    if constexpr (IsSame<T, float>()) {
      return double();
    }
    if constexpr (IsSame<T, double>()) {
      return double();
    }

    if constexpr (IsSame<T, int8_t>()) {
      return int16_t();
    }
    if constexpr (IsSame<T, int16_t>()) {
      return int32_t();
    }
    if constexpr (IsSame<T, int32_t>()) {
      return int64_t();
    }
    if constexpr (IsSame<T, int64_t>()) {
      return int64_t();
    }

    if constexpr (IsSame<T, uint8_t>()) {
      return uint16_t();
    }
    if constexpr (IsSame<T, uint16_t>()) {
      return uint32_t();
    }
    if constexpr (IsSame<T, uint32_t>()) {
      return uint64_t();
    }
    if constexpr (IsSame<T, uint64_t>()) {
      return uint64_t();
    }
  }
  using ExpansionType = decltype(InferExpansionType());
  using ExpansionIntelType = typename Simd<ExpansionType>::IntelType;
  using ExpandsTo = Simd<ExpansionType, kNumElements>;

  template <typename ValidateT>
  SCANN_INLINE Simd<ExpansionType, kNumElements> ExpandTo() const {
    static_assert(IsSame<ValidateT, ExpansionType>());
    const Simd& me = *this;
    Simd<ExpansionType, kNumElements> ret;
    for (size_t j : Seq(kNumElements)) {
      ret[j] = static_cast<ExpansionType>(me[j].Unwrap());
    }
    return ret;
  }

 private:
  std::conditional_t<kNumElements == 1, IntelType, Simd<T, 1>>
      registers_[kNumElements];

  template <typename U, size_t kOther, size_t... kTensorOther>
  friend class Simd;
};

template <typename T, size_t kTensorNumRegisters0,
          size_t... kTensorNumRegisters>
class Simd {
 public:
  using SimdSubArray = Simd<T, kTensorNumRegisters...>;

  Simd(Uninitialized) {}
  Simd() : Simd(Uninitialized()) {}

  SCANN_INLINE Simd(Zeros) {
    for (size_t j : Seq(kTensorNumRegisters0)) {
      tensor_[j] = Zeros();
    }
  }

  SCANN_INLINE SimdSubArray& operator[](size_t idx) {
    DCHECK_LT(idx, kTensorNumRegisters0);
    return tensor_[idx];
  }

  SCANN_INLINE const SimdSubArray& operator[](size_t idx) const {
    DCHECK_LT(idx, kTensorNumRegisters0);
    return tensor_[idx];
  }

  SCANN_INLINE void Store(T* address) const {
    constexpr size_t kStride =
        sizeof(decltype(SimdSubArray().Store())) / sizeof(T);
    for (size_t j : Seq(kTensorNumRegisters0)) {
      tensor_[j].Store(address + j * kStride);
    }
  }

  using StoreResultType =
      array<decltype(SimdSubArray().Store()), kTensorNumRegisters0>;
  SCANN_INLINE StoreResultType Store() const {
    StoreResultType ret;
    for (size_t j : Seq(kTensorNumRegisters0)) {
      ret[j] = tensor_[j].Store();
    }
    return ret;
  }

 private:
  SimdSubArray tensor_[kTensorNumRegisters0];
};

}  // namespace fallback

SCANN_INLINE uint32_t GetComparisonMask(array<bool, 32> cmp) {
  uint32_t ret = 0;
  for (size_t j : ReverseSeq(cmp.size())) {
    ret <<= 1;
    ret += (cmp[j] ? 1 : 0);
  }
  return ret;
}

SCANN_INLINE uint64_t GetComparisonMask(array<bool, 64> cmp) {
  uint64_t ret = 0;
  for (size_t j : ReverseSeq(cmp.size())) {
    ret <<= 1;
    ret |= (cmp[j] ? 1 : 0);
  }
  return ret;
}

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/utils/scann_config_utils.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/utils/scann_config_utils.h"

#include <cstdint>
#include <string>

#include "absl/strings/match.h"
#include "absl/strings/str_format.h"
#include "absl/strings/substitute.h"
#include "scann/data_format/features.pb.h"
#include "scann/distance_measures/distance_measure_base.h"
#include "scann/distance_measures/distance_measure_factory.h"
#include "scann/partitioning/partitioner.pb.h"
#include "scann/proto/brute_force.pb.h"
#include "scann/proto/distance_measure.pb.h"
#include "scann/proto/exact_reordering.pb.h"
#include "scann/proto/hash.pb.h"
#include "scann/proto/incremental_updates.pb.h"
#include "scann/proto/input_output.pb.h"
#include "scann/proto/metadata.pb.h"
#include "scann/proto/partitioning.pb.h"
#include "scann/proto/projection.pb.h"
#include "scann/proto/restricts.pb.h"
#include "scann/proto/scann.pb.h"
#include "scann/utils/common.h"
#include "scann/utils/types.h"
#include "tensorflow/core/lib/core/errors.h"

using absl::StartsWith;

namespace research_scann {

namespace {

Status CanonicalizeDeprecatedFields(ScannConfig* config) {
  if (config->has_partitioning() &&
      config->partitioning().use_float_centers_for_query_tokenization()) {
    config->mutable_partitioning()->set_query_tokenization_type(
        PartitioningConfig::FLOAT);
    config->mutable_partitioning()
        ->clear_use_float_centers_for_query_tokenization();
  }
  if (config->hash().asymmetric_hash().has_min_number_machines() &&
      !config->hash().asymmetric_hash().has_num_machines()) {
    config->mutable_hash()->mutable_asymmetric_hash()->set_num_machines(
        config->hash().asymmetric_hash().min_number_machines());
  }
  if (config->has_partitioning() && config->partitioning().use_flume_kmeans()) {
    LOG(WARNING) << "use_flume_kmeans to be deprecated, use trainner_type to "
                    "specify FLUME_KMEANS_TRAINER instead.";
    config->mutable_partitioning()->set_trainer_type(
        PartitioningConfig::FLUME_KMEANS_TRAINER);
    config->mutable_partitioning()->clear_use_flume_kmeans();
  }

  if (config->has_brute_force() && config->brute_force().scalar_quantized()) {
    auto* bf = config->mutable_brute_force();
    bf->clear_scalar_quantized();
    bf->mutable_fixed_point()->set_enabled(true);
    if (bf->has_scalar_quantization_multiplier_quantile()) {
      bf->mutable_fixed_point()->set_fixed_point_multiplier_quantile(
          bf->scalar_quantization_multiplier_quantile());
    }
  }

  bool restrict_fields_populated = false;
  if (config->has_restricts()) {
    restrict_fields_populated = true;
    if (config->restricts().has_enabled()) {
      config->set_restricts_enabled(config->restricts().enabled());
    }
    if (config->restricts().has_v3_restricts()) {
      *config->mutable_v3_restricts() = config->restricts().v3_restricts();
    }
  }
  if (!restrict_fields_populated && config->restricts_enabled()) {
    auto* restricts = config->mutable_restricts();
    restricts->set_enabled(true);

    if (config->has_v3_restricts()) {
      *restricts->mutable_v3_restricts() = config->v3_restricts();
    }
  }
  return OkStatus();
}

Status CanonicalizeRecallCurves(ScannConfig* config) { return OkStatus(); }

Status CanonicalizeScannConfigImpl(ScannConfig* config,
                                   bool artifact_with_stable_base_name) {
  SCANN_RETURN_IF_ERROR(CanonicalizeDeprecatedFields(config));
  SCANN_RETURN_IF_ERROR(EnsureCorrectNormalizationForDistanceMeasure(config));
  auto io = config->mutable_input_output();

  const bool with_fingerprint_postfix =
      config->input_output().artifacts_naming_option() ==
      InputOutputConfig::HASHED;

  if (io->preprocessed_artifacts_dir().empty()) {
    return CanonicalizeRecallCurves(config);
  }
  const std::string& artifacts_dir = io->preprocessed_artifacts_dir();

  if (config->has_partitioning()) {
    const auto partitioning_config = config->partitioning();
    if (partitioning_config.has_database_spilling()) {
      const auto database_spilling_type =
          partitioning_config.database_spilling().spilling_type();
      const auto trainer_type = partitioning_config.trainer_type();

      if (trainer_type != PartitioningConfig::DEFAULT_SAMPLING_TRAINER &&
          (database_spilling_type == DatabaseSpillingConfig::ADDITIVE ||
           database_spilling_type == DatabaseSpillingConfig::MULTIPLICATIVE)) {
        return InvalidArgumentError(
            "ADDITIVE and MULTIPLICATIVE database spilling is only supported "
            "by the DEFAULT_SAMPLING_TRAINER");
      }
    }

    if (config->has_hash() && config->hash().has_asymmetric_hash()) {
      const auto& asymmetric_hash = config->hash().asymmetric_hash();
      if (asymmetric_hash.use_normalized_residual_quantization()) {
        if (!asymmetric_hash.use_residual_quantization()) {
          return InvalidArgumentError(
              "use_normalized_residual_quantization can only be used when "
              "use_residual_quantization is also turned on");
        }

        config->mutable_partitioning()->set_compute_residual_stdev(true);
      }
    }
  }

  if (!io->has_database_wildcard()) {
    return CanonicalizeRecallCurves(config);
  }

  return OkStatus();
}

}  // namespace

namespace {
Status CanonicalizeScannConfigForRetrieval(ScannConfig* config,
                                           bool artifact_with_stable_base_name,
                                           bool artifact_must_exist) {
  SCANN_RETURN_IF_ERROR(
      CanonicalizeScannConfigImpl(config, artifact_with_stable_base_name));

  auto io = config->mutable_input_output();
  if (io->preprocessed_artifacts_dir().empty()) return OkStatus();
  return OkStatus();
}
}  // namespace

Status CanonicalizeScannConfigForRetrieval(ScannConfig* config) {
  ScannConfig cfg2 = *config;
  SCANN_RETURN_IF_ERROR(
      CanonicalizeScannConfigForRetrieval(config, true, false));
  SCANN_RETURN_IF_ERROR(
      CanonicalizeScannConfigForRetrieval(&cfg2, false, false));
  if (config->partitioning().partitioner_prefix().empty() &&
      !cfg2.partitioning().partitioner_prefix().empty()) {
    config->mutable_partitioning()->set_partitioner_prefix(
        cfg2.partitioning().partitioner_prefix());
  }
  if (config->partitioning().resharded_prefix().empty() &&
      !cfg2.partitioning().resharded_prefix().empty()) {
    config->mutable_partitioning()->set_resharded_prefix(
        cfg2.partitioning().resharded_prefix());
  }
  if (config->hash().asymmetric_hash().centers_filename().empty() &&
      !cfg2.hash().asymmetric_hash().centers_filename().empty()) {
    config->mutable_hash()->mutable_asymmetric_hash()->set_centers_filename(
        cfg2.hash().asymmetric_hash().centers_filename());
  }
  const auto& io1 = config->input_output();
  const auto& io2 = cfg2.input_output();
  if (io1.hashed_database_wildcard().empty() &&
      !io2.hashed_database_wildcard().empty()) {
    config->mutable_input_output()->set_hashed_database_wildcard(
        io2.hashed_database_wildcard());
  }
  if (io1.fixed_point_database_wildcard().empty() &&
      !io2.fixed_point_database_wildcard().empty()) {
    config->mutable_input_output()->set_fixed_point_database_wildcard(
        io2.fixed_point_database_wildcard());
  }
  if (config->exact_reordering().fixed_point().multipliers_filename().empty() &&
      !cfg2.exact_reordering().fixed_point().multipliers_filename().empty()) {
    config->mutable_exact_reordering()
        ->mutable_fixed_point()
        ->set_multipliers_filename(
            cfg2.exact_reordering().fixed_point().multipliers_filename());
  }
  if (io1.tokenized_database_wildcard().empty() &&
      !io2.tokenized_database_wildcard().empty()) {
    config->mutable_input_output()->set_tokenized_database_wildcard(
        io2.tokenized_database_wildcard());
  }
  return OkStatus();
}

StatusOr<InputOutputConfig::InMemoryTypes> TagFromGFVFeatureType(
    const GenericFeatureVector::FeatureType& feature_type) {
  switch (feature_type) {
    case GenericFeatureVector::INT64:
      return InputOutputConfig::INT64;
    case GenericFeatureVector::FLOAT:
      return InputOutputConfig::FLOAT;
    case GenericFeatureVector::DOUBLE:
      return InputOutputConfig::DOUBLE;
    case GenericFeatureVector::BINARY:
      return InputOutputConfig::UINT8;
    default:
      return InvalidArgumentError("Invalid feature_type");
  }
}

StatusOr<InputOutputConfig::InMemoryTypes> DetectInMemoryTypeFromGfv(
    const GenericFeatureVector& gfv) {
  TF_ASSIGN_OR_RETURN(auto ret, TagFromGFVFeatureType(gfv.feature_type()));
  return ret;
}

StatusOr<InputOutputConfig::InMemoryTypes> DetectInMemoryTypeFromDisk(
    const ScannConfig& config, const size_t shard) {
  if (!config.has_input_output()) {
    return InvalidArgumentError("config must have input_output.");
  }

  const auto memory_data_type = config.input_output().in_memory_data_type();

  if (memory_data_type !=
      InputOutputConfig::IN_MEMORY_DATA_TYPE_NOT_SPECIFIED) {
    return memory_data_type;
  }

  if (!config.input_output().has_database_wildcard()) {
    return InvalidArgumentError(
        "config.input_output() must have database_wildcard if "
        "in_memory_data_type is not explicitly specified.");
  }
  return InvalidArgumentError("Input GFV from disk not supported.");
}

StatusOr<Normalization> NormalizationRequired(
    const std::string& distance_measure_name) {
  TF_ASSIGN_OR_RETURN(auto distance, GetDistanceMeasure(distance_measure_name));
  return distance->NormalizationRequired();
}

Status EnsureCorrectNormalizationForDistanceMeasure(ScannConfig* config) {
  std::string main_distance_measure;
  if (config->has_distance_measure()) {
    main_distance_measure = config->distance_measure().distance_measure();
  } else if (config->has_partitioning()) {
    if (config->partitioning().has_partitioning_distance()) {
      main_distance_measure =
          config->partitioning().partitioning_distance().distance_measure();
    } else {
      main_distance_measure = "SquaredL2Distance";
    }
  } else {
    return OkStatus();
  }
  TF_ASSIGN_OR_RETURN(const Normalization expected_normalization,
                      NormalizationRequired(main_distance_measure));
  const bool normalization_user_specified =
      config->input_output().has_norm_type();

  if (expected_normalization != NONE) {
    if (normalization_user_specified) {
      if (static_cast<Normalization>(config->input_output().norm_type()) !=
          expected_normalization) {
        return InvalidArgumentError(
            "Normalization required by the main distance measure %s (%s) "
            "does not match normalization specified in "
            "input_output.norm_type (%s).",
            main_distance_measure.c_str(),
            NormalizationString(expected_normalization),
            NormalizationString(static_cast<Normalization>(
                config->input_output().norm_type())));
      }
    } else {
      config->mutable_input_output()->set_norm_type(
          static_cast<InputOutputConfig::FeatureNorm>(expected_normalization));
    }

    StatusOr<InputOutputConfig::InMemoryTypes> in_memory_type_or_error =
        DetectInMemoryTypeFromDisk(*config);
    if (in_memory_type_or_error.ok()) {
      auto in_memory_type = in_memory_type_or_error.ValueOrDie();
      if (in_memory_type != InputOutputConfig::FLOAT &&
          in_memory_type != InputOutputConfig::DOUBLE) {
        LOG(WARNING) << "Performing "
                     << NormalizationString(expected_normalization)
                     << " normalization with an integral type.";
      }
    }
  }

  auto verify_consistency = [&](const std::string& secondary_distance_measure,
                                const std::string& param_name) -> Status {
    TF_ASSIGN_OR_RETURN(const Normalization secondary_expected,
                        NormalizationRequired(secondary_distance_measure));
    if (secondary_expected != expected_normalization &&
        !(normalization_user_specified && secondary_expected == NONE)) {
      return InvalidArgumentError(
          "Normalization required by main distance measure (%s) does not "
          "match normalization required by %s distance measure "
          "(%s).",
          NormalizationString(expected_normalization), param_name.c_str(),
          NormalizationString(secondary_expected));
    }

    return OkStatus();
  };

  if (config->partitioning().has_partitioning_distance()) {
    SCANN_RETURN_IF_ERROR(verify_consistency(
        config->partitioning().partitioning_distance().distance_measure(),
        "partitioning"));
  }

  if (config->partitioning().has_database_tokenization_distance_override()) {
    SCANN_RETURN_IF_ERROR(
        verify_consistency(config->partitioning()
                               .database_tokenization_distance_override()
                               .distance_measure(),
                           "database tokenization"));
  }

  if (config->partitioning().has_query_tokenization_distance_override()) {
    SCANN_RETURN_IF_ERROR(
        verify_consistency(config->partitioning()
                               .query_tokenization_distance_override()
                               .distance_measure(),
                           "query tokenization"));
  }

  if (config->exact_reordering().has_approx_distance_measure()) {
    SCANN_RETURN_IF_ERROR(verify_consistency(
        config->exact_reordering().approx_distance_measure().distance_measure(),
        "approximate"));
  }

  if (config->has_metadata() && config->metadata().metadata_type_case() ==
                                    MetadataConfig::kExactDistance) {
    SCANN_RETURN_IF_ERROR(verify_consistency(
        config->metadata().exact_distance().distance_measure(), "metadata"));
  }

  return OkStatus();
}

}  // namespace research_scann


<|start_filename|>zebraix/base/dominance.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef ZEBRAIX_BASE_DOMINANCE_H_
#define ZEBRAIX_BASE_DOMINANCE_H_

#include <vector>

namespace zebraix {
namespace base {

// These are indices that are impossible in protobuf input repeated messages.
static constexpr int kSinkImportPseudoIndex = -1;
static constexpr int kSourceImportPseudoIndex = -2;

// DominanceNode indices index the vector of nodes in a DominanceGraph.
struct DominanceNode {
  int prime_rank = -1;
  int obverse_rank = -1;
  // The import_index is used contextually. If the data comes from a proto, it
  // indexes the original proto source so that, for example, one can pull
  // rendering information.
  int import_index = -1;
  std::vector<int> parents;
  std::vector<int> children;
};

struct DominanceGraph {
  // Ranges of ranks that include any imputed source or sink.
  int prime_min;
  int prime_max;
  int obverse_min;
  int obverse_max;

  // Sources and sinks are only imputed if there is not already one.
  int source_index;
  bool source_imputed;
  int sink_index;
  bool sink_imputed;

  std::vector<DominanceNode> d_nodes;
};

inline bool RankCmp(const DominanceNode& i, const DominanceNode& j) {
  return i.prime_rank < j.prime_rank;
}

// Scan nodes for ranges of ranks. Ascertain if source and/or sink nodes need to
// be imputed.
void ScanRankRanges(DominanceGraph* d_graph);
// Create source and sink nodes as required. Sort nodes in increasing order of
// prime rank.
void FleshOutGraphNodes(DominanceGraph* d_graph);

// Apply order-dimension 2 properties to construct graph connections (edges).
void ConnectGraph(DominanceGraph* d_graph);

}  // namespace base
}  // namespace zebraix

#endif  // ZEBRAIX_BASE_DOMINANCE_H_


<|start_filename|>zebraix/render/zebraix_rendersvg.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "render/zebraix_rendersvg.h"

#include <math.h>

#include <iostream>
#include <string>
#include <vector>

#include "glog/logging.h"
#include "absl/container/flat_hash_map.h"
#include "third_party/cairo/cairo-svg.h"
#include "third_party/cairo/cairo.h"  // IWYU pragma: keep
// IWYU pragma: no_include "third_party/cairo/v1_12_16/src/cairo.h"
#include "third_party/fontconfig/fontconfig/fontconfig.h"
// IWYU pragma: no_include "third_party/glib/src/gobject/gobject.h"
#include "base/dominance.h"
#include "misc/misc_proto.h"
#include "misc/misc_render.h"
#include "render/render_structs.h"
#include "render/zebraix_layout.h"
#include "third_party/pango/pango/pangocairo.h"
#include "third_party/pango/pango/pangoft2.h"
// IWYU pragma: no_include "third_party/pango/src/pango/pango-font.h"
// IWYU pragma: no_include "third_party/pango/src/pango/pango-fontmap.h"
// IWYU pragma: no_include "third_party/pango/src/pango/pango-layout.h"
// IWYU pragma: no_include "third_party/pango/src/pango/pango-types.h"

using zebraix::misc::NormalizeVector;
using zebraix::misc::OctantToAnchor;
using zebraix::render::ZebraixCanvas;
using zebraix::render::ZebraixLayout;

namespace zebraix {
namespace render {

namespace {
// TODO(b/193166687): Move constants into base "default" config.
constexpr double kLineWidth = 1.5;
constexpr double kArrowSize = 12.0;
constexpr double kLabelFontSize = 12.0;
constexpr double kDashLength = 6.0;
constexpr double kArrowAngle = 25.0;
constexpr double kArrowOffsetFraction = 1.25;
constexpr double kDotLength = kLineWidth;
constexpr double kGridGray = 0.60;
static double kGridDashLength[2] = {4, 4};

void CairoRenderLabel(double x, double y, const ZebraixLabel& label,
                      cairo_t* cr) {
  CHECK(!label.label_text.empty());
  PangoLayout* layout;
  PangoContext* pango_context;

  PangoFontDescription* font_description;
  font_description = pango_font_description_new();
  pango_font_description_set_family(font_description, "sans");
  pango_font_description_set_absolute_size(font_description,
                                           kLabelFontSize * PANGO_SCALE);

  pango_context = pango_cairo_create_context(cr);
  layout = pango_layout_new(pango_context);
  pango_layout_set_font_description(layout, font_description);
  pango_layout_set_text(layout, label.label_text.c_str(), -1);

  cairo_set_source_rgb(cr, 0.0, 0.0, 1.0);

  double adjust_x = 0.0;
  double adjust_y = 0.0;

  PangoRectangle logical_extents;
  pango_layout_get_extents(layout, nullptr /*=ink_extents*/, &logical_extents);
  switch (label.anchor) {
    case zebraix_proto::TL:
      break;
    case zebraix_proto::T:
      adjust_x = 0.5 * logical_extents.width;
      break;
    case zebraix_proto::TR:
      adjust_x = 1.0 * logical_extents.width;
      break;
    case zebraix_proto::R:
      adjust_x = 1.0 * logical_extents.width;
      adjust_y = -0.5 * logical_extents.height;
      break;
    case zebraix_proto::BR:
      adjust_x = 1.0 * logical_extents.width;
      adjust_y = -1.0 * logical_extents.height;
      break;
    case zebraix_proto::B:
      adjust_x = 0.5 * logical_extents.width;
      adjust_y = -1.0 * logical_extents.height;
      break;
    case zebraix_proto::BL:
      adjust_y = -1.0 * logical_extents.height;
      break;
    case zebraix_proto::L:
      adjust_y = -0.5 * logical_extents.height;
      break;
    default:
      CHECK(false);
      break;
  }

  cairo_move_to(cr, x - adjust_x / PANGO_SCALE, -(y - adjust_y / PANGO_SCALE));
  pango_cairo_show_layout(cr, layout);

  pango_cairo_font_map_set_default(nullptr);
  g_object_unref(layout);
  g_object_unref(pango_context);
  pango_font_description_free(font_description);
}

void CairoRenderEdge(const ZebraixEdge& r_edge,
                     const std::vector<ZebraixNode>& nodes, cairo_t* cr) {
  const ZebraixNode& start_node = nodes[r_edge.start_node_index];
  const ZebraixNode& finish_node = nodes[r_edge.finish_node_index];
  const zebraix_proto::ShowHide display = r_edge.display;
  if (display == zebraix_proto::HIDE) {
    return;
  }
  if (display == zebraix_proto::GHOST) {
    cairo_set_dash(cr, &kDashLength, 1, kArrowOffsetFraction * kDashLength);
  }

  double delta_x = finish_node.centre_x - start_node.centre_x;
  double delta_y = finish_node.centre_y - start_node.centre_y;
  NormalizeVector(&delta_x, &delta_y);
  double forward_tip_x = finish_node.centre_x - delta_x * finish_node.radius;
  double forward_tip_y = finish_node.centre_y - delta_y * finish_node.radius;
  double reverse_tip_x = start_node.centre_x + delta_x * start_node.radius;
  double reverse_tip_y = start_node.centre_y + delta_y * start_node.radius;

  double arrow_1_x = -delta_x * kArrowSize;
  double arrow_1_y = -delta_y * kArrowSize;
  double arrow_2_x = arrow_1_x;
  double arrow_2_y = arrow_1_y;
  cairo_matrix_t arrow_rotate;

  cairo_matrix_init_rotate(&arrow_rotate, kArrowAngle / 180.0 * M_PI);
  cairo_matrix_transform_distance(&arrow_rotate, &arrow_1_x, &arrow_1_y);
  cairo_matrix_init_rotate(&arrow_rotate, -kArrowAngle / 180.0 * M_PI);
  cairo_matrix_transform_distance(&arrow_rotate, &arrow_2_x, &arrow_2_y);

  // Draw main line (of arrow) from destination to source, so that dashed lines
  // are drawn at arrow.
  cairo_set_source_rgb(cr, 0, 0, 0);
  cairo_set_line_width(cr, kLineWidth);

  cairo_move_to(cr, forward_tip_x, -forward_tip_y);
  cairo_line_to(cr, reverse_tip_x, -reverse_tip_y);
  cairo_stroke(cr);

  cairo_set_dash(cr, &kDashLength, 0, 0.0);

  if (r_edge.forward_arrow) {
    cairo_move_to(cr, forward_tip_x, -forward_tip_y);
    cairo_line_to(cr, forward_tip_x + arrow_1_x, -(forward_tip_y + arrow_1_y));
    cairo_stroke(cr);
    cairo_move_to(cr, forward_tip_x, -forward_tip_y);
    cairo_line_to(cr, forward_tip_x + arrow_2_x, -(forward_tip_y + arrow_2_y));
    cairo_stroke(cr);
  }
  if (r_edge.reverse_arrow) {
    cairo_move_to(cr, reverse_tip_x, -reverse_tip_y);
    cairo_line_to(cr, reverse_tip_x - arrow_1_x, -(reverse_tip_y - arrow_1_y));
    cairo_stroke(cr);
    cairo_move_to(cr, reverse_tip_x, -reverse_tip_y);
    cairo_line_to(cr, reverse_tip_x - arrow_2_x, -(reverse_tip_y - arrow_2_y));
    cairo_stroke(cr);
  }
}

void CairoRenderGrid(const ZebraixLayout& graph_layout, cairo_t* cr) {
  const double grid_thickness =
      graph_layout.GetMainLayoutProto().layout().grid_thickness();
  if (grid_thickness != 0.0) {
    cairo_set_source_rgb(cr, 1.0 - kGridGray, 1.0 - kGridGray, 1.0 - kGridGray);
    cairo_set_line_width(cr, grid_thickness);

    for (int x = graph_layout.GetDominance().prime_min;
         x <= graph_layout.GetDominance().prime_max; ++x) {
      if (x == 0) {
        cairo_set_dash(cr, kGridDashLength, 0, 0.0);
      } else {
        cairo_set_dash(cr, kGridDashLength, 2, 0.0);
      }

      double y_begin = graph_layout.GetDominance().obverse_min - 0.5;
      double y_end = graph_layout.GetDominance().obverse_max + 0.5;
      double x_begin = 1.0 * x;
      double x_end = x_begin;

      cairo_matrix_transform_distance(&graph_layout.GetCanvas().global_tran,
                                      &x_begin, &y_begin);
      cairo_matrix_transform_distance(&graph_layout.GetCanvas().global_tran,
                                      &x_end, &y_end);
      cairo_move_to(cr, x_begin, -y_begin);
      cairo_line_to(cr, x_end, -y_end);
      cairo_stroke(cr);
    }
    for (int y = graph_layout.GetDominance().obverse_min;
         y <= graph_layout.GetDominance().obverse_max; ++y) {
      if (y == 0) {
        cairo_set_dash(cr, kGridDashLength, 0, 0.0);
      } else {
        cairo_set_dash(cr, kGridDashLength, 2, 0.0);
      }

      double x_begin = graph_layout.GetDominance().prime_min - 0.5;
      double x_end = graph_layout.GetDominance().prime_max + 0.5;
      double y_begin = 1.0 * y;
      double y_end = y_begin;

      cairo_matrix_transform_distance(&graph_layout.GetCanvas().global_tran,
                                      &x_begin, &y_begin);
      cairo_matrix_transform_distance(&graph_layout.GetCanvas().global_tran,
                                      &x_end, &y_end);
      cairo_move_to(cr, x_begin, -y_begin);
      cairo_line_to(cr, x_end, -y_end);
      cairo_stroke(cr);
    }
    cairo_set_dash(cr, kGridDashLength, 0, 0.0);
  }

  cairo_set_line_width(cr, kLineWidth);
}

void CairoRenderNode(const ZebraixNode& r_node,
                     const ZebraixLayout& graph_layout, cairo_t* cr) {
  if (r_node.display == zebraix_proto::HIDE) {
    return;
  }
  if (r_node.display == zebraix_proto::GHOST) {
    cairo_set_dash(cr, &kDashLength, 1, 0.0);
  } else if (r_node.display == zebraix_proto::WAYPOINT) {
    cairo_set_dash(cr, &kDotLength, 1, 0.0);
  }

  cairo_set_source_rgb(cr, 0, 0, 0);

  if (r_node.radius != 0) {
    cairo_move_to(cr, r_node.centre_x + r_node.radius, -r_node.centre_y);
    cairo_arc(cr, r_node.centre_x, -r_node.centre_y, r_node.radius, 0.0 * M_PI,
              2.0 * M_PI);
    cairo_stroke(cr);
  }

  cairo_set_dash(cr, &kDashLength, 0, 0.0);

  if (!r_node.label.label_text.empty() &&
      (r_node.display != zebraix_proto::HIDE)) {
    CairoRenderLabel(r_node.centre_x + r_node.label_anchor_x,
                     r_node.centre_y + r_node.label_anchor_y, r_node.label, cr);
  }

  if (graph_layout.GetConfig().draw_label_ticks &&
      (r_node.display != zebraix_proto::HIDE)) {
    cairo_move_to(cr, r_node.centre_x, -r_node.centre_y);
    cairo_line_to(cr, r_node.centre_x + r_node.label_anchor_x,
                  -(r_node.centre_y + r_node.label_anchor_y));
    cairo_stroke(cr);
  }
}

void CairoRenderLayout(const ZebraixLayout& graph_layout, cairo_t* cr) {
  CairoRenderGrid(graph_layout, cr);
  cairo_set_line_width(cr, kLineWidth);

  for (const auto& r_node : graph_layout.GetNodes()) {
    CairoRenderNode(r_node, graph_layout, cr);
  }

  for (int p = 0; p < graph_layout.GetNodes().size(); ++p) {
    for (const auto c : graph_layout.GetDominance().d_nodes[p].children) {
      CairoRenderEdge(
          graph_layout.GetEdges().find(std::make_pair(p, c))->second,
          graph_layout.GetNodes(), cr);
    }
  }
}

}  // namespace

// TODO(b/193165551): Remove and create better exemplary samples and tests.
void BuildSampleLayout(ZebraixLayout* graph_layout) {
  constexpr double kNodeRadius = 10.5;
  constexpr double kLabelRadius = 16.5;
  constexpr double kNodeSep = 72.0;

  graph_layout->GetNodes().resize(9);
  graph_layout->GetDominance().d_nodes.resize(9);
  for (int i = 0; i < 8; ++i) {
    ZebraixNode& r_node = graph_layout->GetNodes()[i];
    ZebraixEdge& r_edge = graph_layout->GetEdges()[std::make_pair(8, i)];
    r_node.radius = kNodeRadius;

    cairo_matrix_t rotation_matrix;
    cairo_matrix_init_rotate(&rotation_matrix, 0.25 * i * M_PI);

    r_node.centre_x = kNodeSep;
    r_node.centre_y = 0.0;
    cairo_matrix_transform_distance(&rotation_matrix, &r_node.centre_x,
                                    &r_node.centre_y);

    r_node.label_anchor_x = kLabelRadius;
    r_node.label_anchor_y = 0.0;
    cairo_matrix_transform_distance(&rotation_matrix, &r_node.label_anchor_x,
                                    &r_node.label_anchor_y);

    r_node.label.label_text = "Node label";
    r_node.label.anchor = OctantToAnchor(i + 4);

    r_node.display = zebraix_proto::SHOW;

    r_edge.start_node_index = 8;
    r_edge.finish_node_index = i;
    r_edge.forward_arrow = true;
    r_edge.display = zebraix_proto::SHOW;
    graph_layout->GetDominance().d_nodes[i].parents.resize(1);
    graph_layout->GetDominance().d_nodes[i].parents = {8};
  }
  graph_layout->GetDominance().d_nodes[8].children.resize(8);
  graph_layout->GetDominance().d_nodes[8].children = {0, 1, 2, 3, 4, 5, 6, 7};
  ZebraixNode& r_node = graph_layout->GetNodes()[8];
  r_node.radius = kNodeRadius;
  r_node.centre_x = 0.0;
  r_node.centre_y = 0.0;

  graph_layout->GetNodes()[2].label.label_text = "Hello, world";
  graph_layout->GetNodes()[6].label.label_text = "Eat a fluffy soufflé";

  graph_layout->GetNodes()[1].display = zebraix_proto::GHOST;
  graph_layout->GetNodes()[7].display = zebraix_proto::GHOST;
  graph_layout->GetEdges()[std::make_pair(8, 1)].display = zebraix_proto::GHOST;
  graph_layout->GetEdges()[std::make_pair(8, 7)].display = zebraix_proto::GHOST;

  CHECK(graph_layout->GetNodes()[1].label.anchor == zebraix_proto::BL);
  CHECK(graph_layout->GetNodes()[5].label.anchor == zebraix_proto::TR);

  graph_layout->GetCanvas().canvas_width = 72.0 * 5;
  graph_layout->GetCanvas().canvas_height = 72.0 * 4;
  graph_layout->GetCanvas().canvas_x_offset =
      0.5 * graph_layout->GetCanvas().canvas_width;
  graph_layout->GetCanvas().canvas_y_offset =
      0.5 * graph_layout->GetCanvas().canvas_height;
}

void ZebraixCairoSvg::GlobalTearDown() {
  // Global caches, etc. Explicit destruction required to avoid terminal
  // memory leak errors.
  cairo_debug_reset_static_data();
  FcFini();  // FontConfig.
}

ZebraixCairoSvg::ZebraixCairoSvg(const ZebraixCanvas& overall,
                                 const char* out_file) {
  surface = cairo_svg_surface_create(out_file, overall.canvas_width,
                                     overall.canvas_height);
  CHECK_EQ(cairo_surface_status(surface), CAIRO_STATUS_SUCCESS);

  cr = cairo_create(surface);
  CHECK_EQ(cairo_status(cr), CAIRO_STATUS_SUCCESS);
  cairo_translate(cr, overall.canvas_x_offset, 1.0 * overall.canvas_y_offset);
}

ZebraixCairoSvg::~ZebraixCairoSvg() {
  // Explicit destruction required to avoid terminal memory leak errors.
  cairo_surface_destroy(surface);
  cairo_destroy(cr);
}

void ZebraixCairoSvg::RenderToSvg(const ZebraixLayout& layout) {
  CairoRenderLayout(layout, cr);
  cairo_surface_finish(surface);
}

}  // namespace render
}  // namespace zebraix


<|start_filename|>scann/scann/utils/bit_iterator.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_UTILS_BIT_ITERATOR_H_
#define SCANN_UTILS_BIT_ITERATOR_H_

#include <cstdint>

#include "scann/oss_wrappers/scann_bits.h"
#include "scann/utils/types.h"

namespace research_scann {

template <typename StorageT, typename PositionT = uint32_t>
class BitIterator {
 public:
  BitIterator() {}

  explicit inline BitIterator(StorageT whitelist_storage);

  PositionT value() const { return value_; }

  inline void Next();

  bool Done() const { return cur_whitelist_word_ == 0; }

 private:
  inline void NextNonzeroWord();

  StorageT whitelist_array_;

  size_t whitelist_array_position_ = 0;

  PositionT word_base_ = 0;

  PositionT value_ = 0;

  using WordType = decay_t<decltype(whitelist_array_[0])>;
  WordType cur_whitelist_word_ = 0;
};

namespace bits_internal {

template <typename Int>
inline int FindLSBSetNonZero(Int x) {
  DCHECK_NE(x, 0);
  if (sizeof(x) * CHAR_BIT > 32) {
    return bits::FindLSBSetNonZero64(x);
  } else {
    return bits::FindLSBSetNonZero(x);
  }
}

}  // namespace bits_internal

template <typename StorageT, typename PositionT>
BitIterator<StorageT, PositionT>::BitIterator(StorageT whitelist_storage)
    : whitelist_array_(std::move(whitelist_storage)), value_(0) {
  NextNonzeroWord();
  if (cur_whitelist_word_ == 0) return;
  value_ = word_base_ + bits_internal::FindLSBSetNonZero(cur_whitelist_word_);
}

template <typename StorageT, typename PositionT>
void BitIterator<StorageT, PositionT>::Next() {
  cur_whitelist_word_ &= cur_whitelist_word_ - 1;
  if (cur_whitelist_word_ == 0) {
    ++whitelist_array_position_;
    NextNonzeroWord();
    if (cur_whitelist_word_ == 0) return;
  }
  value_ = word_base_ + bits_internal::FindLSBSetNonZero(cur_whitelist_word_);
}

template <typename StorageT, typename PositionT>
void BitIterator<StorageT, PositionT>::NextNonzeroWord() {
  static constexpr size_t kBitsPerWord = sizeof(cur_whitelist_word_) * CHAR_BIT;
  for (; whitelist_array_position_ < whitelist_array_.size();
       ++whitelist_array_position_) {
    cur_whitelist_word_ = whitelist_array_[whitelist_array_position_];
    if (cur_whitelist_word_ != 0) {
      word_base_ = whitelist_array_position_ * kBitsPerWord;
      return;
    }
  }
}

}  // namespace research_scann

#endif


<|start_filename|>cube_unfoldings/perfect_matchings.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef CUBE_UNFOLDINGS_PERFECT_MATCHINGS_H_
#define CUBE_UNFOLDINGS_PERFECT_MATCHINGS_H_
#include <array>
#include <vector>

#include "absl/functional/function_ref.h"

namespace cube_unfoldings {
using Edge = std::array<int, 2>;
using Matching = std::vector<Edge>;

// Calls `cb` for all the perfect matchings in the graph with `n` nodes and
// edges defined by `edges`.
// Matchings are returned in "canonical form", i.e. with all edges in ascending
// order.
void AllPerfectMatchings(const std::vector<Edge>& edges, size_t n,
                         absl::FunctionRef<void(const Matching&)> cb);
}  // namespace cube_unfoldings

#endif  // CUBE_UNFOLDINGS_PERFECT_MATCHINGS_H_


<|start_filename|>parallel_clustering/parallel/parallel-graph-utils.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef RESEARCH_GRAPH_IN_MEMORY_PARALLEL_PARALLEL_GRAPH_UTILS_H_
#define RESEARCH_GRAPH_IN_MEMORY_PARALLEL_PARALLEL_GRAPH_UTILS_H_

#include <cstdio>

#include "external/gbbs/gbbs/gbbs.h"
#include "external/gbbs/gbbs/graph_io.h"
#include "external/gbbs/gbbs/macros.h"
#include "external/gbbs/pbbslib/seq.h"
#include "external/gbbs/pbbslib/sequence_ops.h"
#include "external/gbbs/pbbslib/utilities.h"
#include "parallel/parallel-sequence-ops.h"

namespace research_graph {

struct OffsetsEdges {
  std::vector<gbbs::uintE> offsets;
  std::unique_ptr<std::tuple<gbbs::uintE, float>[]> edges;
  std::size_t num_edges;
};

// Given get_key, which is nondecreasing, defined for 0, ..., num_keys-1, and
// returns an unsigned integer less than n, return an array of length n + 1
// where array[i] := minimum index k such that get_key(k) >= i.
// Note that array[n] = the total number of keys, num_keys.
std::vector<gbbs::uintE> GetOffsets(
    const std::function<gbbs::uintE(std::size_t)>& get_key,
    gbbs::uintE num_keys, std::size_t n);

// Using parallel sorting, compute inter cluster edges given a set of
// cluster_ids that form the vertices of the new graph. Uses scale_func to first
// scale edge weights, and then uses aggregate_func to combine multiple edges on
// the same cluster ids. Note that aggregate_func must be commutative and
// associative, with 0 as its identity. Returns sorted edges and offsets array
// in edges and offsets respectively. The number of compressed vertices should
// be 1 + the maximum cluster id in cluster_ids.
OffsetsEdges ComputeInterClusterEdgesSort(
    gbbs::symmetric_ptr_graph<gbbs::symmetric_vertex, float>& original_graph,
    const std::vector<gbbs::uintE>& cluster_ids,
    std::size_t num_compressed_vertices,
    const std::function<float(float, float)>& aggregate_func,
    const std::function<bool(gbbs::uintE, gbbs::uintE)>& is_valid_func,
    const std::function<float(std::tuple<gbbs::uintE, gbbs::uintE, float>)>&
        scale_func);

// Given an array of edges (given by a tuple consisting of the second endpoint
// and a weight if the edges are weighted) and the offsets marking the index
// of the first edge corresponding to each vertex (essentially, CSR format),
// return the corresponding graph in GBBS format.
// Note that the returned graph takes ownership of the edges array.
template <typename WeightType>
std::unique_ptr<gbbs::symmetric_ptr_graph<gbbs::symmetric_vertex, WeightType>>
MakeGbbsGraph(
    const std::vector<gbbs::uintE>& offsets, std::size_t num_vertices,
    std::unique_ptr<std::tuple<gbbs::uintE, WeightType>[]> edges_pointer,
    std::size_t num_edges) {
  gbbs::symmetric_vertex<WeightType>* vertices =
      new gbbs::symmetric_vertex<WeightType>[num_vertices];
  auto edges = edges_pointer.release();

  pbbs::parallel_for(0, num_vertices, [&](std::size_t i) {
    gbbs::vertex_data vertex_data{offsets[i], offsets[i + 1] - offsets[i]};
    vertices[i] = gbbs::symmetric_vertex<WeightType>(edges, vertex_data);
  });

  return std::make_unique<
      gbbs::symmetric_ptr_graph<gbbs::symmetric_vertex, WeightType>>(
      num_vertices, num_edges, vertices, [=]() {
        delete[] vertices;
        delete[] edges;
      });
}

// Given new cluster ids in compressed_cluster_ids, remap the original
// cluster ids. A cluster id of UINT_E_MAX indicates that the vertex
// has already been placed into a finalized cluster, and this is
// preserved in the remapping.
std::vector<gbbs::uintE> FlattenClustering(
    const std::vector<gbbs::uintE>& cluster_ids,
    const std::vector<gbbs::uintE>& compressed_cluster_ids);

// Holds a GBBS graph and a corresponding node weights
struct GraphWithWeights {
  GraphWithWeights() {}
  GraphWithWeights(
      std::unique_ptr<gbbs::symmetric_ptr_graph<gbbs::symmetric_vertex, float>>
          graph_,
      std::vector<double> node_weights_)
      : graph(std::move(graph_)), node_weights(std::move(node_weights_)) {}
  std::unique_ptr<gbbs::symmetric_ptr_graph<gbbs::symmetric_vertex, float>>
      graph;
  std::vector<double> node_weights;
};

}  // namespace research_graph

#endif  // RESEARCH_GRAPH_IN_MEMORY_PARALLEL_PARALLEL_GRAPH_UTILS_H_


<|start_filename|>scann/scann/hashes/asymmetric_hashing2/training.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_HASHES_ASYMMETRIC_HASHING2_TRAINING_H_
#define SCANN_HASHES_ASYMMETRIC_HASHING2_TRAINING_H_

#include <utility>

#include "scann/data_format/dataset.h"
#include "scann/hashes/asymmetric_hashing2/training_model.h"
#include "scann/hashes/asymmetric_hashing2/training_options.h"
#include "scann/hashes/internal/asymmetric_hashing_impl.h"
#include "scann/hashes/internal/stacked_quantizers.h"
#include "scann/oss_wrappers/scann_down_cast.h"
#include "scann/oss_wrappers/scann_threadpool.h"
#include "scann/utils/types.h"

namespace research_scann {
namespace asymmetric_hashing2 {

template <typename T>
StatusOr<unique_ptr<Model<T>>> TrainSingleMachine(
    const TypedDataset<T>& dataset, const TrainingOptions<T>& params,
    shared_ptr<ThreadPool> pool = nullptr) {
  if (params.config().quantization_scheme() ==
      AsymmetricHasherConfig::STACKED) {
    if (!dataset.IsDense())
      return InvalidArgumentError(
          "Stacked quantizers can only process dense datasets.");
    const auto& dense = down_cast<const DenseDataset<T>&>(dataset);
    TF_ASSIGN_OR_RETURN(
        auto centers,
        ::research_scann::asymmetric_hashing_internal::StackedQuantizers<
            T>::Train(dense, params, pool));
    return Model<T>::FromCenters(std::move(centers),
                                 params.config().quantization_scheme());
  }
  if (params.config().quantization_scheme() ==
      AsymmetricHasherConfig::PRODUCT_AND_BIAS) {
    const auto& dense = down_cast<const DenseDataset<T>&>(dataset);
    DenseDataset<T> dataset_no_bias;
    dataset_no_bias.set_dimensionality(dense.dimensionality() - 1);
    dataset_no_bias.Reserve(dense.size());
    for (const auto& dp : dense) {
      SCANN_RETURN_IF_ERROR(dataset_no_bias.Append(
          MakeDatapointPtr(dp.values(), dp.dimensionality() - 1)));
    }

    TF_ASSIGN_OR_RETURN(
        auto centers,
        ::research_scann::asymmetric_hashing_internal::TrainAsymmetricHashing(
            dataset_no_bias, params, pool));
    auto converted = asymmetric_hashing_internal::ConvertCentersIfNecessary<T>(
        std::move(centers));
    return Model<T>::FromCenters(std::move(converted),
                                 params.config().quantization_scheme());
  } else {
    TF_ASSIGN_OR_RETURN(
        auto centers,
        ::research_scann::asymmetric_hashing_internal::TrainAsymmetricHashing(
            dataset, params, pool));
    auto converted = asymmetric_hashing_internal::ConvertCentersIfNecessary<T>(
        std::move(centers));
    return Model<T>::FromCenters(std::move(converted),
                                 params.config().quantization_scheme());
  }
}

}  // namespace asymmetric_hashing2
}  // namespace research_scann

#endif


<|start_filename|>scann/scann/utils/hash_leaf_helpers.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/utils/hash_leaf_helpers.h"

#include <cstdint>

#include "absl/synchronization/mutex.h"
#include "scann/distance_measures/distance_measure_factory.h"
#include "scann/hashes/asymmetric_hashing2/searcher.h"
#include "scann/hashes/asymmetric_hashing2/training.h"
#include "scann/hashes/asymmetric_hashing2/training_options.h"
#include "scann/projection/projection_factory.h"
#include "scann/proto/distance_measure.pb.h"
#include "scann/utils/types.h"

namespace research_scann {
namespace internal {
namespace {

template <typename T>
using StatusOrSearcher = StatusOr<unique_ptr<SingleMachineSearcherBase<T>>>;

StatusOr<shared_ptr<const DistanceMeasure>>
CreateOrGetAymmetricHashingQuantizationDistance(
    const AsymmetricHasherConfig& ah_config,
    const GenericSearchParameters& params) {
  if (ah_config.has_quantization_distance()) {
    return GetDistanceMeasure(ah_config.quantization_distance());
  } else {
    return params.pre_reordering_dist;
  }
}

template <typename T, typename Lambda>
shared_ptr<DenseDataset<uint8_t>> IndexDatabase(const TypedDataset<T>& dataset,
                                                Lambda index_datapoint_fn,
                                                shared_ptr<ThreadPool> pool) {
  constexpr size_t kIndexingBlockSize = 128;
  vector<Datapoint<uint8_t>> hashed_vec(dataset.size());
  Status status = OkStatus();
  absl::Mutex status_mutex;
  ParallelFor<kIndexingBlockSize>(
      Seq(dataset.size()), pool.get(), [&](size_t i) {
        Status iter_status = index_datapoint_fn(dataset[i], &hashed_vec[i]);
        if (!iter_status.ok()) {
          absl::MutexLock lock(&status_mutex);
          status = iter_status;
          return;
        }
      });
  if (!status.ok()) {
    LOG(WARNING) << status;
    return nullptr;
  }

  auto mutable_hashed = std::make_shared<DenseDataset<uint8_t>>();

  if (!hashed_vec.empty() &&
      hashed_vec[0].dimensionality() > hashed_vec[0].nonzero_entries()) {
    mutable_hashed->set_packing_strategy(HashedItem::NIBBLE);
    mutable_hashed->set_dimensionality(hashed_vec[0].dimensionality());
  }

  mutable_hashed->Reserve(dataset.size());
  for (size_t i = 0; i < dataset.size(); ++i) {
    mutable_hashed->AppendOrDie(hashed_vec[i].ToPtr(), dataset.GetDocid(i));
    FreeBackingStorage(&hashed_vec[i]);
  }

  return mutable_hashed;
}

}  // namespace

template <typename T>
StatusOr<TrainedAsymmetricHashingResults<T>>
HashLeafHelpers<T>::TrainAsymmetricHashingModel(
    shared_ptr<TypedDataset<T>> dataset, const AsymmetricHasherConfig& config,
    const GenericSearchParameters& params, shared_ptr<ThreadPool> pool) {
  if (params.pre_reordering_dist == nullptr) {
    return InvalidArgumentError(
        "pre_reordering_dist in GenericSearchParameters is not "
        "set.");
  }
  TF_ASSIGN_OR_RETURN(
      auto quantization_distance,
      CreateOrGetAymmetricHashingQuantizationDistance(config, params));
  asymmetric_hashing2::TrainingOptions<T> opts(config, quantization_distance,
                                               *dataset);
  TF_ASSIGN_OR_RETURN(
      shared_ptr<const asymmetric_hashing2::Model<T>> model,
      asymmetric_hashing2::TrainSingleMachine<T>(*dataset, opts, pool));
  internal::TrainedAsymmetricHashingResults<T> result;
  result.indexer = std::make_shared<asymmetric_hashing2::Indexer<T>>(
      opts.projector(), quantization_distance, model);
  result.queryer = std::make_shared<asymmetric_hashing2::AsymmetricQueryer<T>>(
      opts.projector(), params.pre_reordering_dist, model);
  result.lookup_type = config.lookup_type();
  result.fixed_point_lut_conversion_options =
      config.fixed_point_lut_conversion_options();
  result.noise_shaping_threshold = config.noise_shaping_threshold();
  if (config.has_centers_filename()) {
    return InvalidArgumentError("Centers file not supported.");
  }
  return result;
}

template <typename T>
StatusOrSearcher<T> HashLeafHelpers<T>::AsymmetricHasherFactory(
    shared_ptr<TypedDataset<T>> dataset,
    shared_ptr<DenseDataset<uint8_t>> hashed_dataset,
    const TrainedAsymmetricHashingResults<T>& training_results,
    const GenericSearchParameters& params, shared_ptr<ThreadPool> pool) {
  if (!hashed_dataset) {
    if (std::isnan(training_results.noise_shaping_threshold)) {
      hashed_dataset = IndexDatabase<T>(
          *dataset,
          [&](const DatapointPtr<T>& dptr, Datapoint<uint8_t>* dp) {
            return training_results.indexer->Hash(dptr, dp);
          },
          pool);
    } else {
      hashed_dataset = IndexDatabase<T>(
          *dataset,
          [&](const DatapointPtr<T>& dptr, Datapoint<uint8_t>* dp) {
            return training_results.indexer->HashWithNoiseShaping(
                dptr, dp, training_results.noise_shaping_threshold);
          },
          pool);
    }
    if (!hashed_dataset) {
      return UnknownError("Could not index database.");
    }
  }

  asymmetric_hashing2::SearcherOptions<T> opts(training_results.queryer,
                                               training_results.indexer);
  opts.set_asymmetric_lookup_type(training_results.lookup_type);
  opts.set_noise_shaping_threshold(training_results.noise_shaping_threshold);
  opts.set_fixed_point_lut_conversion_options(
      training_results.fixed_point_lut_conversion_options);
  return StatusOrSearcher<T>(make_unique<asymmetric_hashing2::Searcher<T>>(
      std::move(dataset), std::move(hashed_dataset), std::move(opts),
      params.pre_reordering_num_neighbors, params.pre_reordering_epsilon));
}

template <typename T>
StatusOr<TrainedAsymmetricHashingResults<T>>
HashLeafHelpers<T>::LoadAsymmetricHashingModel(
    const AsymmetricHasherConfig& config, const GenericSearchParameters& params,
    shared_ptr<ThreadPool> pool, CentersForAllSubspaces* preloaded_codebook) {
  TF_ASSIGN_OR_RETURN(
      auto quantization_distance,
      CreateOrGetAymmetricHashingQuantizationDistance(config, params));
  shared_ptr<const asymmetric_hashing2::Model<T>> model;
  if (preloaded_codebook) {
    TF_ASSIGN_OR_RETURN(
        model, asymmetric_hashing2::Model<T>::FromProto(*preloaded_codebook));
  } else {
    return InvalidArgumentError("Centers files are not supported.");
  }

  TF_ASSIGN_OR_RETURN(shared_ptr<const ChunkingProjection<T>> projector,
                      ChunkingProjectionFactory<T>(config.projection()));
  internal::TrainedAsymmetricHashingResults<T> result;
  result.indexer = std::make_shared<asymmetric_hashing2::Indexer<T>>(
      projector, quantization_distance, model);
  result.queryer = std::make_shared<asymmetric_hashing2::AsymmetricQueryer<T>>(
      projector, params.pre_reordering_dist, model);
  result.lookup_type = config.lookup_type();
  result.fixed_point_lut_conversion_options =
      config.fixed_point_lut_conversion_options();
  result.noise_shaping_threshold = config.noise_shaping_threshold();
  return result;
}

SCANN_INSTANTIATE_TYPED_CLASS(, TrainedAsymmetricHashingResults);
SCANN_INSTANTIATE_TYPED_CLASS(, HashLeafHelpers);

}  // namespace internal
}  // namespace research_scann


<|start_filename|>zebraix/base/dominance_conversion.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "base/dominance_conversion.h"

#include <vector>

#include "base/dominance.h"
#include "base/zebraix_graph.proto.h"

namespace zebraix {
namespace base {

void SelfMergeLayoutIntoBaseValues(zebraix_proto::ZebraixGraph* layout) {
  for (int i = 0; i < layout->nodes_size(); ++i) {
    zebraix_proto::Node p_node_copy;
    p_node_copy.CopyFrom(layout->nodes(i));
    zebraix_proto::Node* fixedup_node = layout->mutable_nodes(i);
    // Copy default node values, then merge the original node back in.
    fixedup_node->CopyFrom(layout->base_node());
    fixedup_node->MergeFrom(p_node_copy);
  }
}

void FillGraphFromProto(const zebraix_proto::ZebraixGraph& p_graph,
                        DominanceGraph* d_graph) {
  const int p_node_count = p_graph.nodes_size();
  d_graph->d_nodes.resize(p_node_count);

  for (int i = 0; i < p_node_count; ++i) {
    d_graph->d_nodes[i].import_index = i;
  }

  for (int i = 0; i < p_node_count; ++i) {
    int import_index = d_graph->d_nodes[i].import_index;
    const zebraix_proto::Node& p_node = p_graph.nodes(import_index);
    d_graph->d_nodes[i].prime_rank = p_node.prime_rank();
    d_graph->d_nodes[i].obverse_rank = p_node.obverse_rank();
  }
}

}  // namespace base
}  // namespace zebraix


<|start_filename|>scann/scann/hashes/internal/asymmetric_hashing_impl.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/hashes/internal/asymmetric_hashing_impl.h"

#include <cstdint>
#include <numeric>

#include "absl/random/distributions.h"
#include "scann/data_format/datapoint.h"
#include "scann/distance_measures/one_to_many/one_to_many.h"
#include "scann/hashes/internal/asymmetric_hashing_postprocess.h"
#include "scann/oss_wrappers/scann_random.h"
#include "scann/oss_wrappers/scann_status.h"
#include "scann/projection/chunking_projection.h"
#include "scann/utils/common.h"
#include "scann/utils/gmm_utils.h"
#include "scann/utils/noise_shaping_utils.h"
#include "scann/utils/top_n_amortized_constant.h"
#include "scann/utils/types.h"

namespace research_scann {
namespace asymmetric_hashing_internal {

class ParallelPerpendicularDistance : public DistanceMeasure {
 public:
  SCANN_DECLARE_DISTANCE_MEASURE_VIRTUAL_METHODS(NOT_SPECIALLY_OPTIMIZED);

  void set_parallel_cost_multiplier(double x) { parallel_cost_multiplier_ = x; }

 private:
  template <typename T>
  SCANN_INLINE double GetDistanceDenseImpl(
      const DatapointPtr<T>& x_dptr, const DatapointPtr<T>& y_dptr) const {
    DCHECK(x_dptr.IsDense());
    DCHECK(y_dptr.IsDense());
    auto x = x_dptr.values_slice();
    auto y = y_dptr.values_slice();
    DCHECK_EQ(x.size(), y.size());
    double parallel_error = 0.0;
    double residual_squared_norm = 0.0;
    for (size_t i : IndicesOf(x)) {
      const double residual_coord =
          static_cast<double>(y[i]) - static_cast<double>(x[i]);
      parallel_error += residual_coord * x[i];
      residual_squared_norm += Square(residual_coord);
    }
    parallel_error = Square(parallel_error);
    const double perpendicular_error = residual_squared_norm - parallel_error;
    return parallel_cost_multiplier_ * parallel_error + perpendicular_error;
  }

  template <typename T>
  SCANN_INLINE double GetDistanceSparseImpl(
      const DatapointPtr<T>& x_dptr, const DatapointPtr<T>& y_dptr) const {
    LOG(FATAL) << "NOT IMPLEMENTED";
  }

  template <typename T>
  SCANN_INLINE double GetDistanceHybridImpl(
      const DatapointPtr<T>& x_dptr, const DatapointPtr<T>& y_dptr) const {
    LOG(FATAL) << "NOT IMPLEMENTED";
  }

  double parallel_cost_multiplier_ = 1.0;
};

SCANN_DEFINE_DISTANCE_MEASURE_VIRTUAL_METHODS(ParallelPerpendicularDistance,
                                              numeric_limits<size_t>::max());

string_view ParallelPerpendicularDistance::name() const {
  return "ParallelPerpendicularDistance";
}

namespace {

double ComputeNormBiasCorrection(const DenseDataset<double>& db,
                                 DatapointPtr<double> center,
                                 ConstSpan<DatapointIndex> cluster_members) {
  double mean_norm = 0.0;
  for (DatapointIndex idx : cluster_members) {
    mean_norm += std::sqrt(SquaredL2Norm(db[idx]));
  }
  mean_norm /= cluster_members.size();
  const double center_norm = std::sqrt(SquaredL2Norm(center));
  return (center_norm == 0.0) ? 1.0 : (mean_norm / center_norm);
}

}  // namespace

template <typename T>
StatusOr<vector<DenseDataset<double>>> AhImpl<T>::TrainAsymmetricHashing(
    const TypedDataset<T>& dataset, const TrainingOptionsT& opts,
    shared_ptr<ThreadPool> pool) {
  if (dataset.empty()) {
    return InvalidArgumentError("Cannot train AH on an empty dataset.");
  }

  ChunkedDatapoint<double> chunked_vec;

  if (opts.preprocessing_function()) {
    TF_ASSIGN_OR_RETURN(Datapoint<T> preprocessed,
                        opts.preprocessing_function()(dataset[0]));
    SCANN_RETURN_IF_ERROR(
        opts.projector()->ProjectInput(preprocessed.ToPtr(), &chunked_vec));
  } else {
    SCANN_RETURN_IF_ERROR(
        opts.projector()->ProjectInput(dataset[0], &chunked_vec));
  }
  int32_t num_blocks = chunked_vec.size();
  vector<DenseDataset<double>> chunked_dataset(num_blocks);

  const float sampling_fraction =
      opts.config().has_expected_sample_size()
          ? std::min(1.0,
                     static_cast<double>(opts.config().expected_sample_size()) /
                         static_cast<double>(dataset.size()))
          : opts.config().sampling_fraction();

  if (sampling_fraction == 1.0) {
    for (int32_t i = 0; i < num_blocks; ++i) {
      DenseDataset<double>& ds = chunked_dataset[i];
      ds.set_dimensionality(chunked_vec[i].dimensionality());
      ds.Reserve(dataset.size());
    }
  }

  MTRandom rng(kDeterministicSeed * (opts.config().sampling_seed() + 1));
  vector<DatapointIndex> sample;
  for (DatapointIndex i = 0; i < dataset.size(); ++i) {
    if (absl::Uniform<double>(rng, 0, 1.0) < sampling_fraction) {
      sample.push_back(i);
    }
  }

  if (sample.size() > opts.config().max_sample_size()) {
    std::shuffle(sample.begin(), sample.end(), rng);
    sample.resize(opts.config().max_sample_size());
    std::sort(sample.begin(), sample.end());
  }

  if (sample.size() < opts.config().num_clusters_per_block()) {
    return InvalidArgumentError(absl::StrCat(
        "Number of clusters per block (",
        opts.config().num_clusters_per_block(),
        ") is greater than asymmetric hashing training data size (",
        sample.size(), ")."));
  }

  auto append_chunked_blocks = [&] {
    for (size_t j = 0; j < num_blocks; ++j) {
      chunked_dataset[j].AppendOrDie(chunked_vec[j], "");
    }
  };

  if (opts.preprocessing_function()) {
    for (DatapointIndex i : sample) {
      TF_ASSIGN_OR_RETURN(Datapoint<T> preprocessed,
                          opts.preprocessing_function()(dataset[i]));
      SCANN_RETURN_IF_ERROR(
          opts.projector()->ProjectInput(preprocessed.ToPtr(), &chunked_vec));
      append_chunked_blocks();
    }
  } else {
    for (DatapointIndex i : sample) {
      SCANN_RETURN_IF_ERROR(
          opts.projector()->ProjectInput(dataset[i], &chunked_vec));
      append_chunked_blocks();
    }
  }

  auto quantization_distance = opts.quantization_distance();
  GmmUtils::Options gmm_opts;
  gmm_opts.seed = opts.config().clustering_seed();
  gmm_opts.max_iterations = opts.config().max_clustering_iterations();
  gmm_opts.epsilon = opts.config().clustering_convergence_tolerance();
  gmm_opts.parallelization_pool = std::move(pool);
  if (!std::isnan(opts.config().noise_shaping_threshold()) &&
      opts.config().use_noise_shaped_training()) {
    gmm_opts.parallel_cost_multiplier = ComputeParallelCostMultiplier(
        opts.config().noise_shaping_threshold(), 1.0, dataset.dimensionality());
    auto d = make_shared<ParallelPerpendicularDistance>();
    d->set_parallel_cost_multiplier(gmm_opts.parallel_cost_multiplier);
    quantization_distance = d;
  }
  GmmUtils gmm(quantization_distance, gmm_opts);

  vector<DenseDataset<double>> all_centers(num_blocks);
  for (size_t i : Seq(num_blocks)) {
    DenseDataset<double> centers;
    vector<vector<DatapointIndex>> subpartitions;
    SCANN_RETURN_IF_ERROR(gmm.GenericKmeans(
        chunked_dataset[i], opts.config().num_clusters_per_block(), &centers,
        &subpartitions));

    if (opts.config().use_norm_biasing_correction()) {
      for (size_t center_idx : IndicesOf(centers)) {
        const double norm_bias_correction = ComputeNormBiasCorrection(
            chunked_dataset[i], centers[center_idx], subpartitions[center_idx]);
        for (double& d : centers.mutable_data(center_idx)) {
          d *= norm_bias_correction;
        }
      }
    }

    chunked_dataset[i].clear();
    chunked_dataset[i].ShrinkToFit();

    vector<uint32_t> centers_permutation(centers.size());
    std::iota(centers_permutation.begin(), centers_permutation.end(), 0U);
    std::sort(centers_permutation.begin(), centers_permutation.end(),
              [&subpartitions](uint32_t a, uint32_t b) {
                return subpartitions[a].size() > subpartitions[b].size();
              });

    constexpr size_t kAssumedCacheLineSize = 64;
    constexpr size_t kFloatsPerCacheLine =
        kAssumedCacheLineSize / sizeof(float);
    const uint64_t cache_lines_per_row =
        std::max(static_cast<size_t>(1), centers.size() / kFloatsPerCacheLine);
    const size_t num_rotate =
        ((i / 2) % cache_lines_per_row) * kFloatsPerCacheLine;
    std::rotate(centers_permutation.begin(),
                centers_permutation.begin() + num_rotate,
                centers_permutation.end());

    if (i & 1) {
      std::reverse(centers_permutation.begin(), centers_permutation.end());
    }

    for (uint32_t j : centers_permutation) {
      all_centers[i].AppendOrDie(centers[j], "");
    }
  }

  return std::move(all_centers);
}

template <typename T>
Status AhImpl<T>::IndexDatapoint(const DatapointPtr<T>& input,
                                 const ChunkingProjection<T>& projection,
                                 const DistanceMeasure& quantization_distance,
                                 ConstSpan<DenseDataset<FloatT>> centers,
                                 MutableSpan<uint8_t> result) {
  DCHECK(!centers.empty());
  ChunkedDatapoint<FloatT> projected;
  SCANN_RETURN_IF_ERROR(projection.ProjectInput(input, &projected));

  DCHECK_LE(projected.size(), result.size());

  vector<float> distances(centers[0].size());
  DCHECK_GE(distances.size(), 1);
  DCHECK_LE(distances.size(), 256);
  for (size_t i = 0; i < projected.size(); ++i) {
    DCHECK_EQ(centers[0].size(), centers[i].size());
    size_t closest = 0;
    const DatapointPtr<FloatT> projected_ptr = projected[i];
    const DenseDataset<FloatT>& cur_centers = centers[i];

    const size_t centers_size = cur_centers.size();

    if (projected_ptr.IsSparse()) {
      double closest_distance = numeric_limits<double>::infinity();
      for (size_t j = 0; j < centers_size; ++j) {
        const double distance = quantization_distance.GetDistanceHybrid(
            projected_ptr, cur_centers[j]);
        if (ABSL_PREDICT_FALSE(distance < closest_distance)) {
          closest_distance = distance;
          closest = j;
        }
      }
    } else {
      DCHECK_EQ(distances.size(), cur_centers.size());
      DenseDistanceOneToMany(quantization_distance, projected_ptr, cur_centers,
                             MutableSpan<float>(distances));
      auto min_it = std::min_element(distances.begin(), distances.end());
      closest = min_it - distances.begin();
    }
    result[i] = closest;
  }

  return OkStatus();
}

template <typename T>
Status AhImpl<T>::IndexDatapoint(const DatapointPtr<T>& input,
                                 const ChunkingProjection<T>& projection,
                                 const DistanceMeasure& quantization_distance,
                                 ConstSpan<DenseDataset<FloatT>> centers,
                                 Datapoint<uint8_t>* result) {
  DatapointIndex result_size = centers.size();
  DCHECK_EQ(result_size, projection.num_blocks());
  result->clear();
  result->mutable_values()->resize(result_size, 0);

  return AhImpl<T>::IndexDatapoint(input, projection, quantization_distance,
                                   centers, result->mutable_values_slice());
}

namespace {

template <typename T>
T Square(T x) {
  return x * x;
}

double ComputeParallelCostMultiplier(double t, double squared_l2_norm,
                                     DimensionIndex dims) {
  const double parallel_cost = Square(t) / squared_l2_norm;
  const double perpendicular_cost =
      (1.0 - Square(t) / squared_l2_norm) / (dims - 1.0);
  return parallel_cost / perpendicular_cost;
}

struct SubspaceResidualStats {
  double residual_norm = 0.0;

  double parallel_residual_component = 0.0;
};

template <typename T>
SubspaceResidualStats ComputeResidualStatsForCluster(
    ConstSpan<T> maybe_residual_dptr, ConstSpan<T> original_dptr,
    double inv_norm, ConstSpan<FloatingTypeFor<T>> quantized) {
  DCHECK_EQ(maybe_residual_dptr.size(), quantized.size());
  const size_t dims = maybe_residual_dptr.size();
  SubspaceResidualStats result;
  for (size_t i : Seq(dims)) {
    const double residual_coordinate =
        static_cast<double>(maybe_residual_dptr[i]) -
        static_cast<double>(quantized[i]);
    result.residual_norm += Square(residual_coordinate);
    result.parallel_residual_component +=
        residual_coordinate * original_dptr[i] * inv_norm;
  }
  return result;
}

template <typename T>
StatusOr<vector<std::vector<SubspaceResidualStats>>> ComputeResidualStats(
    DatapointPtr<T> maybe_residual_dptr, DatapointPtr<T> original_dptr,
    ConstSpan<DenseDataset<FloatingTypeFor<T>>> centers,
    const ChunkingProjection<T>& projection) {
  const size_t num_subspaces = centers.size();
  DCHECK_GE(num_subspaces, 1);
  vector<std::vector<SubspaceResidualStats>> result(num_subspaces);
  vector<std::vector<SubspaceResidualStats>> residual_stats(num_subspaces);
  const size_t num_clusters_per_block = centers[0].size();

  using FloatT = FloatingTypeFor<T>;
  ChunkedDatapoint<FloatT> maybe_residual_dptr_chunked;
  ChunkedDatapoint<FloatT> original_dptr_chunked;
  SCANN_RETURN_IF_ERROR(projection.ProjectInput(maybe_residual_dptr,
                                                &maybe_residual_dptr_chunked));
  SCANN_RETURN_IF_ERROR(
      projection.ProjectInput(original_dptr, &original_dptr_chunked));
  SCANN_RET_CHECK_EQ(maybe_residual_dptr_chunked.size(), num_subspaces);
  SCANN_RET_CHECK_EQ(original_dptr_chunked.size(), num_subspaces);
  double chunked_norm = 0.0;
  for (size_t subspace_idx : Seq(num_subspaces)) {
    for (FloatT x : original_dptr_chunked[subspace_idx].values_slice()) {
      chunked_norm += Square<double>(x);
    }
  }
  chunked_norm = std::sqrt(chunked_norm);
  double inverse_chunked_norm = 1.0 / chunked_norm;

  for (size_t subspace_idx : Seq(num_subspaces)) {
    auto& cur_subspace_residual_stats = residual_stats[subspace_idx];
    cur_subspace_residual_stats.resize(num_clusters_per_block);
    const DenseDataset<FloatingTypeFor<T>>& cur_subspace_centers =
        centers[subspace_idx];
    for (size_t cluster_idx : Seq(num_clusters_per_block)) {
      ConstSpan<FloatingTypeFor<T>> center =
          cur_subspace_centers[cluster_idx].values_slice();
      ConstSpan<FloatT> maybe_residual_dptr_span =
          maybe_residual_dptr_chunked[subspace_idx].values_slice();
      ConstSpan<FloatT> original_dptr_span =
          original_dptr_chunked[subspace_idx].values_slice();
      cur_subspace_residual_stats[cluster_idx] = ComputeResidualStatsForCluster(
          maybe_residual_dptr_span, original_dptr_span, inverse_chunked_norm,
          center);
    }
  }
  return residual_stats;
}

void InitializeToMinResidualNorm(
    ConstSpan<std::vector<SubspaceResidualStats>> residual_stats,
    MutableSpan<uint8_t> result) {
  DCHECK_EQ(result.size(), residual_stats.size());
  for (size_t subspace_idx : IndicesOf(residual_stats)) {
    auto it = std::min_element(
        residual_stats[subspace_idx].begin(),
        residual_stats[subspace_idx].end(),
        [](const SubspaceResidualStats& a, const SubspaceResidualStats& b) {
          return a.residual_norm < b.residual_norm;
        });
    result[subspace_idx] = it - residual_stats[subspace_idx].begin();
  }
}

double ComputeParallelResidualComponent(
    ConstSpan<uint8_t> quantized,
    ConstSpan<std::vector<SubspaceResidualStats>> residual_stats) {
  double result = 0.0;
  for (size_t subspace_idx : IndicesOf(quantized)) {
    const uint8_t cluster_idx = quantized[subspace_idx];
    result +=
        residual_stats[subspace_idx][cluster_idx].parallel_residual_component;
  }
  return result;
}

struct CoordinateDescentResult {
  uint8_t new_center_idx = 0;
  double cost_delta = 0.0;
  double new_parallel_residual_component = 0.0;
};

CoordinateDescentResult OptimizeSingleSubspace(
    ConstSpan<SubspaceResidualStats> cur_subspace_residual_stats,
    const uint8_t cur_center_idx, const double parallel_residual_component,
    const double parallel_cost_multiplier) {
  CoordinateDescentResult result;
  result.new_center_idx = cur_center_idx;
  result.new_parallel_residual_component = parallel_residual_component;
  const double old_subspace_residual_norm =
      cur_subspace_residual_stats[cur_center_idx].residual_norm;
  const double old_subspace_parallel_component =
      cur_subspace_residual_stats[cur_center_idx].parallel_residual_component;
  for (size_t new_center_idx : IndicesOf(cur_subspace_residual_stats)) {
    if (new_center_idx == cur_center_idx) continue;
    const SubspaceResidualStats& rs =
        cur_subspace_residual_stats[new_center_idx];
    const double new_parallel_residual_component =
        parallel_residual_component - old_subspace_parallel_component +
        rs.parallel_residual_component;
    const double parallel_norm_delta = Square(new_parallel_residual_component) -
                                       Square(parallel_residual_component);
    if (parallel_norm_delta > 0.0) continue;
    const double residual_norm_delta =
        rs.residual_norm - old_subspace_residual_norm;
    const double perpendicular_norm_delta =
        residual_norm_delta - parallel_norm_delta;
    const double cost_delta = parallel_cost_multiplier * parallel_norm_delta +
                              perpendicular_norm_delta;
    if (cost_delta < result.cost_delta) {
      result.new_center_idx = new_center_idx;
      result.cost_delta = cost_delta;
      result.new_parallel_residual_component = new_parallel_residual_component;
    }
  }
  return result;
}

template <typename T>
Status CoordinateDescentAHQuantize(
    DatapointPtr<T> maybe_residual_dptr, DatapointPtr<T> original_dptr,
    ConstSpan<DenseDataset<FloatingTypeFor<T>>> centers,
    const ChunkingProjection<T>& projection, double threshold,
    MutableSpan<uint8_t> result, int* num_changes = nullptr,
    double* residual_ptr = nullptr, double* parallel_residual_ptr = nullptr) {
  SCANN_RET_CHECK_EQ(result.size(), centers.size());
  SCANN_RET_CHECK_EQ(maybe_residual_dptr.dimensionality(),
                     original_dptr.dimensionality());
  TF_ASSIGN_OR_RETURN(auto residual_stats,
                      ComputeResidualStats(maybe_residual_dptr, original_dptr,
                                           centers, projection));

  const double parallel_cost_multiplier = ComputeParallelCostMultiplier(
      threshold, SquaredL2Norm(original_dptr), original_dptr.dimensionality());
  InitializeToMinResidualNorm(residual_stats, result);
  double parallel_residual_component =
      ComputeParallelResidualComponent(result, residual_stats);

  vector<uint16_t> subspace_idxs(result.size());
  std::iota(subspace_idxs.begin(), subspace_idxs.end(), 0U);
  vector<double> subspace_residual_norms(result.size());
  for (size_t subspace_idx : IndicesOf(result)) {
    const uint8_t cluster_idx = result[subspace_idx];
    subspace_residual_norms[subspace_idx] =
        residual_stats[subspace_idx][cluster_idx].residual_norm;
  }
  std::vector<uint8_t> result_sorted(result.begin(), result.end());
  ZipSortBranchOptimized(
      std::greater<double>(), subspace_residual_norms.begin(),
      subspace_residual_norms.end(), result_sorted.begin(), result_sorted.end(),
      subspace_idxs.begin(), subspace_idxs.end());

  enum { kMaxRounds = 10 };
  bool cur_round_changes = true;
  if (num_changes) *num_changes = 0;
  for (int round = 0; cur_round_changes && round < kMaxRounds; ++round) {
    cur_round_changes = false;
    for (size_t i : IndicesOf(subspace_idxs)) {
      const size_t subspace_idx = subspace_idxs[i];
      ConstSpan<SubspaceResidualStats> cur_subspace_residual_stats =
          residual_stats[subspace_idx];
      const uint8_t cur_center_idx = result_sorted[i];
      auto subspace_result = OptimizeSingleSubspace(
          cur_subspace_residual_stats, cur_center_idx,
          parallel_residual_component, parallel_cost_multiplier);
      if (subspace_result.new_center_idx != cur_center_idx) {
        if (num_changes) ++*num_changes;
        parallel_residual_component =
            subspace_result.new_parallel_residual_component;
        result_sorted[i] = subspace_result.new_center_idx;
        cur_round_changes = true;
      }
    }
  }

  double final_residual_norm = 0.0;
  for (size_t i : IndicesOf(result_sorted)) {
    const size_t subspace_idx = subspace_idxs[i];
    const uint8_t center_idx = result_sorted[i];
    result[subspace_idx] = center_idx;
    final_residual_norm +=
        residual_stats[subspace_idx][center_idx].residual_norm;
  }
  if (residual_ptr) *residual_ptr = final_residual_norm;
  if (parallel_residual_ptr) {
    *parallel_residual_ptr = Square(parallel_residual_component);
  }
  return OkStatus();
}

}  // namespace

template <typename T>
Status AhImpl<T>::IndexDatapointNoiseShaped(
    const DatapointPtr<T>& maybe_residual_dptr,
    const DatapointPtr<T>& original_dptr,
    const ChunkingProjection<T>& projection,
    ConstSpan<DenseDataset<FloatingTypeFor<T>>> centers, double threshold,
    MutableSpan<uint8_t> result) {
  return CoordinateDescentAHQuantize<T>(maybe_residual_dptr, original_dptr,
                                        centers, projection, threshold, result);
}

template <typename T>
StatusOr<vector<float>> AhImpl<T>::CreateRawFloatLookupTable(
    const DatapointPtr<T>& query, const ChunkingProjection<T>& projection,
    const DistanceMeasure& lookup_distance,
    ConstSpan<DenseDataset<FloatT>> centers, int32_t num_clusters_per_block) {
  ChunkedDatapoint<FloatT> projected;
  SCANN_RETURN_IF_ERROR(projection.ProjectInput(query, &projected));
  SCANN_RET_CHECK_EQ(centers.size(), projected.size());

  vector<float> result(num_clusters_per_block * projected.size());
  float* result_row_start = result.data();

  for (size_t i = 0; i < centers.size();
       ++i, result_row_start += num_clusters_per_block) {
    const DatapointPtr<FloatT> projected_ptr = projected[i];
    const DenseDataset<FloatT>& cur_centers = centers[i];
    if (projected_ptr.IsSparse()) {
      for (size_t j = 0; j < num_clusters_per_block; ++j) {
        result_row_start[j] = static_cast<float>(
            lookup_distance.GetDistanceHybrid(projected_ptr, cur_centers[j]));
      }
    } else {
      if (lookup_distance.specially_optimized_distance_tag() ==
          DistanceMeasure::LIMITED_INNER_PRODUCT) {
        DenseDistanceOneToMany(
            DotProductDistance(), projected_ptr, cur_centers,
            MutableSpan<float>(result_row_start, num_clusters_per_block));
      } else {
        DenseDistanceOneToMany(
            lookup_distance, projected_ptr, cur_centers,
            MutableSpan<float>(result_row_start, num_clusters_per_block));
      }
    }
  }

  return std::move(result);
}

namespace {
float ComputeMultiplierByQuantile(ConstSpan<float> raw_lookup, float quantile,
                                  int32_t max_integer_value) {
  const size_t k = raw_lookup.size() * (1.0 - quantile) + 1;
  if (k == 1) {
    const float max_abs_lookup_element = std::max(
        std::sqrt(numeric_limits<float>::epsilon()), MaxAbsValue(raw_lookup));
    return max_integer_value / max_abs_lookup_element;
  } else {
    DCHECK_LT(quantile, 1.0f);
    TopNAmortizedConstant<float> tn(k);
    for (auto& elem : raw_lookup) {
      tn.push(std::abs(elem));
    }
    return max_integer_value / tn.exact_bottom();
  }
}

template <typename T, typename Lambda>
inline vector<T> ConvertLookupToFixedPointImpl(ConstSpan<float> raw_lookup,
                                               Lambda convert_to_int_lambda,
                                               float multiplier) {
  constexpr T kBias = FixedPointBias<T>();
  vector<T> result(raw_lookup.size());
  for (size_t i = 0; i < raw_lookup.size(); ++i) {
    result[i] = convert_to_int_lambda(raw_lookup[i] * multiplier) + kBias;
  }
  return result;
}

}  // namespace

template <typename T>
vector<T> ConvertLookupToFixedPoint(
    ConstSpan<float> raw_lookup,
    const AsymmetricHasherConfig::FixedPointLUTConversionOptions&
        conversion_options,
    float* multiplier) {
  DCHECK_GT(conversion_options.multiplier_quantile(), 0.0f);
  DCHECK_LE(conversion_options.multiplier_quantile(), 1.0f);
  using SignedT = make_signed_t<T>;
  *multiplier = ComputeMultiplierByQuantile(
      raw_lookup, conversion_options.multiplier_quantile(),
      numeric_limits<SignedT>::max());
  constexpr int kRound =
      AsymmetricHasherConfig::FixedPointLUTConversionOptions::ROUND;
  if (conversion_options.multiplier_quantile() == 1.0f) {
    if (conversion_options.float_to_int_conversion_method() == kRound) {
      return ConvertLookupToFixedPointImpl<T>(
          raw_lookup, [](float f) { return std::lround(f); }, *multiplier);
    } else {
      return ConvertLookupToFixedPointImpl<T>(
          raw_lookup, [](float f) { return static_cast<SignedT>(f); },
          *multiplier);
    }
  } else {
    auto compress_to_bounds = [](float f) {
      f = std::min<float>(f, numeric_limits<SignedT>::max());
      return std::max<float>(f, numeric_limits<SignedT>::min());
    };
    if (conversion_options.float_to_int_conversion_method() == kRound) {
      return ConvertLookupToFixedPointImpl<T>(
          raw_lookup,
          [&](float f) {
            return static_cast<SignedT>(std::lround(compress_to_bounds(f)));
          },
          *multiplier);
    } else {
      return ConvertLookupToFixedPointImpl<T>(
          raw_lookup,
          [&](float f) { return static_cast<SignedT>(compress_to_bounds(f)); },
          *multiplier);
    }
  }
}

template vector<uint8_t> ConvertLookupToFixedPoint<uint8_t>(
    ConstSpan<float> raw_lookup,
    const AsymmetricHasherConfig::FixedPointLUTConversionOptions&,
    float* multiplier);
template vector<uint16_t> ConvertLookupToFixedPoint<uint16_t>(
    ConstSpan<float> raw_lookup,
    const AsymmetricHasherConfig::FixedPointLUTConversionOptions&,
    float* multiplier);

bool CanUseInt16Accumulator(ConstSpan<uint8_t> lookup_table,
                            size_t num_blocks) {
  const size_t num_centers_per_block = lookup_table.size() / num_blocks;
  DCHECK_EQ(lookup_table.size() % num_blocks, 0);
  if (num_centers_per_block != 16) return false;
  constexpr uint8_t kBias = FixedPointBias<uint8_t>();

  constexpr size_t kMinMin =
      numeric_limits<int16_t>::min() / numeric_limits<int8_t>::min();
  constexpr size_t kMaxMax =
      numeric_limits<int16_t>::max() / numeric_limits<int8_t>::max();
  constexpr size_t kGuaranteedToWork = (kMaxMax < kMinMin) ? kMaxMax : kMinMin;
  if (num_blocks <= kGuaranteedToWork) return true;

  int32_t sum_of_maxes = 0, sum_of_mins = 0;
  auto block_start = lookup_table.begin();
  for (size_t block = num_blocks; block != 0;
       --block, block_start += num_centers_per_block) {
    int8_t max_val = numeric_limits<int8_t>::min();
    int8_t min_val = numeric_limits<int8_t>::max();
    for (size_t i = 0; i < 16; ++i) {
      const int8_t unbiased = block_start[i] - kBias;
      max_val = std::max<int8_t>(unbiased, max_val);
      min_val = std::min<int8_t>(unbiased, min_val);
    }

    sum_of_mins += static_cast<int32_t>(min_val);
    sum_of_maxes += static_cast<int32_t>(max_val);
  }

  return sum_of_maxes <= numeric_limits<int16_t>::max() &&
         sum_of_mins >= numeric_limits<int16_t>::min();
}

vector<uint8_t> CreatePackedDataset(
    const DenseDataset<uint8_t>& hashed_database) {
  vector<uint8_t> packed_dataset;
  if (hashed_database.empty()) {
    return packed_dataset;
  }

  DimensionIndex num_blocks = hashed_database[0].nonzero_entries();
  packed_dataset.resize(num_blocks * ((hashed_database.size() + 31) & (~31)) /
                        2);
  DatapointIndex k = 0;
  for (; k < hashed_database.size() / 32; ++k) {
    size_t start = k * 16 * num_blocks;
    for (size_t j = 0; j < num_blocks; ++j) {
      for (size_t m = 0; m < 16; m++) {
        uint8_t u0 = hashed_database[k * 32 + m].values()[j];
        uint8_t u1 = hashed_database[k * 32 + m + 16].values()[j];
        packed_dataset[start + j * 16 + m] = u1 * 16 + u0;
      }
    }
  }

  if (k * 32 < hashed_database.size()) {
    size_t start = k * 16 * num_blocks;
    for (size_t j = 0; j < num_blocks; ++j) {
      for (size_t m = 0; m < 16; m++) {
        DatapointIndex dp_idx = k * 32 + m;
        dp_idx = dp_idx >= hashed_database.size() ? (hashed_database.size() - 1)
                                                  : dp_idx;
        uint8_t u0 = hashed_database[dp_idx].values()[j];

        dp_idx = k * 32 + m + 16;
        dp_idx = dp_idx >= hashed_database.size() ? (hashed_database.size() - 1)
                                                  : dp_idx;
        uint8_t u1 = hashed_database[dp_idx].values()[j];
        packed_dataset[start + j * 16 + m] = u1 * 16 + u0;
      }
    }
  }

  return packed_dataset;
}

template class UnrestrictedIndexIterator<6, IdentityPostprocessFunctor>;
template class UnrestrictedIndexIterator<6, AddBiasFunctor>;
template class UnrestrictedIndexIterator<6, LimitedInnerFunctor>;
template class PopulateDistancesIterator<6, IdentityPostprocessFunctor>;
template class PopulateDistancesIterator<6, AddBiasFunctor>;
template class PopulateDistancesIterator<6, LimitedInnerFunctor>;

SCANN_INSTANTIATE_TYPED_CLASS(, AhImpl);

}  // namespace asymmetric_hashing_internal
}  // namespace research_scann


<|start_filename|>parallel_clustering/clustering/in-memory-clusterer.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef RESEARCH_GRAPH_IN_MEMORY_CLUSTERING_IN_MEMORY_CLUSTERER_H_
#define RESEARCH_GRAPH_IN_MEMORY_CLUSTERING_IN_MEMORY_CLUSTERER_H_

#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "clustering/config.pb.h"

namespace research_graph {
namespace in_memory {

// Interface of an in-memory clustering algorithm. The classes implementing this
// interface maintain a mutable graph, which can be clustered using a given set
// of parameters.
class InMemoryClusterer {
 public:
  // This is a basic interface for building graphs. Note that the interface only
  // specifies how to build a graph, as different clusterers may use different
  // interfaces for accessing it.
  // The node ids are consecutive, 0-based integers. In particular, adding a
  // node of id k to an empty graph creates k+1 nodes 0, ..., k.
  class Graph {
   public:
    using NodeId = int32_t;

    // Represents a weighted node with weighted outgoing edges.
    struct AdjacencyList {
      NodeId id = -1;
      double weight = 1;
      std::vector<std::pair<NodeId, double>> outgoing_edges;
    };

    virtual ~Graph() = default;

    // Adds a weighted node and its weighted out-edges to the graph. Depending
    // on the Graph implementation, the symmetric edges may be added as well,
    // and edge weights may be adjusted for symmetry.
    //
    // Import must be called at most once for each node. If not called for a
    // node, that node defaults to weight 1.
    //
    // IMPLEMENTATIONS MUST ALLOW CONCURRENT CALLS TO Import()!
    virtual absl::Status Import(AdjacencyList adjacency_list) = 0;

    virtual absl::Status FinishImport();
  };

  using NodeId = Graph::NodeId;
  using AdjacencyList = Graph::AdjacencyList;

  // Represents clustering: each element of the vector contains the set of
  // NodeIds in one cluster. We call a clustering non-overlapping if the
  // elements of the clustering are nonempty vectors that together contain each
  // NodeId exactly once.
  using Clustering = std::vector<std::vector<NodeId>>;

  virtual ~InMemoryClusterer() {}

  // Accessor to the maintained graph. Use it to build the graph.
  virtual Graph* MutableGraph() = 0;

  // Clusters the currently maintained graph using the given set of parameters.
  // Returns a clustering, or an error if the algorithm failed to cluster the
  // given graph.
  // Note that the same clustering may have multiple representations, and the
  // function may return any of them.
  virtual absl::StatusOr<Clustering> Cluster(
      const ClustererConfig& config) const = 0;

  // Same as above, except that it returns a sequence of clusterings. The last
  // element of the sequence is the final clustering. This is primarily used for
  // hierarchical clusterings, but callers should NOT assume that there is a
  // strict hierarchy structure (i.e. that clusters in clustering i are obtained
  // by merging clusters from clustering i-1). The default implementation
  // returns a single-element vector with the result of Cluster().
  virtual absl::StatusOr<std::vector<Clustering>> HierarchicalCluster(
      const ClustererConfig& config) const;

  // Refines a list of clusters and redirects the given pointer to new clusters.
  // This function is useful for methods that can refine / operate on an
  // existing clustering. It does not take ownership of clustering. The default
  // implementation does nothing and returns OkStatus.
  virtual absl::Status RefineClusters(const ClustererConfig& config,
                                      Clustering* clustering) const {
    return absl::OkStatus();
  }

  // Provides a pointer to a vector that contains string ids corresponding to
  // the NodeIds. If set, the ids from the provided map are used in the log and
  // error messages. The vector must live during all method calls of this
  // object. This call does *not* take ownership of the pointee. Using this
  // function is not required. If this function is never called, the ids are
  // converted to strings using absl::StrCat.
  void set_node_id_map(const std::vector<std::string>* node_id_map) {
    node_id_map_ = node_id_map;
  }

 protected:
  // Returns the string id corresponding to a given NodeId. If set_node_id_map
  // was called, uses the map to get the ids. Otherwise, returns the string
  // representation of the id.
  std::string StringId(NodeId id) const;

 private:
  // NodeId map set by set_node_id_map(). May be left to nullptr even after
  // initialization.
  const std::vector<std::string>* node_id_map_ = nullptr;
};

}  // namespace in_memory
}  // namespace research_graph

#endif  // RESEARCH_GRAPH_IN_MEMORY_CLUSTERING_IN_MEMORY_CLUSTERER_H_


<|start_filename|>scann/scann/utils/reduction.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_UTILS_REDUCTION_H_
#define SCANN_UTILS_REDUCTION_H_

#include <sys/types.h>

#include <cstdint>

#include "scann/data_format/datapoint.h"
#include "scann/oss_wrappers/scann_aligned_malloc.h"
#include "scann/utils/types.h"
#include "tensorflow/core/platform/prefetch.h"

namespace research_scann {

template <typename T, typename U, typename Reduce>
auto DensePairAccumulate(T* a, U* b, size_t size, Reduce reduce)
    -> AccumulatorTypeFor<T, U> {
  using AT = AccumulatorTypeFor<T, U>;
  AT result0 = 0;
  AT result1 = 0;
  AT result2 = 0;
  AT result3 = 0;

  auto enda = a + size;
  for (; a + 3 < enda; a += 4, b += 4) {
    reduce(&result0, a[0], b[0]);
    reduce(&result1, a[1], b[1]);
    reduce(&result2, a[2], b[2]);
    reduce(&result3, a[3], b[3]);
  }

  result2 += result3;
  if (a + 1 < enda) {
    reduce(&result0, a[0], b[0]);
    reduce(&result1, a[1], b[1]);
    a += 2;
    b += 2;
  }

  result1 += result2;
  if (a < enda) {
    reduce(&result0, a[0], b[0]);
  }

  return result0 + result1;
}

template <typename I1, typename V1, typename I2, typename V2,
          typename ReduceTwo>
auto SparsePairAccumulateImpl1(I1 indices1, V1 values1, size_t nonzero_entries1,
                               I2 indices2, V2 values2, size_t nonzero_entries2,
                               ReduceTwo reduce_two)
    -> AccumulatorTypeFor<decltype(values1[0]), decltype(values2[0])> {
  using AT = AccumulatorTypeFor<decltype(values1[0]), decltype(values2[0])>;

  if (nonzero_entries1 == 0 || nonzero_entries2 == 0) return 0;

  AT result = 0;

  size_t i1_front = 0, i2_front = 0;
  size_t i1_back = nonzero_entries1 - 1, i2_back = nonzero_entries2 - 1;

  while (i1_front < i1_back && i2_front < i2_back) {
    const size_t to_add_front1 = indices1[i1_front] <= indices2[i2_front];
    const size_t to_add_front2 = indices1[i1_front] >= indices2[i2_front];
    const size_t to_sub_back2 = indices1[i1_back] <= indices2[i2_back];
    const size_t to_sub_back1 = indices1[i1_back] >= indices2[i2_back];
    if (ABSL_PREDICT_FALSE(indices1[i1_front] == indices2[i2_front])) {
      reduce_two(&result, values1[i1_front], values2[i2_front]);
    }

    if (ABSL_PREDICT_FALSE(indices1[i1_back] == indices2[i2_back])) {
      reduce_two(&result, values1[i1_back], values2[i2_back]);
    }

    i1_front += to_add_front1;
    i2_front += to_add_front2;
    i1_back -= to_sub_back1;
    i2_back -= to_sub_back2;
  }

  if (i1_front == i1_back) {
    for (; i2_front <= i2_back; ++i2_front) {
      if (ABSL_PREDICT_FALSE(indices1[i1_front] == indices2[i2_front])) {
        reduce_two(&result, values1[i1_front], values2[i2_front]);
        break;
      }
    }
  } else if (i2_front == i2_back) {
    for (; i1_front <= i1_back; ++i1_front) {
      if (ABSL_PREDICT_FALSE(indices1[i1_front] == indices2[i2_front])) {
        reduce_two(&result, values1[i1_front], values2[i2_front]);
        break;
      }
    }
  }

  return result;
}

template <typename T>
inline T* PunInt(T* arg) {
  return arg;
}
inline int32_t* PunInt(float* arg) { return reinterpret_cast<int32_t*>(arg); }
inline int64_t* PunInt(double* arg) { return reinterpret_cast<int64_t*>(arg); }

template <typename I1, typename V1, typename I2, typename V2,
          typename ReduceTwo, typename ReduceOne>
auto SparsePairAccumulateImpl2(I1 indices1, V1 values1, size_t nonzero_entries1,
                               I2 indices2, V2 values2,
                               const size_t nonzero_entries2,
                               ReduceTwo reduce_two, ReduceOne reduce_one)
    -> AccumulatorTypeFor<decltype(values1[0]), decltype(values2[0])> {
  using AT = AccumulatorTypeFor<decltype(values1[0]), decltype(values2[0])>;
  AT result0 = 0, result1 = 0;

  ssize_t i1_front = 0, i2_front = 0;
  ssize_t i1_back = nonzero_entries1, i2_back = nonzero_entries2;
  --i1_back;
  --i2_back;

  while (i1_front < i1_back && i2_front < i2_back) {
    auto front_left = values1[i1_front];
    auto front_right = values2[i2_front];
    auto back_left = values1[i1_back];
    auto back_right = values2[i2_back];

    const size_t to_add_front1 = indices1[i1_front] <= indices2[i2_front];
    const size_t to_add_front2 = indices1[i1_front] >= indices2[i2_front];
    const size_t to_sub_back2 = indices1[i1_back] <= indices2[i2_back];
    const size_t to_sub_back1 = indices1[i1_back] >= indices2[i2_back];

    *PunInt(&front_left) &= -to_add_front1;
    *PunInt(&front_right) &= -to_add_front2;
    *PunInt(&back_left) &= -to_sub_back1;
    *PunInt(&back_right) &= -to_sub_back2;

    reduce_two(&result0, front_left, front_right);
    reduce_two(&result1, back_left, back_right);
    i1_front += to_add_front1;
    i2_front += to_add_front2;
    i1_back -= to_sub_back1;
    i2_back -= to_sub_back2;
  }

  while (i1_front <= i1_back && i2_front <= i2_back) {
    if (indices1[i1_front] == indices2[i2_front]) {
      reduce_two(&result0, values1[i1_front++], values2[i2_front++]);
    } else if (indices1[i1_front] < indices2[i2_front]) {
      reduce_one(&result0, values1[i1_front++]);
    } else {
      reduce_one(&result0, values2[i2_front++]);
    }
  }

  if (i1_front > i1_back) {
    for (; i2_front <= i2_back; ++i2_front) {
      reduce_one(&result0, values2[i2_front]);
    }
  } else {
    for (; i1_front <= i1_back; ++i1_front) {
      reduce_one(&result0, values1[i1_front]);
    }
  }

  return result0 + result1;
}

template <typename T, typename U, typename ReduceTwo, typename ReduceOne>
inline AccumulatorTypeFor<T, U> SparsePairAccumulate(const DatapointPtr<T>& a,
                                                     const DatapointPtr<U>& b,
                                                     ReduceTwo reduce_two,
                                                     ReduceOne reduce_one) {
  if (reduce_one.IsNoop()) {
    return SparsePairAccumulateImpl1(
        a.indices(), a.values(), a.nonzero_entries(), b.indices(), b.values(),
        b.nonzero_entries(), reduce_two);
  } else {
    return SparsePairAccumulateImpl2(
        a.indices(), a.values(), a.nonzero_entries(), b.indices(), b.values(),
        b.nonzero_entries(), reduce_two, reduce_one);
  }
}

template <typename I1, typename V1, typename I2, typename V2,
          typename ReduceTwo, typename ReduceOne>
auto SparsePairAccumulate(I1 indices1, V1 values1, size_t nonzero_entries1,
                          I2 indices2, V2 values2,
                          const size_t nonzero_entries2, ReduceTwo reduce_two,
                          ReduceOne reduce_one)
    -> AccumulatorTypeFor<decltype(values1[0]), decltype(values2[0])> {
  if (reduce_one.IsNoop()) {
    return SparsePairAccumulateImpl1(indices1, values1, nonzero_entries1,
                                     indices2, values2, nonzero_entries2,
                                     reduce_two);
  } else {
    return SparsePairAccumulateImpl2(indices1, values1, nonzero_entries1,
                                     indices2, values2, nonzero_entries2,
                                     reduce_two, reduce_one);
  }
}

template <typename T, typename U, typename ReduceTwo>
inline AccumulatorTypeFor<T, U> HybridPairAccumulateImpl1(
    const DatapointPtr<T>& a, const DatapointPtr<U>& b, ReduceTwo reduce_two) {
  DCHECK(a.IsSparse());
  DCHECK(b.IsDense());
  auto sparse_index_ptr = a.indices();
  auto sparse_values_ptr = a.values();
  const auto sparse_index_end = sparse_index_ptr + a.nonzero_entries();
  auto dense_values_ptr = b.values();
  using AT = AccumulatorTypeFor<T, U>;
  AT result0 = 0;
  AT result1 = 0;
  AT result2 = 0;
  AT result3 = 0;

  for (; sparse_index_ptr + 3 < sparse_index_end;
       sparse_index_ptr += 4, sparse_values_ptr += 4) {
    reduce_two(&result0, dense_values_ptr[sparse_index_ptr[0]],
               sparse_values_ptr[0]);
    reduce_two(&result1, dense_values_ptr[sparse_index_ptr[1]],
               sparse_values_ptr[1]);
    reduce_two(&result2, dense_values_ptr[sparse_index_ptr[2]],
               sparse_values_ptr[2]);
    reduce_two(&result3, dense_values_ptr[sparse_index_ptr[3]],
               sparse_values_ptr[3]);
  }

  result2 += result3;
  if (sparse_index_ptr + 1 < sparse_index_end) {
    reduce_two(&result0, dense_values_ptr[sparse_index_ptr[0]],
               sparse_values_ptr[0]);
    reduce_two(&result1, dense_values_ptr[sparse_index_ptr[1]],
               sparse_values_ptr[1]);
    sparse_index_ptr += 2;
    sparse_values_ptr += 2;
  }

  result1 += result2;
  if (sparse_index_ptr < sparse_index_end) {
    reduce_two(&result0, dense_values_ptr[sparse_index_ptr[0]],
               sparse_values_ptr[0]);
  }

  return result0 + result1;
}

template <typename T, typename U, typename ReduceTwo, typename ReduceOne>
AccumulatorTypeFor<T, U> HybridPairAccumulateImpl2(const DatapointPtr<T>& a,
                                                   const DatapointPtr<U>& b,
                                                   ReduceTwo reduce_two,
                                                   ReduceOne reduce_one) {
  DCHECK(a.IsSparse());
  DCHECK(b.IsDense());
  using AT = AccumulatorTypeFor<T, U>;
  AT result0 = 0;
  AT result1 = 0;
  AT result2 = 0;
  AT result3 = 0;

  auto cur_dense = b.values();
  auto end_dense = cur_dense + b.nonzero_entries();
  for (; cur_dense + 3 < end_dense; cur_dense += 4) {
    reduce_one(&result0, *cur_dense);
    reduce_one(&result1, cur_dense[1]);
    reduce_one(&result2, cur_dense[2]);
    reduce_one(&result3, cur_dense[3]);
  }

  if (cur_dense + 1 < end_dense) {
    reduce_one(&result0, *cur_dense);
    reduce_one(&result1, cur_dense[1]);
    cur_dense += 2;
  }

  if (cur_dense < end_dense) {
    reduce_one(&result0, *cur_dense);
  }

  AT throw_away0 = 0;
  AT throw_away1 = 0;
  AT throw_away2 = 0;
  AT throw_away3 = 0;

  auto sparse_index_ptr = a.indices();
  auto sparse_value_ptr = a.values();
  auto sparse_index_end = a.indices() + a.nonzero_entries();
  for (; sparse_index_ptr + 3 < sparse_index_end;
       sparse_index_ptr += 4, sparse_value_ptr += 4) {
    reduce_one(&throw_away0, b.values()[sparse_index_ptr[0]]);
    reduce_two(&result0, b.values()[sparse_index_ptr[0]], *sparse_value_ptr);
    reduce_one(&throw_away1, b.values()[sparse_index_ptr[1]]);
    reduce_two(&result1, b.values()[sparse_index_ptr[1]], sparse_value_ptr[1]);
    reduce_one(&throw_away2, b.values()[sparse_index_ptr[2]]);
    reduce_two(&result2, b.values()[sparse_index_ptr[2]], sparse_value_ptr[2]);
    reduce_one(&throw_away3, b.values()[sparse_index_ptr[3]]);
    reduce_two(&result3, b.values()[sparse_index_ptr[3]], sparse_value_ptr[3]);
  }

  if (sparse_index_ptr + 1 < sparse_index_end) {
    reduce_one(&throw_away0, b.values()[sparse_index_ptr[0]]);
    reduce_two(&result0, b.values()[sparse_index_ptr[0]], *sparse_value_ptr);
    reduce_one(&throw_away1, b.values()[sparse_index_ptr[1]]);
    reduce_two(&result1, b.values()[sparse_index_ptr[1]], sparse_value_ptr[1]);
    sparse_index_ptr += 2;
    sparse_value_ptr += 2;
  }

  if (sparse_index_ptr < sparse_index_end) {
    reduce_one(&throw_away0, b.values()[sparse_index_ptr[0]]);
    reduce_two(&result0, b.values()[sparse_index_ptr[0]], *sparse_value_ptr);
  }

  return ((result0 - throw_away0) + (result1 - throw_away1)) +
         ((result2 - throw_away2) + (result3 - throw_away3));
}

template <typename T, typename U, typename ReduceTwo, typename ReduceOne>
AccumulatorTypeFor<T, U> HybridPairAccumulate(const DatapointPtr<T>& a,
                                              const DatapointPtr<U>& b,
                                              ReduceTwo reduce_two,
                                              ReduceOne reduce_one) {
  DCHECK(a.IsSparse() != b.IsSparse());
  if (reduce_one.IsNoop()) {
    if (a.IsSparse()) {
      return HybridPairAccumulateImpl1(a, b, reduce_two);
    } else {
      return HybridPairAccumulateImpl1(b, a, reduce_two);
    }
  } else {
    if (a.IsSparse()) {
      return HybridPairAccumulateImpl2(a, b, reduce_two, reduce_one);
    } else {
      return HybridPairAccumulateImpl2(b, a, reduce_two, reduce_one);
    }
  }
}

template <typename T, typename Reduce>
AccumulatorTypeFor<T> DenseSingleAccumulate(ConstSpan<T> values,
                                            Reduce reduce) {
  using AT = AccumulatorTypeFor<T>;
  AT result0 = 0;
  AT result1 = 0;
  AT result2 = 0;
  AT result3 = 0;

  const T* ptr = values.data();
  const T* const end = values.data() + values.size();
  for (; ptr + 4 <= end; ptr += 4) {
    reduce(&result0, ptr[0]);
    reduce(&result1, ptr[1]);
    reduce(&result2, ptr[2]);
    reduce(&result3, ptr[3]);
  }

  result2 += result3;
  if (ptr + 2 <= end) {
    reduce(&result0, ptr[0]);
    reduce(&result1, ptr[1]);
    ptr += 2;
  }

  result1 += result2;
  if (ptr < end) {
    reduce(&result0, ptr[0]);
  }

  return result0 + result1;
}

struct DoNothingReduce {
  template <typename... T>
  void operator()(T... args) {}
  bool IsNoop() { return true; }
};

}  // namespace research_scann

#endif


<|start_filename|>zebraix/tools/zvx_svg.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include <iostream>
#include <vector>


#include "absl/flags/flag.h"
#include "absl/flags/parse.h"
#include "absl/strings/string_view.h"
#include "render/render_structs.h"
#include "render/zebraix_layout.h"
#include "render/zebraix_rendersvg.h"

using zebraix::render::BuildSampleLayout;
using zebraix::render::InboundDumpChoices;
using zebraix::render::ZebraixCairoSvg;
using zebraix::render::ZebraixLayout;

ABSL_FLAG(int, dump_inbound_graph, 0,
          "Dump debug string of graph, often in proto form");
ABSL_FLAG(
    bool, generate_sample_graph, false,
    "Generate a simple sample of drawing elements, and do not read input");
ABSL_FLAG(bool, draw_label_ticks, false,
          "Draw short lines from node centres to label anchor points");
ABSL_FLAG(bool, label_with_ranks, false,
          "Replace labels with (prime-rank,obverse-rank)");
ABSL_FLAG(bool, vanish_waypoints, false,
          "Make waypoints shrink to nothing, and no default arrows to them");

int main(int argc, char** argv) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);
  std::vector<char*> positional_args = absl::ParseCommandLine(argc, argv);

  if (positional_args.size() != 3) {
    std::cerr << "Usage:  " << argv[0] << " <input-file> <output-file>"
              << std::endl;
    return -1;
  }

  ZebraixLayout graph_layout;

  graph_layout.GetConfig().dump_inbound_graph =
      static_cast<InboundDumpChoices>(absl::GetFlag(FLAGS_dump_inbound_graph));
  graph_layout.GetConfig().draw_label_ticks =
      absl::GetFlag(FLAGS_draw_label_ticks);
  graph_layout.GetConfig().label_with_ranks =
      absl::GetFlag(FLAGS_label_with_ranks);
  graph_layout.GetConfig().vanish_waypoints =
      absl::GetFlag(FLAGS_vanish_waypoints);

  if (absl::GetFlag(FLAGS_generate_sample_graph)) {
    BuildSampleLayout(&graph_layout);
  } else {
    graph_layout.BuildLayoutFromFile(positional_args[1]);
  }

  const char* out_file = positional_args[2];

  {
    // Scoped to control destruction.
    ZebraixCairoSvg zcs(graph_layout.GetCanvas(), out_file);
    zcs.RenderToSvg(graph_layout);
    // ZebraixCairoSvg destructor performs important tear-down.
  }

  // Tear down required to avoid sanitizer errors over global caches.
  ZebraixCairoSvg::GlobalTearDown();
  return 0;
}


<|start_filename|>aptamers_mlpd/preprocess/scam_featurize.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Reads DNA sequences from the keys of an sstable and write another sstable
// with DNA sequence keys and values given by N-gram GenericFeatureVectors.
//
// Since the DNA alphabet has only four characters, and typical N-gram sizes are
// 4-6, we use a simple mathematical formula to map N-grams to sparse dimensions
// instead of hashing. We simply encode each character in the N-gram as two
// bits:
//
// A = 00
// C = 01
// G = 10
// T = 11
//
// Then, for example, the four-gram TCAG would be encoded as:
//
// (Leading zeroes) 11 01 00 10
//
// This allows up to 32-grams to be used, which should be more than sufficient.
#include "xxx/preprocess/scam_featurize.h"

#include <vector>

#include "absl/strings/string_view.h"
#include "xxx/scam/data_format/features.proto.h"
#include "xxx/scam/utils/gfv_normalization.h"

namespace research_biology {
namespace aptamers {

namespace {

std::vector<uint32> SequenceToNgrams(absl::string_view dna, size_t ngram_size) {
  std::vector<uint32> result;
  for (size_t i = 0; i + ngram_size <= dna.ssize(); ++i) {
    uint32 dim = 0;
    for (size_t j = 0; j < ngram_size; ++j) {
      const char c = dna[i + j];
      switch (c) {
        case 'A':
        case 'a':
          break;  // 0
        case 'C':
        case 'c':
          dim |= (1 << (2 * j));
          break;
        case 'G':
        case 'g':
          dim |= (2 << (2 * j));
          break;
        case 'T':
        case 't':
          dim |= (3 << (2 * j));
          break;
        default:
          LOG(FATAL) << "Non-DNA Character:  " << c;
      }
    }

    result.push_back(dim);
  }

  return result;
}

research_scam::GenericFeatureVector DnaToGfv(absl::string_view dna) {
  const size_t ngram_size = 6;
  research_scam::GenericFeatureVector gfv;
  std::vector<uint32> dims = SequenceToNgrams(dna, ngram_size);
  std::sort(dims.begin(), dims.end());
  gfv.set_feature_dim(std::pow(4, ngram_size));
  CHECK(!dims.empty()) << "No valid N-grams for sequence: " << dna;
  gfv.add_feature_index(dims[0]);
  gfv.add_feature_value_float(1);
  for (size_t i = 1; i < dims.size(); ++i) {
    if (dims[i] == dims[i - 1]) {
      ++(*gfv.mutable_feature_value_float()->rbegin());
    } else {
      gfv.add_feature_index(dims[i]);
      gfv.add_feature_value_float(1);
    }
  }

  gfv.set_feature_type(research_scam::GenericFeatureVector::FLOAT);
  return research_scam::NormalizeUnitL2(gfv);
}

class FeaturizeFn : public flume::MapFn<string, FeatureVectorEntry> {
 public:
  FeatureVectorEntry Map(const string& dna) override {
    return flume::make_kv(dna, DnaToGfv(dna));
  }

 private:
  REGISTER_AS_STATELESS_FN(FeaturizeFn);
};

}  // namespace

FeatureVectorTable FeaturizeSequences(const SequenceCollection& in) {
  return in.ParDo("featurize_dna", new FeaturizeFn());
}

}  // namespace aptamers
}  // namespace research_biology


<|start_filename|>scann/scann/oss_wrappers/scann_aligned_malloc.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/oss_wrappers/scann_aligned_malloc.h"

namespace research_scann {

#if (defined(__STD_C_VERSION__) && (__STD_C_VERSION__ >= 201112L)) || \
    (__cplusplus >= 201703L) || defined(_ISOC11_SOURCE)
void *aligned_malloc(size_t size, size_t minimum_alignment) {
  size = (size + minimum_alignment - 1) / minimum_alignment * minimum_alignment;
  return aligned_alloc(minimum_alignment, size);
}

void aligned_free(void *aligned_memory) { free(aligned_memory); }

#elif defined(_MSC_VER)
#include <malloc.h>

void *aligned_malloc(size_t size, size_t minimum_alignment) {
  return _aligned_alloc(size, minimum_alignment);
}

void aligned_free(void *aligned_memory) { _aligned_free(aligned_memory); }

#else
#endif

}  // namespace research_scann


<|start_filename|>scann/scann/utils/intrinsics/flags.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/utils/intrinsics/flags.h"

#include "tensorflow/core/platform/cpu_info.h"

ABSL_FLAG(bool, ignore_avx512, false,
          "Ignore the presence of AVX512 instructions when assigning "
          "function pointers at ScaNN startup.  Useful for testing and "
          "debugging.  NOTE:  AVX512 support is currently experimental and "
          "therefore disabled by default.");

ABSL_FLAG(bool, ignore_avx2, false,
          "Ignore the presence of AVX2/FMA instructions when assigning "
          "function pointers at ScaNN startup.  Useful for testing and "
          "debugging.");

ABSL_FLAG(bool, ignore_avx, false,
          "Ignore the presence of AVX and higher instructions when assigning "
          "function pointers at ScaNN startup.  Useful for testing and "
          "debugging.");

ABSL_RETIRED_FLAG(bool, ignore_sse4, false, "Ignore SSE4");

namespace research_scann {
namespace flags_internal {

bool should_use_sse4 =
    tensorflow::port::TestCPUFeature(tensorflow::port::SSE4_2);
bool should_use_avx1 = tensorflow::port::TestCPUFeature(tensorflow::port::AVX);
bool should_use_avx2 = tensorflow::port::TestCPUFeature(tensorflow::port::AVX2);
bool should_use_avx512 =
    tensorflow::port::TestCPUFeature(tensorflow::port::AVX512F) &&
    tensorflow::port::TestCPUFeature(tensorflow::port::AVX512DQ) &&
    tensorflow::port::TestCPUFeature(tensorflow::port::AVX512BW);

}  // namespace flags_internal

ScopedPlatformOverride::ScopedPlatformOverride(PlatformGeneration generation) {
  original_avx1_ = flags_internal::should_use_avx1;
  original_avx2_ = flags_internal::should_use_avx2;
  original_avx512_ = flags_internal::should_use_avx512;
  original_sse4_ = flags_internal::should_use_sse4;
  flags_internal::should_use_sse4 = false;
  flags_internal::should_use_avx1 = false;
  flags_internal::should_use_avx2 = false;
  flags_internal::should_use_avx512 = false;
  switch (generation) {
    case kSkylakeAvx512:
      flags_internal::should_use_avx512 = true;
      ABSL_FALLTHROUGH_INTENDED;

    case kHaswellAvx2:
      flags_internal::should_use_avx2 = true;
      ABSL_FALLTHROUGH_INTENDED;

    case kSandyBridgeAvx1:
      flags_internal::should_use_avx1 = true;
      ABSL_FALLTHROUGH_INTENDED;

    case kBaselineSse4:
      flags_internal::should_use_sse4 = true;
      break;

    case kFallbackForNonX86:
      break;

    default:
      LOG(FATAL) << "Unexpected Case: " << generation;
  }
}

ScopedPlatformOverride::~ScopedPlatformOverride() {
  flags_internal::should_use_avx1 = original_avx1_;
  flags_internal::should_use_avx2 = original_avx2_;
  flags_internal::should_use_avx512 = original_avx512_;
  flags_internal::should_use_sse4 = original_sse4_;
}

bool ScopedPlatformOverride::IsSupported() {
  if (flags_internal::should_use_avx512 &&
      !(tensorflow::port::TestCPUFeature(tensorflow::port::AVX512F) &&
        tensorflow::port::TestCPUFeature(tensorflow::port::AVX512DQ))) {
    LOG(WARNING) << "The CPU lacks AVX512 support! (skipping some tests)";
    return false;
  }
  if (flags_internal::should_use_avx2 &&
      !tensorflow::port::TestCPUFeature(tensorflow::port::AVX2)) {
    LOG(WARNING) << "The CPU lacks AVX2 support! (skipping some tests)";
    return false;
  }
  if (flags_internal::should_use_avx1 &&
      !tensorflow::port::TestCPUFeature(tensorflow::port::AVX)) {
    LOG(WARNING) << "The CPU lacks AVX1 support! (skipping some tests)";
    return false;
  }
  if (flags_internal::should_use_sse4 &&
      !tensorflow::port::TestCPUFeature(tensorflow::port::SSE4_2)) {
    LOG(WARNING) << "This CPU lacks SSE4.2 support! (skipping some tests)";
  }
  return true;
}

ScopedPlatformOverride TestHookOverridePlatform(PlatformGeneration generation) {
  return ScopedPlatformOverride(generation);
}

}  // namespace research_scann


<|start_filename|>zebraix/base/dominance.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "base/dominance.h"

#include <algorithm>
#include <limits>
#include <list>
#include <vector>

#include "glog/logging.h"

namespace zebraix {
namespace base {

void ScanRankRanges(DominanceGraph* d_graph) {
  d_graph->prime_max = std::numeric_limits<int>::min();
  d_graph->prime_min = std::numeric_limits<int>::max();
  d_graph->obverse_max = std::numeric_limits<int>::min();
  d_graph->obverse_min = std::numeric_limits<int>::max();
  bool has_source = false;
  bool has_sink = false;
  for (int i = 0; i < d_graph->d_nodes.size(); ++i) {
    bool new_prime_max = d_graph->d_nodes[i].prime_rank > d_graph->prime_max;
    bool new_prime_min = d_graph->d_nodes[i].prime_rank < d_graph->prime_min;
    bool new_obverse_max =
        d_graph->d_nodes[i].obverse_rank > d_graph->obverse_max;
    bool new_obverse_min =
        d_graph->d_nodes[i].obverse_rank < d_graph->obverse_min;

    if (new_prime_max) {
      d_graph->prime_max = d_graph->d_nodes[i].prime_rank;
    }
    if (new_prime_min) {
      d_graph->prime_min = d_graph->d_nodes[i].prime_rank;
    }
    if (new_obverse_max) {
      d_graph->obverse_max = d_graph->d_nodes[i].obverse_rank;
    }
    if (new_obverse_min) {
      d_graph->obverse_min = d_graph->d_nodes[i].obverse_rank;
    }

    if (new_prime_max || new_obverse_max) {
      has_sink = new_prime_max && new_obverse_max;
    }
    if (new_prime_min || new_obverse_min) {
      has_source = new_prime_min && new_obverse_min;
    }
  }

  d_graph->source_imputed = !has_source;
  d_graph->sink_imputed = !has_sink;

  if (d_graph->source_imputed) {
    --d_graph->prime_min;
    --d_graph->obverse_min;
  }
  if (d_graph->sink_imputed) {
    ++d_graph->prime_max;
    ++d_graph->obverse_max;
  }
  const int node_count = d_graph->prime_max - d_graph->prime_min + 1;
  CHECK_EQ(node_count, d_graph->obverse_max - d_graph->obverse_min + 1);
}

void FleshOutGraphNodes(DominanceGraph* d_graph) {
  if (d_graph->source_imputed) {
    DominanceNode dnode;
    dnode.import_index = kSourceImportPseudoIndex;
    dnode.prime_rank = d_graph->prime_min;
    dnode.obverse_rank = d_graph->obverse_min;
    d_graph->d_nodes.push_back(dnode);
  }
  if (d_graph->sink_imputed) {
    DominanceNode dnode;
    dnode.import_index = kSinkImportPseudoIndex;
    dnode.prime_rank = d_graph->prime_max;
    dnode.obverse_rank = d_graph->obverse_max;
    d_graph->d_nodes.push_back(dnode);
  }

  std::sort(d_graph->d_nodes.begin(), d_graph->d_nodes.end(), &RankCmp);
}

void ConnectGraph(DominanceGraph* d_graph) {
  const int node_count = d_graph->prime_max - d_graph->prime_min + 1;
  std::vector<std::list<int>> scratch_parents(node_count);
  std::vector<std::list<int>> scratch_children(node_count);

  for (int c = node_count - 1; c >= 0; --c) {
    const int child_obverse = d_graph->d_nodes[c].obverse_rank;
    int max_parent_obverse = -1;  // For each node, keep max parent rank.
    for (int p = c - 1; p >= 0; --p) {
      const int parent_obverse = d_graph->d_nodes[p].obverse_rank;
      if ((parent_obverse < child_obverse) &&
          (parent_obverse > max_parent_obverse)) {
        max_parent_obverse = parent_obverse;
        scratch_parents[c].push_back(p);
        scratch_children[p].push_back(c);
      }
    }
  }

  // Create efficient shrink-wrapped vector edge structures.
  for (int i = 0; i < node_count; ++i) {
    d_graph->d_nodes[i].parents = std::vector<int>(scratch_parents[i].cbegin(),
                                                   scratch_parents[i].cend());
    d_graph->d_nodes[i].parents.shrink_to_fit();
    d_graph->d_nodes[i].children = std::vector<int>(
        scratch_children[i].cbegin(), scratch_children[i].cend());
    d_graph->d_nodes[i].children.shrink_to_fit();
  }
}

}  // namespace base
}  // namespace zebraix


<|start_filename|>scann/scann/utils/internal/avx_funcs.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_UTILS_INTERNAL_AVX_FUNCS_H_
#define SCANN_UTILS_INTERNAL_AVX_FUNCS_H_
#ifdef __x86_64__

#include "scann/utils/intrinsics/avx1.h"
#include "scann/utils/types.h"

namespace research_scann {

class AvxFunctionsAvx {
 public:
  SCANN_AVX1_INLINE static __m256 Int8ToFloatLower(__m128i x) {
    __m128i avals0 = _mm_cvtepi8_epi32(x);
    __m128i avals1 = _mm_cvtepi8_epi32(_mm_srli_si128(x, 4));
    __m256i combined =
        _mm256_insertf128_si256(_mm256_castsi128_si256(avals0), avals1, 1);
    return _mm256_cvtepi32_ps(combined);
  }

  SCANN_AVX1_INLINE static __m256 Int8ToFloatUpper(__m128i x) {
    return Int8ToFloatLower(_mm_srli_si128(x, 8));
  }

  SCANN_AVX1_INLINE static __m256 MultiplyAdd(__m256 mul1, __m256 mul2,
                                              __m256 add) {
    return _mm256_add_ps(add, _mm256_mul_ps(mul1, mul2));
  }
  SCANN_AVX1_INLINE static __m256d MultiplyAdd(__m256d mul1, __m256d mul2,
                                               __m256d add) {
    return _mm256_add_pd(add, _mm256_mul_pd(mul1, mul2));
  }
  SCANN_AVX1_INLINE static __m128d MultiplyAdd(__m128d mul1, __m128d mul2,
                                               __m128d add) {
    return _mm_add_pd(add, _mm_mul_pd(mul1, mul2));
  }

  SCANN_AVX1_INLINE static __m256 MultiplySub(__m256 mul1, __m256 mul2,
                                              __m256 base) {
    return _mm256_sub_ps(base, _mm256_mul_ps(mul1, mul2));
  }

  SCANN_AVX1_INLINE static __m256 SseToAvx(__m128 x) {
    return _mm256_insertf128_ps(_mm256_setzero_ps(), x, 0);
  }

  SCANN_AVX1_INLINE static float Sum8(__m256 x) {
    const __m128 upper = _mm256_extractf128_ps(x, 1);
    const __m128 lower = _mm256_castps256_ps128(x);
    __m128 sum = _mm_add_ps(upper, lower);
    sum = _mm_add_ps(
        sum, _mm_castsi128_ps(_mm_srli_si128(_mm_castps_si128(sum), 8)));
    return sum[0] + sum[1];
  }
};

}  // namespace research_scann

#endif
#endif


<|start_filename|>zebraix/render/zebraix_rendersvg.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef ZEBRAIX_RENDER_ZEBRAIX_RENDERSVG_H_
#define ZEBRAIX_RENDER_ZEBRAIX_RENDERSVG_H_

#include "third_party/cairo/cairo.h"  // IWYU pragma: keep
#include "base/dominance.h"
#include "render/render_structs.h"
#include "render/zebraix_layout.h"

namespace zebraix {
namespace render {

class ZebraixCairoSvg {
 public:
  ZebraixCairoSvg(const ZebraixCanvas& overall, const char* out_file);
  ~ZebraixCairoSvg();

  void RenderToSvg(const ZebraixLayout& layout);

  static void GlobalTearDown();

 private:
  cairo_surface_t* surface;
  cairo_t* cr;
};

void BuildSampleLayout(ZebraixLayout* graph_layout);

}  // namespace render
}  // namespace zebraix

#endif  // ZEBRAIX_RENDER_ZEBRAIX_RENDERSVG_H_


<|start_filename|>aptamers_mlpd/preprocess/merge/merge_measurements_main.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// A FlumeC++ pipeline that merges measurements tables.

#include <gflags/gflags.h>

#include "pipeline/flume/public/flume.h"
#include "pipeline/flume/public/sstableio.h"
#include "xxx/preprocess/merge/merge_measurements.h"

DEFINE_string(output_filename,
              "",
              "output filename for sstable");

namespace research_biology {
namespace aptamers {

// Run the merge on the input files, specified by <argc, argv>.
void Run(int argc, char **argv) {
  // Initialize Flume now, and automatically clean up at the end of
  // this scope.
  flume::Flume flume;

  // Open the sstable inputs and construct the join table.
  std::vector<JoinTag> tags;
  flume::JoinOp<string> join("Merge");
  for (int i = 1; i < argc; i++) {
    const string key(argv[i]);
    MeasurementTable input =
        MeasurementTable::Read(
            key,
            flume::SSTableSource<string, Measurement>(
                key,
                flume::Strings(),
                flume::Protos<Measurement>()));

    JoinTag tag(key);
    tags.push_back(tag);
    join = join.With(flume::JoinArg::Of(tag, input));
  }
  JoinTable table = join.Join();

  // Transform the collection.
  MeasurementTable merged = MergeMeasurements(tags, table);

  // Write to the specified output file.
  flume::PSink<MeasurementEntry> output =
      flume::SSTableSink<string, Measurement>(
              FLAGS_output_filename,
              flume::Strings(),
              flume::Protos<Measurement>());

  // Write it out.
  merged.Write("Write", output);
}

}  // namespace aptamers
}  // namespace research_biology

int main(int argc, char** argv) {
  gflags::ParseCommandLineFlags(&argc, &argv, true);

  if (base::GetFlag(FLAGS_output_filename).empty()) {
    LOG(ERROR) << "--output_filename is required";
    return 1;
  }

  // Copy the file.
  research_biology::aptamers::Run(argc, argv);

  return 0;
}


<|start_filename|>scann/scann/utils/common.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_UTILS_COMMON_H_
#define SCANN_UTILS_COMMON_H_

#include <stddef.h>

#include <array>
#include <cstdint>
#include <memory>
#include <string>
#include <type_traits>
#include <vector>

#include "absl/base/attributes.h"
#include "absl/container/flat_hash_map.h"
#include "absl/container/flat_hash_set.h"
#include "absl/container/node_hash_map.h"
#include "absl/container/node_hash_set.h"
#include "absl/flags/flag.h"
#include "absl/memory/memory.h"
#include "absl/strings/match.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_replace.h"
#include "absl/strings/str_split.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "scann/oss_wrappers/scann_down_cast.h"
#include "scann/oss_wrappers/scann_serialize.h"
#include "scann/oss_wrappers/scann_status.h"
#include "tensorflow/core/lib/core/errors.h"
#include "tensorflow/core/lib/core/status.h"
#include "tensorflow/core/platform/logging.h"
#include "tensorflow/core/platform/prefetch.h"
#include "tensorflow/core/platform/types.h"

namespace research_scann {

using ::std::array;
using ::std::make_shared;
using ::std::numeric_limits;
using ::std::pair;
using ::std::shared_ptr;
using ::std::string;
using ::std::tuple;
using ::std::unique_ptr;
using ::std::vector;

using std::vector;

using ::absl::make_unique;

using ::absl::flat_hash_map;
using ::absl::flat_hash_set;
using ::absl::node_hash_map;
using ::absl::node_hash_set;

using ::absl::PrintF;
using ::absl::StrAppend;
using ::absl::StrAppendFormat;
using ::absl::StrCat;
using ::absl::StrFormat;
using ::absl::string_view;
using ::absl::StrJoin;
using ::absl::StrSplit;

using ::absl::GetFlag;

using ::absl::Mutex;
using ::absl::MutexLock;
using ::absl::ReaderMutexLock;

using ::tensorflow::Status;
using OkStatus = Status;
using internal::StatusOr;

#define MAKE_TF_ERROR_FORWARDER(ERRNAME)                                      \
  ABSL_MUST_USE_RESULT inline Status ERRNAME##Error(absl::string_view s) {    \
    return ::tensorflow::errors::ERRNAME(std::forward<absl::string_view>(s)); \
  }

MAKE_TF_ERROR_FORWARDER(Aborted);
MAKE_TF_ERROR_FORWARDER(AlreadyExists);
MAKE_TF_ERROR_FORWARDER(Cancelled);
MAKE_TF_ERROR_FORWARDER(FailedPrecondition);
MAKE_TF_ERROR_FORWARDER(Internal);
MAKE_TF_ERROR_FORWARDER(InvalidArgument);
MAKE_TF_ERROR_FORWARDER(NotFound);
MAKE_TF_ERROR_FORWARDER(OutOfRange);
MAKE_TF_ERROR_FORWARDER(Unauthenticated);
MAKE_TF_ERROR_FORWARDER(Unavailable);
MAKE_TF_ERROR_FORWARDER(Unimplemented);
MAKE_TF_ERROR_FORWARDER(Unknown);

#undef MAKE_TF_ERROR_FORWARDER

template <typename T>
using ConstSpan = absl::Span<const T>;

using ::absl::MakeConstSpan;

template <typename T>
using MutableSpan = absl::Span<T>;

template <typename... Args>
auto MakeMutableSpan(Args&&... args)
    -> decltype(absl::MakeSpan(std::forward<Args>(args)...)) {
  return absl::MakeSpan(std::forward<Args>(args)...);
}

enum : uint32_t {
  kInvalidIdx32 = std::numeric_limits<uint32_t>::max(),
};
enum : uint64_t {
  kInvalidIdx64 = std::numeric_limits<uint64_t>::max(),
};

#define SCANN_INLINE inline ABSL_ATTRIBUTE_ALWAYS_INLINE

#define SCANN_INLINE_LAMBDA ABSL_ATTRIBUTE_ALWAYS_INLINE

#define SCANN_OUTLINE ABSL_ATTRIBUTE_NOINLINE

const std::string& EmptyString();

template <typename CollectionT>
bool IsEmpty(const CollectionT& c) {
  return c.empty();
}

inline bool IsEmpty(const absl::Flag<std::string>& flag) {
  return absl::GetFlag(flag).empty();
}

template <typename CollectionT>
bool NotEmpty(const CollectionT& c) {
  return !c.empty();
}

inline bool NotEmpty(const absl::Flag<std::string>& flag) {
  return !absl::GetFlag(flag).empty();
}

template <typename StatusT>
inline bool IsOk(const StatusT& status) {
  return status.ok();
}

template <typename T>
inline bool IsOk(const StatusOr<T>& statusor) {
  return statusor.ok();
}

template <typename StatusT>
inline bool NotOk(const StatusT& status) {
  return !status.ok();
}

template <typename T>
inline bool NotOk(const StatusOr<T>& statusor) {
  return !statusor.ok();
}

template <typename... Args>
ABSL_MUST_USE_RESULT Status FailedPreconditionError(
    const absl::FormatSpec<Args...>& fmt, const Args&... args) {
  return FailedPreconditionError(StrFormat(fmt, args...));
}

template <typename... Args>
ABSL_MUST_USE_RESULT Status InternalError(const absl::FormatSpec<Args...>& fmt,
                                          const Args&... args) {
  return InternalError(StrFormat(fmt, args...));
}

template <typename... Args>
ABSL_MUST_USE_RESULT Status InvalidArgumentError(
    const absl::FormatSpec<Args...>& fmt, const Args&... args) {
  return InvalidArgumentError(StrFormat(fmt, args...));
}

template <typename... Args>
ABSL_MUST_USE_RESULT Status NotFoundError(const absl::FormatSpec<Args...>& fmt,
                                          const Args&... args) {
  return NotFoundError(StrFormat(fmt, args...));
}

template <typename... Args>
ABSL_MUST_USE_RESULT Status
OutOfRangeError(const absl::FormatSpec<Args...>& fmt, const Args&... args) {
  return OutOfRangeError(StrFormat(fmt, args...));
}

template <typename... Args>
ABSL_MUST_USE_RESULT Status
UnavailableError(const absl::FormatSpec<Args...>& fmt, const Args&... args) {
  return UnavailableError(StrFormat(fmt, args...));
}

template <typename... Args>
ABSL_MUST_USE_RESULT Status
UnimplementedError(const absl::FormatSpec<Args...>& fmt, const Args&... args) {
  return UnimplementedError(StrFormat(fmt, args...));
}

template <typename... Args>
ABSL_MUST_USE_RESULT Status UnknownError(const absl::FormatSpec<Args...>& fmt,
                                         const Args&... args) {
  return UnknownError(StrFormat(fmt, args...));
}

template <typename ContainerT>
void FreeBackingStorage(ContainerT* c) {
  DCHECK(c);
  *c = ContainerT();
}

template <typename Int, typename DenomInt>
constexpr Int DivRoundUp(Int num, DenomInt denom) {
  return (num + static_cast<Int>(denom) - static_cast<Int>(1)) /
         static_cast<Int>(denom);
}

template <typename Int, typename DenomInt>
constexpr Int NextMultipleOf(Int num, DenomInt denom) {
  return DivRoundUp(num, denom) * denom;
}

template <typename Int, typename DenomInt>
constexpr bool IsDivisibleBy(Int num, DenomInt denom) {
  return num % denom == 0;
}

template <typename T>
T ValueOrDie(StatusOr<T> statusor) {
  if (!statusor.ok()) {
    LOG(FATAL) << "VALUE_OR_DIE_FAILURE: " << statusor.status();
  }
  return std::move(statusor).ValueOrDie();
}

template <ssize_t kStride = 1>
class SeqWithStride {
 public:
  static constexpr size_t Stride() { return kStride; }

  SCANN_INLINE explicit SeqWithStride(size_t end) : begin_(0), end_(end) {}

  SCANN_INLINE SeqWithStride(size_t begin, size_t end)
      : begin_(begin), end_(end) {
    static_assert(kStride != 0);
    if (kStride > 0) {
      DCHECK_LE(begin, end);
      DCHECK_GE(static_cast<ssize_t>(begin), 0);
      DCHECK_GE(static_cast<ssize_t>(end), 0);
    } else {
      DCHECK_GE(begin, end);
      DCHECK_GE(static_cast<ssize_t>(begin + 1), 0);
      DCHECK_GE(static_cast<ssize_t>(end), 0);
    }
  }

  class SizeT {
   public:
    SCANN_INLINE SizeT(size_t val) : val_(val) {}
    SCANN_INLINE operator size_t() const { return val_; }

    SCANN_INLINE ~SizeT() {}

   private:
    size_t val_;
  };

  class Iterator {
   public:
    SCANN_INLINE explicit Iterator(size_t idx) : idx_(idx) {}

    SCANN_INLINE SizeT operator*() const {
      DCHECK_GE(static_cast<ssize_t>(idx_), 0);
      return idx_;
    }

    SCANN_INLINE Iterator& operator++() {
      idx_ += kStride;
      return *this;
    }

    SCANN_INLINE bool operator!=(Iterator end) const {
      if constexpr (kStride > 0) {
        return idx_ < end.idx_;
      }

      return static_cast<ssize_t>(idx_) >= static_cast<ssize_t>(end.idx_);
    }

   private:
    size_t idx_;
  };
  using iterator = Iterator;

  SCANN_INLINE Iterator begin() const { return Iterator(begin_); }
  SCANN_INLINE Iterator end() const { return Iterator(end_); }

 private:
  size_t begin_;
  size_t end_;
};

SCANN_INLINE auto Seq(size_t end) { return SeqWithStride<1>(0, end); }

SCANN_INLINE auto Seq(size_t begin, size_t end) {
  return SeqWithStride<1>(begin, end);
}

SCANN_INLINE auto ReverseSeq(size_t end) {
  return SeqWithStride<-1>(end - 1, 0);
}

SCANN_INLINE auto ReverseSeq(size_t begin, size_t end) {
  return SeqWithStride<-1>(end - 1, begin);
}

template <size_t kStride>
SCANN_INLINE auto ReverseSeqWithStride(size_t end) {
  end = (end - 1) / kStride * kStride;
  return SeqWithStride<-kStride>(end, 0);
}

template <size_t kStride>
SCANN_INLINE auto ReverseSeqWithStride(size_t begin, size_t end) {
  end = begin + (end - begin - 1) / kStride * kStride;
  return SeqWithStride<-kStride>(end, begin);
}

template <typename Container>
SeqWithStride<1> IndicesOf(const Container& container) {
  return Seq(container.size());
}

template <typename T, typename IdxType = size_t,
          typename TIter = decltype(std::begin(std::declval<T>())),
          typename = decltype(std::end(std::declval<T>()))>
constexpr auto Enumerate(T&& iterable) {
  class IteratorWithIndex {
   public:
    IteratorWithIndex(IdxType idx, TIter it) : idx_(idx), it_(it) {}
    bool operator!=(const IteratorWithIndex& other) const {
      return it_ != other.it_;
    }
    void operator++() { idx_++, it_++; }
    auto operator*() const { return std::tie(idx_, *it_); }

   private:
    IdxType idx_;
    TIter it_;
  };
  struct iterator_wrapper {
    T iterable;
    auto begin() { return IteratorWithIndex{0, std::begin(iterable)}; }
    auto end() { return IteratorWithIndex{0, std::end(iterable)}; }
  };
  return iterator_wrapper{std::forward<T>(iterable)};
}

template <typename FloatT>
Status VerifyAllFiniteImpl(ConstSpan<FloatT> span) {
  for (size_t j : Seq(span.size())) {
    if (!ABSL_PREDICT_TRUE(std::isfinite(span[j]))) {
      return InternalError("Element not finite (dim idx = %d, value = %f)", j,
                           span[j]);
    }
  }
  return OkStatus();
}

inline Status VerifyAllFinite(ConstSpan<float> span) {
  return VerifyAllFiniteImpl<float>(span);
}

inline Status VerifyAllFinite(ConstSpan<double> span) {
  return VerifyAllFiniteImpl<double>(span);
}

struct VirtualDestructor {
  virtual ~VirtualDestructor() {}
};

struct MoveOnly {
  MoveOnly() = default;

  MoveOnly(MoveOnly&&) = default;
  MoveOnly& operator=(MoveOnly&&) = default;

  MoveOnly(const MoveOnly&) = delete;
  MoveOnly& operator=(const MoveOnly&) = delete;
};

#define SCANN_DECLARE_COPYABLE_CLASS(ClassName) \
  ClassName(ClassName&&) = default;             \
  ClassName& operator=(ClassName&&) = default;  \
  ClassName(const ClassName&) = default;        \
  ClassName& operator=(const ClassName&) = default

#define SCANN_DECLARE_MOVE_ONLY_CLASS(ClassName) \
  ClassName(ClassName&&) = default;              \
  ClassName& operator=(ClassName&&) = default;   \
  ClassName(const ClassName&) = delete;          \
  ClassName& operator=(const ClassName&) = delete

#define SCANN_DECLARE_MOVE_ONLY_CLASS_CUSTOM_IMPL(ClassName) \
  ClassName(const ClassName&) = delete;                      \
  ClassName& operator=(const ClassName&) = delete

#define SCANN_DECLARE_IMMOBILE_CLASS(ClassName) \
  ClassName(ClassName&&) = delete;              \
  ClassName& operator=(ClassName&&) = delete;   \
  ClassName(const ClassName&) = delete;         \
  ClassName& operator=(const ClassName&) = delete

using ::absl::conditional_t;
using ::absl::decay_t;
using ::absl::enable_if_t;
using ::absl::make_signed_t;
using ::absl::make_unsigned_t;
using ::absl::remove_cv_t;
using ::absl::remove_pointer_t;
using ::std::declval;

template <typename T, typename U>
inline constexpr bool IsSame() {
  return std::is_same<T, U>::value;
}

namespace is_same_any_impl {

template <typename T, typename... UU>
struct IsSameAny;

template <typename T>
struct IsSameAny<T> : std::false_type {};

template <typename T, typename U, typename... UU>
struct IsSameAny<T, U, UU...> {
  static constexpr bool value = IsSame<T, U>() || IsSameAny<T, UU...>::value;
};

}  // namespace is_same_any_impl

template <typename T, typename... UU>
inline constexpr bool IsSameAny() {
  return is_same_any_impl::IsSameAny<T, UU...>::value;
}

template <typename T>
inline constexpr bool IsUint8() {
  return IsSame<T, uint8_t>();
}

template <typename T>
inline constexpr bool IsFloat() {
  return IsSame<T, float>();
}

template <typename T>
inline constexpr bool IsDouble() {
  return IsSame<T, double>();
}

template <typename T>
inline constexpr bool IsFixedPointType() {
  return IsSameAny<T, int8_t, uint8_t, int16_t, uint16_t>();
}

template <typename T>
inline constexpr bool IsFloatingType() {
  return std::is_floating_point<decay_t<T>>::value;
}

template <typename T>
inline constexpr bool IsIntegerType() {
  return std::is_integral<T>::value;
}

template <typename T>
inline constexpr bool IsSignedType() {
  return std::is_signed<T>::value;
}

template <typename T>
constexpr bool IsStatusType() {
  return std::is_convertible<Status, T>::value;
}

template <typename T>
inline constexpr bool IsPod() {
  return std::is_pod<decay_t<T>>::value;
}

template <typename T>
using ConstRefOrPod = conditional_t<IsPod<T>(), decay_t<T>, const decay_t<T>&>;

template <typename T>
struct RecursivelyRemoveCVImpl {
  using type = remove_cv_t<T>;
};

template <typename T>
using RecursivelyRemoveCV = typename RecursivelyRemoveCVImpl<T>::type;

template <typename K, typename V>
struct RecursivelyRemoveCVImpl<pair<K, V>> {
  using type = pair<RecursivelyRemoveCV<K>, RecursivelyRemoveCV<V>>;
};

template <typename... Ts>
struct RecursivelyRemoveCVImpl<tuple<Ts...>> {
  using type = tuple<RecursivelyRemoveCV<Ts>...>;
};

static_assert(IsSame<pair<const int, volatile float>,
                     remove_cv_t<pair<const int, volatile float>>>(),
              "");
static_assert(IsSame<pair<int, float>,
                     RecursivelyRemoveCV<pair<const int, volatile float>>>(),
              "");

}  // namespace research_scann

#endif


<|start_filename|>fairness_and_bias_in_online_selection/secretary_eval.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "secretary_eval.h"

#include <iostream>

#include "utils.h"

namespace fair_secretary {

using std::cout;
using std::endl;
using std::vector;

void SecretaryEval::Eval(const vector<SecretaryInstance>& instance,
                         const vector<SecretaryInstance>& answer,
                         int num_colors) {
  vector<double> max_value(num_colors, 0);
  vector<int> correct_answer(num_colors, 0);
  vector<int> num_answer(num_colors, 0);
  double total_max = 0;
  int not_picked = 0;
  int total_correct_answer = 0;
  for (int i = 0; i < instance.size(); i++) {
    max_value[instance[i].color] =
        std::max(max_value[instance[i].color], instance[i].value);
  }
  for (int i = 0; i < num_colors; i++) {
    total_max = std::max(total_max, max_value[i]);
  }
  cout << "Color Distribution: " << std::endl;
  vector<int> color(num_colors, 0);
  for (int i = 0; i < instance.size(); i++) {
    color[instance[i].color]++;
  }
  for (int i = 0; i < num_colors; i++) {
    cout << color[i] << " ";
  }
  cout << endl;
  for (const auto& element : answer) {
    if (element.color == -1) {
      not_picked++;
      continue;
    }
    num_answer[element.color]++;
    if (max_value[element.color] - element.value < 0.0000001) {
      correct_answer[element.color]++;
    }
    if (total_max - element.value < 0.0000001) {
      total_correct_answer++;
    }
  }
  cout << "Answer Distribution:" << endl;
  for (int i = 0; i < num_colors; i++) {
    cout << num_answer[i] << " ";
  }
  cout << endl;
  cout << "Correct Answer Distribution:" << endl;
  for (int i = 0; i < num_colors; i++) {
    cout << correct_answer[i] << " ";
  }
  cout << endl;
  cout << "Total Correct Answer: " << total_correct_answer << endl;
  cout << "Probability Correct Answer: "
       << static_cast<double>(total_correct_answer) / answer.size() << endl;
  cout << "Total Not Picked: " << not_picked << endl;
}

void SecretaryEval::ThEval(const vector<SecretaryInstance>& instance,
                           const vector<vector<SecretaryInstance>>& answers,
                           int num_colors) {
  double total_max = 0;
  int not_picked = 0;
  vector<double> total_correct_answer(answers.size(), 0);
  for (int i = 0; i < instance.size(); i++) {
    total_max = std::max(total_max, instance[i].value);
  }
  for (int i = 0; i < answers.size(); i++) {
    for (const auto& element : answers[i]) {
      if (element.color == -1) {
        not_picked++;
        continue;
      }
      if (total_max - element.value < 0.0000001) {
        total_correct_answer[i]++;
      }
    }
    cout << static_cast<double>(total_correct_answer[i]) / answers[i].size()
         << " ";
  }
  cout << endl;
}

void SecretaryEval::InnerUnbalanced(
    const vector<SecretaryInstance>& instance, const SecretaryInstance& ans,
    vector<int>& correct_answer, vector<int>& num_answer, vector<int>& max_dist,
    const int num_colors, int& not_picked, int& total_correct_answer) {
  vector<double> max_value(num_colors, 0);
  double total_max = 0;

  int max_color = 0;
  for (const auto& element : instance) {
    max_value[element.color] =
        std::max(max_value[element.color], element.value);
    double old_max = total_max;
    total_max = std::max(total_max, element.value);
    if (old_max < total_max) {
      max_color = element.color;
    }
  }
  max_dist[max_color]++;
  if (ans.color == -1) {
    not_picked++;
    return;
  }
  num_answer[ans.color]++;
  if (max_value[ans.color] - ans.value < 10) {
    correct_answer[ans.color]++;
  }
  if (total_max - ans.value < 10) {
    total_correct_answer++;
  }
}

}  // namespace fair_secretary


<|start_filename|>scann/scann/oss_wrappers/scann_status.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/oss_wrappers/scann_status.h"

#include "absl/strings/str_cat.h"

namespace research_scann {

Status AnnotateStatus(const Status& s, absl::string_view msg) {
  if (s.ok() || msg.empty()) return s;

  absl::string_view new_msg = msg;
  std::string annotated;
  if (!s.error_message().empty()) {
    absl::StrAppend(&annotated, s.error_message(), "; ", msg);
    new_msg = annotated;
  }
  return Status(s.code(), new_msg);
}

StatusBuilder RetCheckFail(absl::string_view msg) {
  return InternalErrorBuilder() << msg;
}

}  // namespace research_scann


<|start_filename|>fairness_and_bias_in_online_selection/secretary_eval.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef FAIRNESS_AND_BIAS_SELECTION_SECRETARY_EVAL_H_
#define FAIRNESS_AND_BIAS_SELECTION_SECRETARY_EVAL_H_

#include <vector>

#include "utils.h"

namespace fair_secretary {

// Evaluates the quality of the solution, and provide stats to the STDOUT.
class SecretaryEval {
 public:
  static void Eval(const std::vector<SecretaryInstance>& instance,
                   const std::vector<SecretaryInstance>& answer, int num_color);
  static void ThEval(const std::vector<SecretaryInstance>& instance,
                     const std::vector<std::vector<SecretaryInstance>>& answers,
                     int num_colors);
  static void InnerUnbalanced(const std::vector<SecretaryInstance>& instance,
                              const SecretaryInstance& ans,
                              std::vector<int>& correct_answer,
                              std::vector<int>& num_answer,
                              std::vector<int>& max_dist, const int num_colors,
                              int& not_picked, int& total_correct_answer);
};

}  // namespace fair_secretary

#endif  // FAIRNESS_AND_BIAS_SELECTION_SECRETARY_EVAL_H_


<|start_filename|>scann/scann/data_format/docid_collection_interface.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_DATA_FORMAT_DOCID_COLLECTION_INTERFACE_H_
#define SCANN_DATA_FORMAT_DOCID_COLLECTION_INTERFACE_H_

#include "scann/utils/types.h"

namespace research_scann {

class DocidCollectionInterface {
 public:
  virtual ~DocidCollectionInterface() {}

  virtual Status Append(string_view docid) = 0;

  virtual size_t size() const = 0;

  virtual bool empty() const = 0;

  virtual string_view Get(size_t i) const = 0;

  virtual size_t capacity() const = 0;

  virtual size_t MemoryUsage() const = 0;

  virtual void Clear() = 0;

  virtual void Reserve(DatapointIndex n_elements) = 0;

  virtual void ShrinkToFit() = 0;

  virtual unique_ptr<DocidCollectionInterface> Copy() const = 0;

  class Mutator {
   public:
    virtual ~Mutator() {}

    virtual Status AddDatapoint(string_view docid) = 0;

    virtual Status RemoveDatapoint(string_view docid) = 0;

    virtual bool LookupDatapointIndex(string_view docid,
                                      DatapointIndex* idx) const = 0;

    virtual void Reserve(size_t size) = 0;

    virtual Status RemoveDatapoint(DatapointIndex idx) = 0;
  };

  virtual StatusOr<Mutator*> GetMutator() const = 0;
};

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/metadata/metadata_getter.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_METADATA_METADATA_GETTER_H_
#define SCANN_METADATA_METADATA_GETTER_H_

#include "absl/synchronization/mutex.h"
#include "scann/data_format/datapoint.h"
#include "scann/data_format/dataset.h"
#include "scann/data_format/features.pb.h"
#include "scann/utils/types.h"

namespace research_scann {

template <typename T>
class MetadataGetter;

class UntypedMetadataGetter {
 public:
  virtual Status AppendMetadata(const GenericFeatureVector& gfv);

  virtual Status UpdateMetadata(DatapointIndex idx,
                                const GenericFeatureVector& gfv);

  virtual Status RemoveMetadata(DatapointIndex removed_idx);

  virtual bool needs_dataset() const;

  virtual research_scann::TypeTag TypeTag() const = 0;

  virtual ~UntypedMetadataGetter();

 private:
  absl::Mutex mutex_;
};

template <typename T>
class MetadataGetter : public UntypedMetadataGetter {
 public:
  MetadataGetter() {}

  research_scann::TypeTag TypeTag() const final { return TagForType<T>(); }

  virtual Status GetMetadata(const TypedDataset<T>* dataset,
                             const DatapointPtr<T>& query,
                             DatapointIndex neighbor_index,
                             std::string* result) const = 0;

  virtual StatusOr<std::string> GetByDatapointIndex(
      const TypedDataset<T>* dataset, DatapointIndex dp_idx) const {
    return UnimplementedError(
        StrCat("Cannot get metadata by datapoint index for "
               "metadata getter type ",
               typeid(*this).name(), "."));
  }

 private:
  TF_DISALLOW_COPY_AND_ASSIGN(MetadataGetter);
};

}  // namespace research_scann

#endif


<|start_filename|>sliding_window_clustering/sliding_window_clustering_main.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Main program to run the experiments of the paper.
// The program reads a file representing the stream of points as a sequence of
// rows one point per row. All rows contain d+1 fields separated by tab '\t'
// where the first field is the timestamp of the point (which must start from 0
// and be a consecutive integer). The remaining fields are the dimensions of the
// point.
#include <cstring>
#include <fstream>
#include <iostream>
#include <vector>

using std::string;
using std::vector;

#include "absl/flags/commandlineflag.h"
#include "absl/flags/flag.h"
#include "absl/flags/parse.h"
#include "absl/random/random.h"
#include "absl/strings/str_cat.h"
#include "base.h"
#include "estimate_optimum_range.h"
#include "io.h"
#include "meyerson_algorithm.h"
#include "sliding_window_framework.h"

ABSL_FLAG(std::string, input_file, "",
          "Path to the input file containing the stream of the points. The "
          "file format is a text file where each line is a point (in order of "
          "the stream). Each line contains d floating point numbers seprated "
          "by tab character defining the d dimensions of the point.");
ABSL_FLAG(
    std::string, output_file, "",
    "Path to the output text file with statistics on the experiment. The "
    "file will contain one line for every solution computed (i.e., one "
    "every output_every_n points). Each line contains in order the following "
    "tab-separated numerical fields: position in the stream; number of dist "
    "calls for one update, number of items stored, cost of the solution of our "
    "algorithm; cost of the baseline algorithm storing the entire window; cost "
    "of the baseline algorithm storing the same number of points of our "
    "algorithm.");
ABSL_FLAG(
    std::string, output_centers_file, "",
    "(Optional) Path to the output text file with the centers computed by our "
    "algorithm. If the path is non-empty, the file will contain one line for "
    "every solution computed (i.e., one every output_every_n points). Each line"
    " contains at most k comma-separated vectors. Each vector is a sequence of "
    "d tab-separated numbers.");
ABSL_FLAG(int32_t, window_size, 10000, "The size of the sliding window.");
ABSL_FLAG(int32_t, output_every_n, 100, "print solution every n steps.");
ABSL_FLAG(double, delta_grid, 0.2,
          "The delta parameter used int the grid for guessing the optimum.");
ABSL_FLAG(int32_t, k, 10, "the number of centers.");

namespace sliding_window {

// Define the global counters.
// global  count of the *current* number of points stored in any summary.
int64_t ITEMS_STORED;
// global count of all calls to the distance function.
int64_t CALLS_DIST_FN;

const char SEPARATOR[] = "\t";

void Main() {
  CHECK_GT(absl::GetFlag(FLAGS_window_size), 0);
  CHECK(!absl::GetFlag(FLAGS_input_file).empty());
  CHECK(!absl::GetFlag(FLAGS_output_file).empty());
  CHECK_GT(absl::GetFlag(FLAGS_k), 0);
  CHECK_GT(absl::GetFlag(FLAGS_delta_grid), 0);
  CHECK_LT(absl::GetFlag(FLAGS_k), absl::GetFlag(FLAGS_window_size));

  // Initialize the globals counters.
  ITEMS_STORED = 0;
  CALLS_DIST_FN = 0;

  absl::BitGen shared_genetor;

  const auto& [lower_bound, upper_bound] = guess_optimum_range_bounds(
      absl::GetFlag(FLAGS_input_file), absl::GetFlag(FLAGS_window_size),
      /*num_samples=*/100, absl::GetFlag(FLAGS_k), &shared_genetor);
  CHECK_GT(lower_bound, 0);
  CHECK_GT(upper_bound, 0);

  std::cout << "Estimated range for opt: [" << lower_bound << ", "
            << upper_bound << "]" << std::endl;

  std::ofstream output_file(absl::GetFlag(FLAGS_output_file), std::ios::out);

  absl::optional<std::ofstream> output_centers_file;
  if (!absl::GetFlag(FLAGS_output_centers_file).empty()) {
    output_centers_file.emplace(absl::GetFlag(FLAGS_output_centers_file),
                                std::ios::out);
  }

  // Solution found
  std::vector<TimePointPair> centers;

  // Our algorithm.
  FrameworkAlg<KMeansSummary> framework(
      absl::GetFlag(FLAGS_window_size), absl::GetFlag(FLAGS_k),
      absl::GetFlag(FLAGS_delta_grid), lower_bound, upper_bound,
      &shared_genetor);

  // The sampling based baseline.
  BaselineKmeansOverSampleWindow baseline(
      absl::GetFlag(FLAGS_window_size), absl::GetFlag(FLAGS_k),
      /* number_tries=*/10, &shared_genetor);

  std::string line;
  std::ifstream input_file(absl::GetFlag(FLAGS_input_file), std::ios::in);
  int32_t lines_read = 0;
  CHECK(input_file.is_open());
  while (std::getline(input_file, line)) {
    const auto& point = parse_point_from_string(line);

    // Execute the main algorithm.
    // Stores the number of calls before execting the algorithm.
    int64_t num_calls_dist_before = CALLS_DIST_FN;
    framework.process_point(lines_read, point);
    // New distance calls.
    int64_t num_calls_dist = CALLS_DIST_FN - num_calls_dist_before;

    // Executes the baseline.
    baseline.process_point(lines_read, point);

    // The current solution is evalauted only every FLAGS_print_every_n steps
    // for efficiency reasons (computing the correct cost of the solution
    // requires O(W) operations).
    if (lines_read >= absl::GetFlag(FLAGS_window_size) &&
        lines_read % absl::GetFlag(FLAGS_output_every_n) == 0) {
      // Estimate of the cost of the solution.
      double cost_estimate = 0;
      framework.solution(&centers, &cost_estimate);

      // If the path is passed, output the solution.
      if (output_centers_file.has_value()) {
        CHECK(output_centers_file.value().is_open());
        CHECK(!centers.empty());
        string centers_output;
        for (int i = 0; i < centers.size(); i++) {
          for (int j = 0; j < centers[i].second.size(); j++) {
            absl::StrAppend(&centers_output, centers[i].second[j]);
            if (j < centers[i].second.size() - 1) {
              absl::StrAppend(&centers_output, "\t");
            }
          }
          if (i < centers.size() - 1) {
            absl::StrAppend(&centers_output, ",");
          }
        }
        output_centers_file.value() << centers_output << "\n";
      }

      // Get the actual points in the window.
      const auto& points_window = baseline.points_window();

      // Actual cost of the solution for our algorithm.
      double cost = cost_solution(points_window, centers);
      CHECK_LE(centers.size(), absl::GetFlag(FLAGS_k));

      std::vector<TimePointPair> unused_centers;
      // Cost of the solution of the baseline algorithm using the whole window.
      double baseline_cost = 0;
      baseline.solution(&unused_centers, &baseline_cost);

      // Cost of the solution of the baseline algorithm using the same space of
      // our algorithm.
      double cost_baseline_subsampling;
      baseline.solution_subsampling(ITEMS_STORED, &unused_centers,
                                    &cost_baseline_subsampling);

      string output_str =
          absl::StrCat(lines_read, SEPARATOR, num_calls_dist, SEPARATOR,
                       ITEMS_STORED, SEPARATOR, cost, SEPARATOR, baseline_cost,
                       SEPARATOR, cost_baseline_subsampling, "\n");

      output_file << output_str.c_str();
    }

    if (lines_read % 500 == 0) {
      std::cout << "Points processed: " << lines_read << std::endl;
    }

    ++lines_read;
  }

  output_file.close();
  if (output_centers_file.has_value()) {
    output_centers_file.value().close();
  }
}

}  // namespace sliding_window

int main(int argc, char* argv[]) {
  absl::ParseCommandLine(argc, argv);
  sliding_window::Main();
  return 0;
}


<|start_filename|>scann/scann/distance_measures/one_to_one/nonzero_intersect_distance.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_DISTANCE_MEASURES_ONE_TO_ONE_NONZERO_INTERSECT_DISTANCE_H_
#define SCANN_DISTANCE_MEASURES_ONE_TO_ONE_NONZERO_INTERSECT_DISTANCE_H_

#include "scann/data_format/datapoint.h"
#include "scann/distance_measures/distance_measure_base.h"
#include "scann/distance_measures/one_to_one/common.h"
#include "scann/distance_measures/one_to_one/dot_product.h"
#include "scann/utils/reduction.h"
#include "tensorflow/core/platform/logging.h"

namespace research_scann {

class NonzeroIntersectDistance final : public DistanceMeasure {
 public:
  SCANN_DECLARE_DISTANCE_MEASURE_VIRTUAL_METHODS(NOT_SPECIALLY_OPTIMIZED);

 private:
  template <typename T>
  SCANN_INLINE double GetDistanceDenseImpl(const DatapointPtr<T>& a,
                                           const DatapointPtr<T>& b) const {
    DCHECK_EQ(a.dimensionality(), b.dimensionality());
    DimensionIndex result = 0;
    for (DimensionIndex i : Seq(a.nonzero_entries())) {
      result += a.values()[i] && b.values()[i];
    }
    return -static_cast<double>(result);
  }

  template <typename T>
  SCANN_INLINE double GetDistanceSparseImpl(const DatapointPtr<T>& a,
                                            const DatapointPtr<T>& b) const {
    struct Sub1Reduce {
      void operator()(AccumulatorTypeFor<T>* acc, const T a, const T b) {
        *acc += (a && b);
      }
    };
    return -static_cast<double>(
        SparsePairAccumulate(a, b, Sub1Reduce(), DoNothingReduce()));
  }

  template <typename T>
  SCANN_INLINE double GetDistanceHybridImpl(const DatapointPtr<T>& a,
                                            const DatapointPtr<T>& b) const {
    DCHECK(a.IsSparse() != b.IsSparse());
    if (b.IsSparse()) return GetDistanceHybridImpl(b, a);
    DCHECK(a.IsSparse());
    DCHECK(b.IsDense());
    DimensionIndex result = 0;
    for (DimensionIndex i : Seq(a.nonzero_entries())) {
      const DimensionIndex dim = a.indices()[i];
      result += (a.values()[i] && b.values()[dim]);
    }
    return -static_cast<double>(result);
  }
};

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/distance_measures/one_to_one/jaccard_distance.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#ifndef SCANN_DISTANCE_MEASURES_ONE_TO_ONE_JACCARD_DISTANCE_H_
#define SCANN_DISTANCE_MEASURES_ONE_TO_ONE_JACCARD_DISTANCE_H_

#include <cstdint>

#include "scann/distance_measures/distance_measure_base.h"
#include "scann/distance_measures/one_to_one/binary_distance_measure_base.h"
#include "scann/distance_measures/one_to_one/common.h"
#include "scann/utils/reduction.h"
#include "scann/utils/types.h"

namespace research_scann {

class GeneralJaccardDistance final : public DistanceMeasure {
 public:
  SCANN_DECLARE_DISTANCE_MEASURE_VIRTUAL_METHODS(NOT_SPECIALLY_OPTIMIZED);

 private:
  template <typename T>
  SCANN_INLINE double GetDistanceDenseImpl(const DatapointPtr<T>& a,
                                           const DatapointPtr<T>& b) const;

  template <typename T>
  SCANN_INLINE double GetDistanceSparseImpl(const DatapointPtr<T>& a,
                                            const DatapointPtr<T>& b) const;

  template <typename T>
  SCANN_INLINE double GetDistanceHybridImpl(const DatapointPtr<T>& a,
                                            const DatapointPtr<T>& b) const {
    LOG(FATAL) << "Not implemented yet.";
  }
};

class BinaryJaccardDistance final : public BinaryDistanceMeasureBase {
 public:
  string_view name() const final;

  using BinaryDistanceMeasureBase::GetDistanceDense;
  using BinaryDistanceMeasureBase::GetDistanceHybrid;
  using BinaryDistanceMeasureBase::GetDistanceSparse;

  double GetDistanceDense(const DatapointPtr<uint8_t>& a,
                          const DatapointPtr<uint8_t>& b) const final;
  double GetDistanceSparse(const DatapointPtr<uint8_t>& a,
                           const DatapointPtr<uint8_t>& b) const final;
  double GetDistanceHybrid(const DatapointPtr<uint8_t>& a,
                           const DatapointPtr<uint8_t>& b) const final {
    return GetDistanceSparse(a, b);
  }
};

template <typename T>
double GeneralJaccardDistance::GetDistanceDenseImpl(
    const DatapointPtr<T>& a, const DatapointPtr<T>& b) const {
  double intersection = 0.0;
  double sum = 0.0;
  for (size_t i = 0; i < a.dimensionality(); ++i) {
    T en1 = a.values()[i];
    T en2 = b.values()[i];
    intersection += std::min(en1, en2);
    sum += std::max(en1, en2);
  }
  if (!sum) {
    return 0;
  }
  double ratio = intersection / sum;
  return 1.0 - ratio;
}

template <typename T>
double GeneralJaccardDistance::GetDistanceSparseImpl(
    const DatapointPtr<T>& a, const DatapointPtr<T>& b) const {
  double intersection = 0.0;
  double sum = 0.0;
  auto s1 = a.nonzero_entries();
  auto s2 = b.nonzero_entries();
  std::vector<pair<DimensionIndex, double> > w1;
  std::vector<pair<DimensionIndex, double> > w2;
  std::vector<pair<DimensionIndex, double> > merged(s1 + s2);
  w1.reserve(s1);
  for (size_t i = 0; i < s1; ++i) {
    w1.push_back(std::make_pair(a.indices()[i], a.values()[i]));
  }
  w2.reserve(s2);
  for (size_t i = 0; i < s2; ++i) {
    w2.push_back(std::make_pair(b.indices()[i], b.values()[i]));
  }
  std::merge(w1.begin(), w1.end(), w2.begin(), w2.end(), merged.begin());
  while (true) {
    if (merged.empty()) break;
    if (merged.size() == 1) {
      sum += merged[0].second;
      break;
    }
    auto in1 = merged[merged.size() - 1].first;
    auto in2 = merged[merged.size() - 2].first;
    if (in1 == in2) {
      sum += std::max(merged[merged.size() - 1].second,
                      merged[merged.size() - 2].second);
      intersection += std::min(merged[merged.size() - 1].second,
                               merged[merged.size() - 2].second);
      merged.pop_back();
      merged.pop_back();
    } else {
      sum += merged[merged.size() - 1].second;
      merged.pop_back();
    }
  }
  if (!sum) {
    return 0;
  }
  double ratio = intersection / sum;
  return 1.0 - ratio;
}

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/utils/zip_sort_impl.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_UTILS_ZIP_SORT_IMPL_H_
#define SCANN_UTILS_ZIP_SORT_IMPL_H_

#include <algorithm>
#include <array>
#include <type_traits>

#include "scann/oss_wrappers/scann_aligned_malloc.h"
#include "scann/oss_wrappers/scann_bits.h"
#include "scann/utils/types.h"
#include "tensorflow/core/platform/prefetch.h"

namespace research_scann {
namespace zip_sort_internal {

enum : size_t { kInsertionSortSize = 20, kSelectionSortSize = 15 };

class DefaultComparator {
 public:
  template <typename T>
  bool operator()(const T& a, const T& b) const {
    return a < b;
  }
};

template <typename T, typename Comparator>
inline size_t MedianOf3(Comparator comp, T iter, size_t index1, size_t index2) {
  DCHECK_GT(index2, index1);
  const size_t mid_index = index1 + (index2 - index1) / 2;
  const auto& mid = iter[mid_index];
  ::tensorflow::port::prefetch<::tensorflow::port::PREFETCH_HINT_NTA>(&mid);
  const auto& first = iter[index1];
  ::tensorflow::port::prefetch<::tensorflow::port::PREFETCH_HINT_NTA>(&first);
  const auto& last = iter[index2 - 1];
  ::tensorflow::port::prefetch<::tensorflow::port::PREFETCH_HINT_NTA>(&last);

  if (comp(mid, first)) {
    if (comp(last, mid)) {
      return mid_index;
    } else if (comp(first, last)) {
      return index1;
    } else {
      return index2 - 1;
    }
  } else {
    if (comp(mid, last)) {
      return mid_index;
    } else if (comp(last, first)) {
      return index1;
    } else {
      return index2 - 1;
    }
  }
}

template <typename T>
void PrefetchIter(T iter) {
  ::tensorflow::port::prefetch<::tensorflow::port::PREFETCH_HINT_NTA>(&(*iter));
}

template <typename T, typename Comparator>
size_t MedianOf9(Comparator comp, T iter, size_t index1, size_t index2) {
  DCHECK_GE(index2 - index1, 9);

  std::array<T, 9> iterators;
  iterators[0] = iter + index1;
  PrefetchIter(iterators[0]);
  iterators[1] = iter + index1 + 1;
  PrefetchIter(iterators[1]);
  iterators[2] = iter + index1 + 2;
  PrefetchIter(iterators[2]);
  const size_t mid = index1 + (index2 - index1) / 2;
  iterators[3] = iter + mid - 1;
  PrefetchIter(iterators[3]);
  iterators[4] = iter + mid;
  PrefetchIter(iterators[4]);
  iterators[5] = iter + mid + 1;
  PrefetchIter(iterators[5]);
  iterators[6] = iter + index2 - 3;
  PrefetchIter(iterators[6]);
  iterators[7] = iter + index2 - 2;
  PrefetchIter(iterators[7]);
  iterators[8] = iter + index2 - 1;
  PrefetchIter(iterators[8]);

  for (size_t i = 0; i < 5; ++i) {
    auto to_swap = iterators.begin() + i;
    for (auto it = iterators.begin() + i + 1; it != iterators.end(); ++it) {
      if (comp(**it, **to_swap)) to_swap = it;
    }
    using std::swap;
    swap(iterators[i], *to_swap);
  }
  return iterators[4] - iter;
}

template <typename T, typename... U>
inline void ZipSwap(size_t index1, size_t index2, T begin) {
  using std::swap;
  swap(begin[index1], begin[index2]);
}

template <typename T, typename... U>
void ZipSwap(size_t index1, size_t index2, T begin, U... rest);

template <typename T, typename... U>
inline void ZipSwapRecursionHelper(size_t index1, size_t index2, T begin, T end,
                                   U... rest) {
  ZipSwap(index1, index2, begin, rest...);
}

template <typename T, typename... U>
inline void ZipSwap(size_t index1, size_t index2, T begin, U... rest) {
  using std::swap;
  swap(begin[index1], begin[index2]);
  ZipSwapRecursionHelper(index1, index2, rest...);
}

template <typename Comparator, typename T, typename... U>
void InsertionZipSort(Comparator comp, size_t index1, size_t index2, T begin,
                      U... rest) {
  const size_t size = index2 - index1;
  if (size < 2) return;
  const size_t max_j = index2 - 1;

  for (size_t i = index2 - 1; i > index1; --i) {
    for (size_t j = i - 1; j < max_j && comp(begin[j + 1], begin[j]); ++j) {
      ZipSwap(j + 1, j, begin, rest...);
    }
  }
}

template <typename Comparator, typename T, typename... U>
void SelectionZipSort(Comparator comp, size_t index1, size_t index2, T begin,
                      U... rest) {
  const T end = begin + index2;
  for (size_t i = index1; i + 1 < index2; ++i) {
    T to_swap = begin + i;
    for (T it = begin + i + 1; it != end; ++it) {
      if (comp(*it, *to_swap)) to_swap = it;
    }
    const size_t to_swap_index = to_swap - begin;
    ZipSwap(i, to_swap_index, begin, rest...);
  }
}

template <typename T>
using ConstRefIfNotPod = conditional_t<std::is_pod<decay_t<T>>::value,
                                       decay_t<T>, const decay_t<T>&>;

template <typename Comparator, typename T, typename... U>
inline size_t PivotPartition(Comparator comp, size_t index1, size_t index2,
                             T begin, U... rest) {
  DCHECK_GE(index2 - index1, 3);
  const size_t median_index = MedianOf3(comp, begin, index1, index2);
  ZipSwap(median_index, index2 - 1, begin, rest...);
  using Elem = decltype(begin[0]);
  ConstRefIfNotPod<Elem> pivot = begin[index2 - 1];
  T less_iter = begin + index1 - 1;
  T greater_iter = begin + index2 - 1;

  while (true) {
    while (comp(*(++less_iter), pivot)) {
      DCHECK(less_iter < begin + index2);
    }

    while (comp(pivot, *(--greater_iter))) {
      DCHECK(greater_iter >= begin);
    }

    if (ABSL_PREDICT_TRUE(less_iter < greater_iter)) {
      DCHECK_GE(greater_iter - begin, index1);
      DCHECK_LT(less_iter - begin, index2 - 1);
      ZipSwap(less_iter - begin, greater_iter - begin, begin, rest...);
    } else {
      break;
    }
  }

  const size_t pivot_swap_index = less_iter - begin;
  DCHECK_LT(pivot_swap_index, index2);
  ZipSwap(index2 - 1, pivot_swap_index, begin, rest...);
  return pivot_swap_index;
}

template <typename Comparator, typename T, typename... U>
size_t PivotPartitionBranchOptimized(Comparator comp, size_t index1,
                                     size_t index2, T begin, U... rest) {
  DCHECK_GE(index2 - index1, 3);
  const size_t range_start = index1, range_end = index2;
  const size_t median_index = (index2 - index1 >= 1000)
                                  ? MedianOf9(comp, begin, index1, index2)
                                  : MedianOf3(comp, begin, index1, index2);
  DCHECK_GE(median_index, range_start);
  DCHECK_LT(median_index, range_end);
  ZipSwap(median_index, index2 - 1, begin, rest...);
  using Elem = decltype(begin[0]);
  const size_t pivot_index = index2 - 1;

  ConstRefIfNotPod<Elem> pivot = begin[pivot_index];

  index2 -= 2;

  static constexpr ssize_t kNumSwapSlots = 32;
  std::array<size_t, kNumSwapSlots> left_swaps, right_swaps;
  ssize_t left_swaps_slot;
  ssize_t right_swaps_slot;
  for (;;) {
    DCHECK_LT(index1, index2);
    left_swaps_slot = -kNumSwapSlots;
    right_swaps_slot = -kNumSwapSlots;
    do {
      left_swaps.end()[left_swaps_slot] = index1;
      right_swaps.end()[right_swaps_slot] = index2;

      left_swaps_slot += !comp(begin[index1++], pivot);
      right_swaps_slot += !comp(pivot, begin[index2--]);

      if (ABSL_PREDICT_FALSE(left_swaps_slot == 0)) {
        if (ABSL_PREDICT_FALSE(right_swaps_slot == 0)) break;
        while (ABSL_PREDICT_TRUE(index1 < index2)) {
          right_swaps.end()[right_swaps_slot] = index2;
          if ((right_swaps_slot += !comp(pivot, begin[index2--])) == 0) break;
        }
        break;
      } else if (ABSL_PREDICT_FALSE(right_swaps_slot == 0)) {
        while (ABSL_PREDICT_TRUE(index1 < index2)) {
          left_swaps.end()[left_swaps_slot] = index1;
          if ((left_swaps_slot += !comp(begin[index1++], pivot)) == 0) break;
        }
        break;
      }
    } while (ABSL_PREDICT_TRUE(index1 < index2));

    const ssize_t num_swaps = std::min(kNumSwapSlots + left_swaps_slot,
                                       kNumSwapSlots + right_swaps_slot);
    for (ssize_t i = 0; i < num_swaps; ++i) {
      const size_t left_swap = left_swaps[i];
      const size_t right_swap = right_swaps[i];
      DCHECK_LT(left_swap, right_swap);
      DCHECK(!comp(begin[left_swap], begin[right_swap]));
      ZipSwap(left_swap, right_swap, begin, rest...);
    }

    if (ABSL_PREDICT_FALSE(index1 >= index2)) break;

    DCHECK_EQ(left_swaps_slot, 0);
    DCHECK_EQ(right_swaps_slot, 0);
  }
  DCHECK_LE(index1 - index2, 1);

  size_t pivot_swap_index = index1;

  pivot_swap_index +=
      (index1 == index2 && comp(begin[pivot_swap_index], pivot));

  for (; left_swaps_slot > right_swaps_slot; --left_swaps_slot) {
    DCHECK_GT(pivot_swap_index, range_start);
    --pivot_swap_index;
    const size_t left_swap = left_swaps.end()[left_swaps_slot - 1];

    DCHECK_GE(pivot_swap_index, left_swap);

    DCHECK(!comp(begin[left_swap], begin[pivot_swap_index]));
    ZipSwap(pivot_swap_index, left_swap, begin, rest...);
  }

  for (; right_swaps_slot > left_swaps_slot; --right_swaps_slot) {
    const size_t right_swap = right_swaps.end()[right_swaps_slot - 1];

    DCHECK_LE(pivot_swap_index, right_swap);

    DCHECK(!comp(begin[pivot_swap_index], begin[right_swap]));
    ZipSwap(pivot_swap_index, right_swap, begin, rest...);
    ++pivot_swap_index;
    DCHECK_LT(pivot_swap_index, range_end);
  }
  DCHECK_EQ(left_swaps_slot, right_swaps_slot);
  DCHECK(!comp(begin[pivot_swap_index], pivot));
  DCHECK(pivot_swap_index == range_start ||
         !comp(pivot, begin[pivot_swap_index - 1]));
  ZipSwap(pivot_swap_index, pivot_index, begin, rest...);
  return pivot_swap_index;
}

template <typename Comparator, typename T, typename... U>
void ZipSiftFrontDown(Comparator comp, size_t index1, size_t index2,
                      size_t abs_root, T begin, U... rest) {
  const size_t sz = index2 - index1;
  if (sz < 2) return;

  const auto it = begin + index1;

  size_t rel_root = abs_root - index1;
  for (size_t child = rel_root * 2 + 1; child < sz; child = rel_root * 2 + 1) {
    size_t to_swap = rel_root;
    if (comp(it[rel_root], it[child])) {
      to_swap = child;
    }
    if (child + 1 < sz && comp(it[to_swap], it[child + 1])) {
      to_swap = child + 1;
    }
    if (to_swap == rel_root) return;
    ZipSwap(rel_root + index1, to_swap + index1, begin, rest...);
    rel_root = to_swap;
  }
}

template <typename Comparator, typename T, typename... U>
ABSL_ATTRIBUTE_NOINLINE void ZipMakeHeap(Comparator comp, size_t index1,
                                         size_t index2, T begin, U... rest) {
  const size_t sz = index2 - index1;
  if (sz < 2) return;
  const size_t first_non_leaf_rel_index = (sz - 2) / 2;
  DCHECK_LT(first_non_leaf_rel_index * 2 + 1, sz);
  DCHECK_GE((first_non_leaf_rel_index + 1) * 2 + 1, sz);
  const size_t first_non_leaf_abs_index = first_non_leaf_rel_index + index1;
  DCHECK_GE(first_non_leaf_abs_index, index1);
  for (size_t i = first_non_leaf_abs_index;; --i) {
    ZipSiftFrontDown(comp, index1, index2, i, begin, rest...);
    if (i == index1) break;
  }
}

template <typename Comparator, typename T, typename... U>
void ZipPopHeap(Comparator comp, size_t index1, size_t index2, T begin,
                U... rest) {
  DCHECK_GT(index2, index1);
  ZipSwap(index1, index2 - 1, begin, rest...);
  ZipSiftFrontDown(comp, index1, index2 - 1, index1, begin, rest...);
}

template <typename Comparator, typename T, typename... U>
ABSL_ATTRIBUTE_NOINLINE void ZipHeapSortImpl(Comparator comp, size_t index1,
                                             size_t index2, T begin,
                                             U... rest) {
  ZipMakeHeap<Comparator, T, U...>(comp, index1, index2, begin, rest...);
  for (; index2 > index1; --index2) {
    ZipPopHeap(comp, index1, index2, begin, rest...);
  }
}

inline size_t ComputeDepthLimit(const size_t range_size) {
  static_assert(sizeof(size_t) == 4 || sizeof(size_t) == 8,
                "Unknown architecture.  Size_t is neither 32 nor 64-bit.");
  if (range_size == 0) {
    return 0;
  } else if (sizeof(size_t) == 4) {
    return 2 * bits::Log2Ceiling(range_size);
  } else {
    return 2 * bits::Log2Ceiling64(range_size);
  }
}

template <typename Comparator, typename T, typename... U>
ABSL_ATTRIBUTE_NOINLINE void ZipSortImpl(Comparator comp, size_t index1,
                                         size_t index2, size_t depth_limit,
                                         T begin, U... rest) {
  while (true) {
    if (index2 - index1 < kInsertionSortSize) {
      InsertionZipSort<Comparator, T, U...>(comp, index1, index2, begin,
                                            rest...);
      return;
    } else if (depth_limit == 0) {
      ZipHeapSortImpl(comp, index1, index2, begin, rest...);
      return;
    }

    --depth_limit;
    const size_t pivot_index = PivotPartition<Comparator, T, U...>(
        comp, index1, index2, begin, rest...);

    const size_t index1_less = index1;
    const size_t index2_less = pivot_index;
    DCHECK_GE(index2_less, index1_less);
    const size_t index1_greater = pivot_index + 1;
    const size_t index2_greater = index2;
    DCHECK_GE(index2_greater, index1_greater);

    if (index2_greater - index1_greater > index2_less - index1_less) {
      ZipSortImpl(comp, index1_less, index2_less, depth_limit, begin, rest...);
      index1 = index1_greater;
      index2 = index2_greater;
    } else {
      ZipSortImpl(comp, index1_greater, index2_greater, depth_limit, begin,
                  rest...);
      index1 = index1_less;
      index2 = index2_less;
    }
  }
}

template <typename Comparator, typename T, typename... U>
ABSL_ATTRIBUTE_NOINLINE void ZipSortImplBranchOptimized(Comparator comp,
                                                        size_t index1,
                                                        size_t index2,
                                                        size_t depth_limit,
                                                        T begin, U... rest) {
  while (true) {
    if (index2 - index1 < kSelectionSortSize) {
      SelectionZipSort<Comparator, T, U...>(comp, index1, index2, begin,
                                            rest...);
      return;
    } else if (depth_limit == 0) {
      ZipHeapSortImpl(comp, index1, index2, begin, rest...);
      return;
    }

    --depth_limit;
    const size_t pivot_index =
        PivotPartitionBranchOptimized<Comparator, T, U...>(comp, index1, index2,
                                                           begin, rest...);

    const size_t index1_less = index1;
    const size_t index2_less = pivot_index;
    DCHECK_GE(index2_less, index1_less);
    const size_t index1_greater = pivot_index + 1;
    const size_t index2_greater = index2;
    DCHECK_GE(index2_greater, index1_greater);

    if (index2_greater - index1_greater > index2_less - index1_less) {
      ZipSortImplBranchOptimized(comp, index1_less, index2_less, depth_limit,
                                 begin, rest...);
      index1 = index1_greater;
      index2 = index2_greater;
    } else {
      ZipSortImplBranchOptimized(comp, index1_greater, index2_greater,
                                 depth_limit, begin, rest...);
      index1 = index1_less;
      index2 = index2_less;
    }
  }
}

template <typename Comparator, typename T, typename... U>
void ZipNthElementImpl(Comparator comp, size_t n, T begin, T end, U... rest) {
  size_t index1 = 0;
  size_t index2 = end - begin;

  while (index2 - index1 > 3) {
    const size_t pivot_index = PivotPartition<Comparator, T, U...>(
        comp, index1, index2, begin, rest...);
    if (pivot_index == n) {
      return;
    } else if (pivot_index < n) {
      index1 = pivot_index + 1;
    } else {
      index2 = pivot_index;
    }
  }

  InsertionZipSort<Comparator, T, U...>(comp, index1, index2, begin, rest...);
}

template <typename Comparator, typename T, typename... U>
void ZipNthElementImplBranchOptimized(Comparator comp, size_t n, T begin, T end,
                                      U... rest) {
  size_t index1 = 0;
  size_t index2 = end - begin;

  while (index2 - index1 > 3) {
    const size_t pivot_index =
        PivotPartitionBranchOptimized<Comparator, T, U...>(comp, index1, index2,
                                                           begin, rest...);
    if (pivot_index == n) {
      return;
    } else if (pivot_index < n) {
      index1 = pivot_index + 1;
    } else {
      index2 = pivot_index;
    }
  }

  SelectionZipSort<Comparator, T, U...>(comp, index1, index2, begin, rest...);
}

}  // namespace zip_sort_internal

template <typename Predicate, typename T, typename... U>
size_t ZipPartition(Predicate pred, T begin, T end, U... rest) {
  using zip_sort_internal::ZipSwap;
  T less_iter = begin - 1;
  T greater_iter = end;

  while (++less_iter < end && pred(*less_iter)) {
  }

  while (greater_iter > less_iter && !pred(*(--greater_iter))) {
  }

  if (less_iter < greater_iter) {
    ZipSwap(less_iter - begin, greater_iter - begin, begin, rest...);

    while (true) {
      while (pred(*(++less_iter))) {
      }
      while (greater_iter > less_iter && !pred(*(--greater_iter))) {
      }

      if (less_iter < greater_iter) {
        ZipSwap(less_iter - begin, greater_iter - begin, begin, rest...);
      } else {
        break;
      }
    }
  }

  return less_iter - begin;
}

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/utils/infinite_one_array.h<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.



#ifndef SCANN_UTILS_INFINITE_ONE_ARRAY_H_
#define SCANN_UTILS_INFINITE_ONE_ARRAY_H_

#include "scann/utils/types.h"

namespace research_scann {

template <typename T>
class InfiniteOneIterator {
 public:
  T operator*() const { return T(1); }
  T operator[](size_t i) const { return T(1); }
  void operator++() {}
  void operator++(int dummy) {}
  InfiniteOneIterator operator+=(size_t i) { return *this; }
  InfiniteOneIterator operator-=(size_t i) { return *this; }

  bool operator<(const InfiniteOneIterator<T>& x) const { return true; }
};

template <typename T>
class InfiniteOneArray {
 public:
  InfiniteOneArray() {}
  size_t size() const { return numeric_limits<size_t>::max(); }
  bool empty() const { return false; }
  T operator[](size_t i) const { return T(1); }
  T front() const { return T(1); }
  T back() const { return T(1); }
  InfiniteOneIterator<T> begin() { return InfiniteOneIterator<T>(); }
  InfiniteOneIterator<T> end() { return InfiniteOneIterator<T>(); }
  InfiniteOneArray data() const { return *this; }
};

template <typename T>
class InfiniteOneMatrix {
 public:
  InfiniteOneMatrix() {}
  size_t size() const { return numeric_limits<size_t>::max(); }
  bool empty() const { return false; }
  InfiniteOneArray<T> operator[](size_t i) const {
    return InfiniteOneArray<T>();
  }

  InfiniteOneArray<T> front() const { return InfiniteOneArray<T>(); }
  InfiniteOneArray<T> back() const { return InfiniteOneArray<T>(); }
};

}  // namespace research_scann

#endif


<|start_filename|>scann/scann/hashes/asymmetric_hashing2/querying.cc<|end_filename|>
// Copyright 2021 The Google Research Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#include "scann/hashes/asymmetric_hashing2/querying.h"

#include <cstdint>

#include "scann/utils/common.h"
#include "scann/utils/intrinsics/flags.h"

using std::shared_ptr;

namespace research_scann {
namespace asymmetric_hashing2 {

PackedDataset CreatePackedDataset(
    const DenseDataset<uint8_t>& hashed_database) {
  PackedDataset result;
  result.bit_packed_data =
      asymmetric_hashing_internal::CreatePackedDataset(hashed_database);
  result.num_datapoints = hashed_database.size();
  result.num_blocks =
      (!hashed_database.empty()) ? (hashed_database[0].nonzero_entries()) : 0;
  return result;
}

DenseDataset<uint8_t> UnpackDataset(const PackedDataset& packed) {
  const size_t num_dim = packed.num_blocks, num_dp = packed.num_datapoints;

  vector<uint8_t> unpacked(num_dim * num_dp);

  int idx = 0;
  for (int dp_block = 0; dp_block < num_dp / 32; dp_block++) {
    const int out_idx = 32 * dp_block;
    for (int dim = 0; dim < num_dim; dim++) {
      for (int offset = 0; offset < 16; offset++) {
        uint8_t data = packed.bit_packed_data[idx++];
        unpacked[(out_idx | offset) * num_dim + dim] = data & 15;
        unpacked[(out_idx | 16 | offset) * num_dim + dim] = data >> 4;
      }
    }
  }

  if (num_dp % 32 != 0) {
    const int out_idx = num_dp - (num_dp % 32);
    for (int dim = 0; dim < num_dim; dim++) {
      for (int offset = 0; offset < 16; offset++) {
        uint8_t data = packed.bit_packed_data[idx++];
        int idx1 = out_idx | offset, idx2 = out_idx | 16 | offset;
        if (idx1 < num_dp) unpacked[idx1 * num_dim + dim] = data & 15;
        if (idx2 < num_dp) unpacked[idx2 * num_dim + dim] = data >> 4;
      }
    }
  }
  return DenseDataset<uint8_t>(unpacked, packed.num_datapoints);
}

template <typename T>
AsymmetricQueryer<T>::AsymmetricQueryer(
    shared_ptr<const ChunkingProjection<T>> projector,
    shared_ptr<const DistanceMeasure> lookup_distance,
    shared_ptr<const Model<T>> model)
    : projector_(std::move(projector)),
      lookup_distance_(std::move(lookup_distance)),
      model_(std::move(model)) {}

template <typename T>
StatusOr<LookupTable> AsymmetricQueryer<T>::CreateLookupTable(
    const DatapointPtr<T>& query,
    AsymmetricHasherConfig::LookupType lookup_type,
    AsymmetricHasherConfig::FixedPointLUTConversionOptions
        float_int_conversion_options) const {
  switch (lookup_type) {
    case AsymmetricHasherConfig::FLOAT:
      return CreateLookupTable<float>(query, float_int_conversion_options);
    case AsymmetricHasherConfig::INT8:
    case AsymmetricHasherConfig::INT8_LUT16:
      return CreateLookupTable<int8_t>(query, float_int_conversion_options);
    case AsymmetricHasherConfig::INT16:
      return CreateLookupTable<int16_t>(query, float_int_conversion_options);
    default:
      return InvalidArgumentError("Unrecognized lookup type.");
  }
}

SCANN_INSTANTIATE_TYPED_CLASS(, AsymmetricQueryer);

}  // namespace asymmetric_hashing2
}  // namespace research_scann
