text: # Copyright 1999-2020 Alibaba Group Holding Ltd.
  #
  # 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.
  
  import atexit
  import os
  import operator
  import tempfile
  from enum import Enum
  
  import numpy as np
  try:
      import faiss
  except ImportError:  # pragma: no cover
      faiss = None
  
  from ... import opcodes as OperandDef
  from ...context import RunningMode
  from ...core import TilesError
  from ...core.operand import OperandStage
  from ...serialize import KeyField, StringField, Int64Field, \
      Int32Field, BoolField, Int8Field
  from ...tensor import tensor as astensor
  from ...tensor.core import TensorOrder
  from ...tensor.random import RandomState
  from ...tensor.array_utils import as_same_device, device
  from ...tensor.utils import check_random_state, gen_random_seeds
  from ...utils import check_chunks_unknown_shape, require_not_none, recursive_tile
  from ..operands import LearnOperand, LearnOperandMixin, OutputType
  
  
  class MemoryRequirementGrade(Enum):
      minimum = 0
      low = 1
      high = 2
      maximum = 3
  
  
  if faiss is not None:
      METRIC_TO_FAISS_METRIC_TYPE = {
          'l2': faiss.METRIC_L2,
          'euclidean': faiss.METRIC_L2,
          'innerproduct': faiss.METRIC_INNER_PRODUCT,
          'cosine': faiss.METRIC_INNER_PRODUCT,
      }
  else:  # pragma: no cover
      METRIC_TO_FAISS_METRIC_TYPE = {}
  
  
  @require_not_none(faiss)
  class FaissBuildIndex(LearnOperand, LearnOperandMixin):
      _op_type_ = OperandDef.FAISS_BUILD_INDEX
  
      _input = KeyField('input')
      _metric = StringField('metric')
      _faiss_index = StringField('faiss_index')
      _n_sample = Int64Field('n_sample')
      _seed = Int32Field('seed')
      _same_distribution = BoolField('same_distribution')
      _accuracy = BoolField('accuracy')
      _memory_require = Int8Field('memory_require',
                                  on_serialize=operator.attrgetter('value'),
                                  on_deserialize=MemoryRequirementGrade)
      # for test purpose, could be 'object', 'filename' or 'bytes'
      _return_index_type = StringField('return_index_type')
  
      def __init__(self, metric=None, faiss_index=None, n_sample=None, seed=None,
                   same_distribution=None, return_index_type=None,
                   accuracy=None, memory_require=None,
                   output_types=None, **kw):
          super().__init__(_metric=metric, _faiss_index=faiss_index, _n_sample=n_sample,
                           _seed=seed, _same_distribution=same_distribution,
                           _return_index_type=return_index_type,
                           _accuracy=accuracy, _memory_require=memory_require,
                           _output_types=output_types, **kw)
          if self.output_types is None:
              self.output_types = [OutputType.object]
  
      @property
      def input(self):
          return self._input
  
      @property
      def metric(self):
          return self._metric
  
      @property
      def faiss_metric_type(self):
          return METRIC_TO_FAISS_METRIC_TYPE[self._metric]
  
      @property
      def faiss_index(self):
          return self._faiss_index
  
      @property
      def n_sample(self):
          return self._n_sample
  
      @property
      def seed(self):
          return self._seed
  
      @property
      def same_distribution(self):
          return self._same_distribution
  
      @property
      def accuracy(self):
          return self._accuracy
  
      @property
      def memory_require(self):
          return self._memory_require
  
      @property
      def return_index_type(self):
          return self._return_index_type
  
      def _set_inputs(self, inputs):
          super()._set_inputs(inputs)
          self._input = self._inputs[0]
  
      def __call__(self, X):
          return self.new_tileable([X])
  
      @classmethod
      def tile(cls, op):
          check_chunks_unknown_shape(op.inputs, TilesError)
  
          in_tensor = astensor(op.input, np.dtype(np.float32))._inplace_tile()
          if op.faiss_index == 'auto':
              faiss_index, n_sample = _gen_index_string_and_sample_count(
                  in_tensor.shape, op.n_sample, op.accuracy, op.memory_require,
                  gpu=op.gpu, **op.extra_params)
              op._n_sample = n_sample
          else:
              faiss_index, n_sample = op.faiss_index, op.n_sample
  
          if len(in_tensor.chunks) == 1:
              return cls._tile_one_chunk(op, faiss_index, n_sample)
  
          if in_tensor.chunk_shape[1] != 1:
              # make sure axis 1 has 1 chunk
              in_tensor = in_tensor.rechunk({1: in_tensor.shape[1]})._inplace_tile()
          return cls._tile_chunks(op, in_tensor, faiss_index, n_sample)
  
      @classmethod
      def _tile_one_chunk(cls, op, faiss_index, n_sample):
          in_chunk = op.input.chunks[0]
          chunk_op = op.copy().reset_key()
          chunk_op._faiss_index = faiss_index
          chunk_op._n_sample = n_sample
          chunk = chunk_op.new_chunk([in_chunk], index=in_chunk.index)
  
          new_op = op.copy()
          kw = op.outputs[0].params
          kw['chunks'] = [chunk]
          kw['nsplits'] = ((1,),)
          return new_op.new_tileables(op.inputs, kws=[kw])
  
      @classmethod
      def _tile_chunks(cls, op, in_tensor, faiss_index, n_sample):
          """
          If the distribution on each chunk is the same,
          refer to:
          https://github.com/facebookresearch/faiss/wiki/FAQ#how-can-i-distribute-index-building-on-several-machines
  
          1. train an IndexIVF* on a representative sample of the data, store it.
          2. for each node, load the trained index, add the local data to it, store the resulting populated index
          3. on a central node, load all the populated indexes and merge them.
          """
          faiss_index_ = faiss.index_factory(in_tensor.shape[1], faiss_index,
                                             op.faiss_metric_type)
          # Training on sample data when two conditions meet
          # 1. the index type requires for training, e.g. Flat does not require
          # 2. distributions of chunks are the same, in not,
          #    train separately on each chunk data
          need_sample_train = not faiss_index_.is_trained and op.same_distribution
          need_merge_index = hasattr(faiss_index_, 'merge_from') if need_sample_train else False
  
          train_chunk = None
          if need_sample_train:
              # sample data to train
              rs = RandomState(op.seed)
              sampled_index = rs.choice(in_tensor.shape[0], size=n_sample,
                                        replace=False, chunk_size=n_sample)
              sample_tensor = recursive_tile(in_tensor[sampled_index])
              assert len(sample_tensor.chunks) == 1
              sample_chunk = sample_tensor.chunks[0]
              train_op = FaissTrainSampledIndex(faiss_index=faiss_index, metric=op.metric,
                                                return_index_type=op.return_index_type)
              train_chunk = train_op.new_chunk([sample_chunk])
          elif op.gpu:  # pragma: no cover
              # if not need train, and on gpu, just merge data together to train
              in_tensor = in_tensor.rechunk(in_tensor.shape)._inplace_tile()
  
          # build index for each input chunk
          build_index_chunks = []
          for i, chunk in enumerate(in_tensor.chunks):
              build_index_op = op.copy().reset_key()
              build_index_op.stage = OperandStage.map
              build_index_op._faiss_index = faiss_index
              if train_chunk is not None:
                  build_index_chunk = build_index_op.new_chunk(
                      [chunk, train_chunk], index=(i,))
              else:
                  build_index_chunk = build_index_op.new_chunk([chunk], index=(i,))
              build_index_chunks.append(build_index_chunk)
  
          out_chunks = []
          if need_merge_index:
              assert op.n_sample is not None
              # merge all indices into one, do only when trained on sample data
              out_chunk_op = op.copy().reset_key()
              out_chunk_op._faiss_index = faiss_index
              out_chunk_op.stage = OperandStage.agg
              out_chunk = out_chunk_op.new_chunk(build_index_chunks, index=(0,))
              out_chunks.append(out_chunk)
          else:
              out_chunks.extend(build_index_chunks)
  
          new_op = op.copy()
          return new_op.new_tileables(op.inputs, chunks=out_chunks,
                                      nsplits=((len(out_chunks),),))
  
      @classmethod
      def _execute_one_chunk(cls, ctx, op):
          (inp,), device_id, xp = as_same_device(
              [ctx[c.key] for c in op.inputs], device=op.device, ret_extra=True)
  
          with device(device_id):
              inp = inp.astype(np.float32, copy=False)
              # create index
              index = faiss.index_factory(inp.shape[1], op.faiss_index,
                                          op.faiss_metric_type)
              # GPU
              if device_id >= 0:  # pragma: no cover
                  index = _index_to_gpu(index, device_id)
  
              # train index
              if not index.is_trained:
                  assert op.n_sample is not None
                  sample_indices = xp.random.choice(inp.shape[0],
                                                    size=op.n_sample, replace=False)
                  sampled = inp[sample_indices]
                  index.train(sampled)
  
              if op.metric == 'cosine':
                  # faiss does not support cosine distances directly,
                  # data needs to be normalize before adding to index,
                  # refer to:
                  # https://github.com/facebookresearch/faiss/wiki/FAQ#how-can-i-index-vectors-for-cosine-distance
                  faiss.normalize_L2(inp)
              # add vectors to index
              if device_id >= 0:  # pragma: no cover
                  # gpu
                  index.add_c(inp.shape[0], _swig_ptr_from_cupy_float32_array(inp))
              else:
                  index.add(inp)
  
              ctx[op.outputs[0].key] = _store_index(ctx, op, index, device_id)
  
      @classmethod
      def _execute_map(cls, ctx, op):
          (data,), device_id, xp = as_same_device(
              [ctx[op.inputs[0].key]], device=op.device, ret_extra=True)
          index = ctx[op.inputs[1].key] if len(op.inputs) == 2 else None
  
          with device(device_id):
              data = xp.ascontiguousarray(data)
              if index is not None:
                  # fetch the trained index
                  trained_index = _load_index(ctx, op, index, device_id)
                  return_index_type = _get_index_type(op.return_index_type, ctx)
                  if return_index_type == 'object':
                      # clone a new one,
                      # because faiss does not ensure thread-safe for operations that change index
                      # https://github.com/facebookresearch/faiss/wiki/Threads-and-asynchronous-calls#thread-safety
                      trained_index = faiss.clone_index(trained_index)
              else:
                  trained_index = faiss.index_factory(data.shape[1], op.faiss_index,
                                                      op.faiss_metric_type)
                  if op.same_distribution:
                      # no need to train, just create index
                      pass
                  else:
                      # distribution no the same, train on each chunk
                      trained_index.train(data)
  
                  if device_id >= 0:  # pragma: no cover
                      trained_index = _index_to_gpu(trained_index, device_id)
              if op.metric == 'cosine':
                  # faiss does not support cosine distances directly,
                  # data needs to be normalize before adding to index,
                  # refer to:
                  # https://github.com/facebookresearch/faiss/wiki/FAQ#how-can-i-index-vectors-for-cosine-distance
                  faiss.normalize_L2(data)
  
              # add data into index
              if device_id >= 0:  # pragma: no cover
                  # gpu
                  trained_index.add_c(data.shape[0], _swig_ptr_from_cupy_float32_array(data))
              else:
                  trained_index.add(data)
  
              ctx[op.outputs[0].key] = _store_index(ctx, op, trained_index, device_id)
  
      @classmethod
      def _execute_agg(cls, ctx, op):
          device_id = op.device
          if device_id is None:
              device_id = -1
          inputs = [ctx[inp.key] for inp in op.inputs]
  
          with device(device_id):
              merged_index = None
              indexes = []
              for index in inputs:
                  index = _load_index(ctx, op, index, device_id)
                  indexes.append(index)
                  assert hasattr(index, 'merge_from')
                  if merged_index is None:
                      merged_index = index
                  else:
                      merged_index.merge_from(index, index.ntotal)
  
              ctx[op.outputs[0].key] = _store_index(ctx, op, merged_index, device_id)
  
      @classmethod
      def execute(cls, ctx, op):
          if op.stage == OperandStage.map:
              cls._execute_map(ctx, op)
          elif op.stage == OperandStage.agg:
              cls._execute_agg(ctx, op)
          else:
              assert op.stage is None
              cls._execute_one_chunk(ctx, op)
  
  
  def _get_index_type(return_index_type, ctx):
      if return_index_type is None:  # pragma: no cover
          if ctx.running_mode == RunningMode.local:
              return_index_type = 'object'
          elif ctx.running_mode == RunningMode.local_cluster:
              return_index_type = 'filename'
          else:
              return_index_type = 'bytes'
      return return_index_type
  
  
  def _store_index(ctx, op, index, device_id):
      return_index_type = _get_index_type(op.return_index_type, ctx)
  
      if return_index_type == 'object':
          # no need to serialize
          return index
      elif return_index_type == 'filename':
          # save to file, then return filename
          if device_id >= 0:  # pragma: no cover
              # for gpu, convert to cpu first
              index = faiss.index_gpu_to_cpu(index)
          fn = tempfile.mkstemp('.index', prefix='faiss_')[1]
          faiss.write_index(index, fn)
  
          atexit.register(lambda: os.remove(fn))
  
          return fn
      else:
          if device_id >= 0:  # pragma: no cover
              # for gpu, convert to cpu first
              index = faiss.index_gpu_to_cpu(index)
          # distributed, save to file, then return in memory bytes
          fn = tempfile.mkstemp('.index', prefix='faiss_')[1]
          faiss.write_index(index, fn)
          try:
              with open(fn, 'rb') as f:
                  return f.read()
          finally:
              os.remove(fn)
  
  
  def _load_index(ctx, op, index, device_id):
      return_index_type = _get_index_type(op.return_index_type, ctx)
  
      if return_index_type == 'object':
          # local
          return index
      elif return_index_type == 'filename':
          # local cluster
          return faiss.read_index(index)
      else:
          # distributed
          fn = tempfile.mkstemp('.index', prefix='faiss_')[1]
          with open(fn, 'wb') as f:
              f.write(index)
          index = faiss.read_index(f.name)
          if device_id >= 0:  # pragma: no cover
              index = _index_to_gpu(index, device_id)
          return index
  
  
  def _index_to_gpu(index, device_id):  # pragma: no cover
      res = faiss.StandardGpuResources()
      return faiss.index_cpu_to_gpu(res, device_id, index)
  
  
  def _swig_ptr_from_cupy_float32_array(x):  # pragma: no cover
      assert x.flags.c_contiguous
      assert x.dtype == np.float32
      data_ptr = x.__cuda_array_interface__['data'][0]
      return faiss.cast_integer_to_float_ptr(data_ptr)
  
  
  def _swig_ptr_from_cupy_int64_array(x):  # pragma: no cover
      assert x.flags.c_contiguous
      assert x.dtype == np.int64
      data_ptr = x.__cuda_array_interface__['data'][0]
      return faiss.cast_integer_to_long_ptr(data_ptr)
  
  
  @require_not_none(faiss)
  class FaissTrainSampledIndex(LearnOperand, LearnOperandMixin):
      _op_type_ = OperandDef.FAISS_TRAIN_SAMPLED_INDEX
  
      _input = KeyField('input')
      _metric = StringField('metric')
      _faiss_index = StringField('faiss_index')
      # for test purpose, could be 'object', 'filename' or 'bytes'
      _return_index_type = StringField('return_index_type')
  
      def __init__(self, faiss_index=None, metric=None,
                   return_index_type=None, output_types=None, **kw):
          super().__init__(_faiss_index=faiss_index, _metric=metric,
                           _return_index_type=return_index_type,
                           _output_types=output_types, **kw)
          if self.output_types is None:
              self.output_types = [OutputType.object]
  
      @property
      def input(self):
          return self._input
  
      @property
      def metric(self):
          return self._metric
  
      @property
      def faiss_metric_type(self):
          return METRIC_TO_FAISS_METRIC_TYPE[self.metric]
  
      @property
      def faiss_index(self):
          return self._faiss_index
  
      @property
      def return_index_type(self):
          return self._return_index_type
  
      def _set_inputs(self, inputs):
          super()._set_inputs(inputs)
          self._input = self._inputs[0]
  
      @classmethod
      def execute(cls, ctx, op):
          (data,), device_id, _ = as_same_device(
              [ctx[op.input.key]], device=op.device, ret_extra=True)
  
          with device(device_id):
              index = faiss.index_factory(data.shape[1], op.faiss_index,
                                          op.faiss_metric_type)
  
              if device_id >= 0:  # pragma: no cover
                  # GPU
                  index = _index_to_gpu(index, device_id)
                  index.train_c(data.shape[0], _swig_ptr_from_cupy_float32_array(data))
              else:
                  index.train(data)
  
              ctx[op.outputs[0].key] = _store_index(
                  ctx, op, index, device_id)
  
  
  def _gen_index_string_and_sample_count(shape, n_sample, accuracy, memory_require, gpu=False, **kw):
      """
      Generate index string and sample count according to guidance of faiss:
      https://github.com/facebookresearch/faiss/wiki/Guidelines-to-choose-an-index
      """
      size, dim = shape
      memory_require = _get_memory_require(memory_require)
  
      if accuracy or size < 10 ** 5:
          # Flat is the only index that guarantees exact results
          # no need to train, thus sample count is None
          return 'Flat', None
  
      if memory_require == MemoryRequirementGrade.maximum and not gpu:
          x = kw.get('M', 32)  # get medium number by default
          if x < 4 or x > 64:
              raise ValueError(f'HNSWx requires M that between 4 and 64, got {x}')
          return f'HNSW{x}', None
  
      if memory_require in (MemoryRequirementGrade.high, MemoryRequirementGrade.maximum):
          basement = '{},Flat'
      elif memory_require == MemoryRequirementGrade.low:
          x = kw.get('dim', dim // 2)
          basement = f'PCAR{x},{{}},SQ8'
      elif memory_require == MemoryRequirementGrade.minimum:
          x = kw.get('M', min(64, dim // 2))
          if x > 64:
              raise ValueError(f'PQx requires M <= 64, got {x}')
          y = kw.get('dim', None)
          if y is not None and y % x != 0:
              raise ValueError(f'OPQx_y requires dim is a multiple of M({x}), got dim: {y}')
          y = min(dim, 4 * x)
          y = x * (y // x)  # make sure y is a multiple of x
          basement = f'OPQ{x}_{y},{{}},PQ{x}'
      else:  # pragma: no cover
          raise ValueError('unknown memory require')
  
      # now choose the clustering options
      if size < 10 ** 6 or (size < 10 ** 7 and gpu):
          # < 1M, or <10M but need GPU
          k = kw.get('k', 5 * int(np.sqrt(size)))
          if k < 4 * int(np.sqrt(size)) or k > 16 * int(np.sqrt(size)):
              raise ValueError(f'k should be between 4 * sqrt(N) and 16 * sqrt(N), got {k}')
          index_str = basement.format(f'IVF{k}')
          if n_sample is None:
              # 30 * k - 256 * k
              n_sample = min(30 * k, size)
      elif size < 10 ** 7 and not gpu:
          # 1M - 10M
          index_str = basement.format('IVF65536_HNSW32')
          if n_sample is None:
              # between 30 * 65536 and 256 * 65536
              n_sample = 32 * 65536
      elif size < 10 ** 8:
          index_str = basement.format('IVF65536_HNSW32')
          n_sample = 64 * 65536 if n_sample is None else n_sample
      else:
          index_str = basement.format('IVF1048576_HNSW32')
          n_sample = 64 * 65536 if n_sample is None else n_sample
  
      return index_str, n_sample
  
  
  def _get_memory_require(memory_require):
      if isinstance(memory_require, str):
          return getattr(MemoryRequirementGrade, memory_require)
      elif isinstance(memory_require, MemoryRequirementGrade):
          return memory_require
      return MemoryRequirementGrade(memory_require)
  
  
  @require_not_none(faiss)
  def build_faiss_index(X, index_name='auto', n_sample=None, metric="euclidean",
                        random_state=None, same_distribution=True,
                        accuracy=False, memory_require=None, **kw):
      X = astensor(X)
  
      if metric not in METRIC_TO_FAISS_METRIC_TYPE:
          raise ValueError(f'unknown metric: {metric}')
      if index_name != 'auto':
          try:
              faiss.index_factory(X.shape[1], index_name,
                                  METRIC_TO_FAISS_METRIC_TYPE[metric])
          except RuntimeError:
              raise ValueError(f'illegal faiss index: {index_name}')
  
      rs = check_random_state(random_state)
      if isinstance(rs, RandomState):
          rs = rs.to_numpy()
      seed = gen_random_seeds(1, rs)[0]
      if memory_require is None:
          memory_require = MemoryRequirementGrade.low
      else:
          memory_require = _get_memory_require(memory_require)
      op = FaissBuildIndex(faiss_index=index_name, metric=metric,
                           n_sample=n_sample, gpu=X.op.gpu, seed=seed,
                           same_distribution=same_distribution,
                           accuracy=accuracy, memory_require=memory_require, **kw)
      return op(X)
  
  
  class FaissQuery(LearnOperand, LearnOperandMixin):
      _op_type_ = OperandDef.FAISS_QUERY
  
      _input = KeyField('input')
      _faiss_index = KeyField('faiss_index')
      _metric = StringField('metric')
      _n_neighbors = Int32Field('n_neighbors')
      _return_distance = BoolField('return_distance')
      _nprobe = Int64Field('nprobe')
      # for test purpose, could be 'object', 'filename' or 'bytes'
      _return_index_type = StringField('return_index_type')
  
      def __init__(self, faiss_index=None, metric=None, n_neighbors=None,
                   return_distance=None, return_index_type=None,
                   nprobe=None, output_types=None, gpu=None, **kw):
          super().__init__(_faiss_index=faiss_index, _n_neighbors=n_neighbors, _metric=metric,
                           _return_distance=return_distance, _output_types=output_types,
                           _nprobe=nprobe, _return_index_type=return_index_type, _gpu=gpu, **kw)
          if self.output_types is None:
              self.output_types = [OutputType.tensor] * self.output_limit
  
      @property
      def input(self):
          return self._input
  
      @property
      def faiss_index(self):
          return self._faiss_index
  
      @property
      def metric(self):
          return self._metric
  
      @property
      def n_neighbors(self):
          return self._n_neighbors
  
      @property
      def nprobe(self):
          return self._nprobe
  
      @property
      def return_distance(self):
          return self._return_distance
  
      @property
      def return_index_type(self):
          return self._return_index_type
  
      @property
      def output_limit(self):
          return 2 if self._return_distance else 1
  
      def _set_inputs(self, inputs):
          super()._set_inputs(inputs)
          self._input = self._inputs[0]
          if self._faiss_index is not None:
              self._faiss_index = self._inputs[1]
  
      def __call__(self, y):
          kws = []
          if self._return_distance:
              kws.append({'shape': (y.shape[0], self._n_neighbors),
                          'dtype': np.dtype(np.float32),
                          'order': TensorOrder.C_ORDER,
                          'type': 'distance'})
          kws.append({
              'shape': (y.shape[0], self._n_neighbors),
              'dtype': np.dtype(np.int64),
              'order': TensorOrder.C_ORDER,
              'type': 'indices'
          })
          return self.new_tileables([y, self._faiss_index], kws=kws)
  
      @classmethod
      def tile(cls, op):
          in_tensor = astensor(op.input)
  
          if in_tensor.chunk_shape[1] != 1:
              check_chunks_unknown_shape([in_tensor], TilesError)
              in_tensor = in_tensor.rechunk({1: in_tensor.shape[1]})._inplace_tile()
  
          out_chunks = [], []
          for chunk in in_tensor.chunks:
              chunk_op = op.copy().reset_key()
              chunk_kws = []
              if op.return_distance:
                  chunk_kws.append({
                      'shape': (chunk.shape[0], op.n_neighbors),
                      'dtype': np.dtype(np.float32),
                      'order': TensorOrder.C_ORDER,
                      'index': chunk.index,
                      'type': 'distance'
                  })
              chunk_kws.append({
                  'shape': (chunk.shape[0], op.n_neighbors),
                  'dtype': np.dtype(np.int64),
                  'order': TensorOrder.C_ORDER,
                  'index': chunk.index,
                  'type': 'indices'
              })
              in_chunks = [chunk]
              in_chunks.extend(op.faiss_index.chunks)
              chunks = chunk_op.new_chunks(in_chunks, kws=chunk_kws)
              if op.return_distance:
                  out_chunks[0].append(chunks[0])
              out_chunks[1].append(chunks[-1])
  
          new_op = op.copy()
          kws = [out.params for out in op.outputs]
          if op.return_distance:
              kws[0]['chunks'] = out_chunks[0]
              kws[0]['nsplits'] = (in_tensor.nsplits[0], (op.n_neighbors,))
          kws[-1]['chunks'] = out_chunks[1]
          kws[-1]['nsplits'] = (in_tensor.nsplits[0], (op.n_neighbors,))
          return new_op.new_tileables(op.inputs, kws=kws)
  
      @classmethod
      def execute(cls, ctx, op):
          (y,), device_id, xp = as_same_device(
              [ctx[op.input.key]], device=op.device, ret_extra=True)
          indexes = [_load_index(ctx, op, ctx[index.key], device_id)
                     for index in op.inputs[1:]]
  
          with device(device_id):
              y = xp.ascontiguousarray(y, dtype=np.float32)
  
              if len(indexes) == 1:
                  index = indexes[0]
              else:
                  index = faiss.IndexShards(indexes[0].d)
                  [index.add_shard(ind) for ind in indexes]
  
              if op.metric == 'cosine':
                  # faiss does not support cosine distances directly,
                  # data needs to be normalize before searching,
                  # refer to:
                  # https://github.com/facebookresearch/faiss/wiki/FAQ#how-can-i-index-vectors-for-cosine-distance
                  faiss.normalize_L2(y)
  
              if op.nprobe is not None:
                  index.nprobe = op.nprobe
  
              if device_id >= 0:  # pragma: no cover
                  n = y.shape[0]
                  k = op.n_neighbors
                  distances = xp.empty((n, k), dtype=xp.float32)
                  indices = xp.empty((n, k), dtype=xp.int64)
                  index.search_c(n, _swig_ptr_from_cupy_float32_array(y),
                                 k, _swig_ptr_from_cupy_float32_array(distances),
                                 _swig_ptr_from_cupy_int64_array(indices))
              else:
                  distances, indices = index.search(y, op.n_neighbors)
              if op.return_distance:
                  if index.metric_type == faiss.METRIC_L2:
                      # make it equivalent to `pairwise.euclidean_distances`
                      distances = xp.sqrt(distances, out=distances)
                  elif op.metric == 'cosine':
                      # make it equivalent to `pairwise.cosine_distances`
                      distances = xp.subtract(1, distances, out=distances)
                  ctx[op.outputs[0].key] = distances
              ctx[op.outputs[-1].key] = indices
  
  
  @require_not_none(faiss)
  def faiss_query(faiss_index, data, n_neighbors, return_distance=True, nprobe=None):
      data = astensor(data)
      op = FaissQuery(faiss_index=faiss_index, n_neighbors=n_neighbors,
                      metric=faiss_index.op.metric, return_distance=return_distance,
                      return_index_type=faiss_index.op.return_index_type,
                      nprobe=nprobe, gpu=data.op.gpu)
      ret = op(data)
      if not return_distance:
          return ret[0]
      return ret
  