# Buggy code
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):
        # 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
            inp = inp.astype(np.float32, copy=False)
            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)


# 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 _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)