# Buggy code
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

    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_sample_train:
        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),),)
    )


# Traceback (AssertionError)
In [1]: from sklearn.datasets import make_classification

In [2]: x, y = make_classification()

In [3]: import mars.tensor as mt
/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
return f(*args, **kwds)
/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
return f(*args, **kwds)

In [4]: x = mt.tensor(x, chunk_size=20)

In [5]: x.shape
Out[5]: (100, 20)

In [6]: y = mt.tensor(y, chunk_size=20)

In [7]: y.shape
Out[7]: (100,)

/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
return f(*args, **kwds)
/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
return f(*args, **kwds)
/Users/qinxuye/miniconda3/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject
return f(*args, **kwds)
In [8]: from mars.learn.neighbors._faiss import build_faiss_index

In [39]: index = build_faiss_index(x, index_name='PCAR6,IVF8_HNSW32,SQ8', n_samp
...: le=10)

In [40]: index.execute()
WARNING clustering 10 points to 8 centroids: please provide at least 312 training points
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-40-bd4069985a62> in <module>
----> 1 index.execute()

~/Workspace/mars/mars/core.py in execute(self, session, **kw)
373
374         if wait:
--> 375             return run()
376         else:
377             # leverage ThreadPoolExecutor to submit task,

~/Workspace/mars/mars/core.py in run()
368         def run():
369             # no more fetch, thus just fire run
--> 370             session.run(self, **kw)
371             # return Tileable or ExecutableTuple itself
372             return self

~/Workspace/mars/mars/session.py in run(self, *tileables, **kw)
498         tileables = tuple(mt.tensor(t) if not isinstance(t, (Entity, Base)) else t
499                           for t in tileables)
--> 500         result = self._sess.run(*tileables, **kw)
501
502         for t in tileables:

~/Workspace/mars/mars/session.py in run(self, *tileables, **kw)
106             # set number of running cores
107             self.context.set_ncores(kw['n_parallel'])
--> 108             res = self._executor.execute_tileables(tileables, **kw)
109             return res
110

~/Workspace/mars/mars/utils.py in _inner(*args, **kwargs)
437         def _inner(*args, **kwargs):
438             with self:
--> 439                 return func(*args, **kwargs)
440
441         return _inner

~/Workspace/mars/mars/executor.py in execute_tileables(self, tileables, fetch, n_parallel, n_thread, print_progress, mock, compose, name)
879                                    n_parallel=n_parallel or n_thread,
880                                    print_progress=print_progress, mock=mock,
--> 881                                    chunk_result=chunk_result)
882
883                 # update shape of tileable and its chunks whatever it's successful or not

~/Workspace/mars/mars/executor.py in execute_graph(self, graph, keys, n_parallel, print_progress, mock, no_intermediate, compose, retval, chunk_result)
691             print_progress=print_progress, mock=mock, mock_max_memory=self._mock_max_memory,
692             fetch_keys=fetch_keys, no_intermediate=no_intermediate)
--> 693         res = graph_execution.execute(retval)
694         self._mock_max_memory = max(self._mock_max_memory, graph_execution._mock_max_memory)
695         if mock:

~/Workspace/mars/mars/executor.py in execute(self, retval)
572         # wait until all the futures completed
573         for future in executed_futures:
--> 574             future.result()
575
576         if retval:

~/miniconda3/lib/python3.7/concurrent/futures/_base.py in result(self, timeout)
433                 raise CancelledError()
434             elif self._state == FINISHED:
--> 435                 return self.__get_result()
436             else:
437                 raise TimeoutError()

~/miniconda3/lib/python3.7/concurrent/futures/_base.py in __get_result(self)
382     def __get_result(self):
383         if self._exception:
--> 384             raise self._exception
385         else:
386             return self._result

~/miniconda3/lib/python3.7/concurrent/futures/thread.py in run(self)
55
56         try:
---> 57             result = self.fn(*self.args, **self.kwargs)
58         except BaseException as exc:
59             self.future.set_exception(exc)

~/Workspace/mars/mars/utils.py in _inner(*args, **kwargs)
437         def _inner(*args, **kwargs):
438             with self:
--> 439                 return func(*args, **kwargs)
440
441         return _inner

~/Workspace/mars/mars/executor.py in _execute_operand(self, op)
444             # so we pass the first operand's first output to Executor.handle
445             first_op = ops[0]
--> 446             self.handle_op(first_op, results, self._mock)
447
448             # update maximal memory usage during execution

~/Workspace/mars/mars/executor.py in handle_op(self, *args, **kw)
376
377     def handle_op(self, *args, **kw):
--> 378         return Executor.handle(*args, **kw)
379
380     def _order_starts(self):

~/Workspace/mars/mars/executor.py in handle(cls, op, results, mock)
642                 # The `UFuncTypeError` was introduced by numpy#12593 since v1.17.0.
643                 try:
--> 644                     return runner(results, op)
645                 except UFuncTypeError as e:
646                     raise TypeError(str(e)).with_traceback(sys.exc_info()[2]) from None

~/Workspace/mars/mars/learn/neighbors/_faiss.py in execute(cls, ctx, op)
340             cls._execute_map(ctx, op)
341         elif op.stage == OperandStage.agg:
--> 342             cls._execute_agg(ctx, op)
343         else:
344             assert op.stage is None

~/Workspace/mars/mars/learn/neighbors/_faiss.py in _execute_agg(cls, ctx, op)
327                 index = _load_index(ctx, op, index, device_id)
328                 indexes.append(index)
--> 329                 assert hasattr(index, 'merge_from')
330                 if merged_index is None:
331                     merged_index = index

AssertionError:

# Fixed code
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),),)
    )