From 05c7d15d5993415eedd0a4fafbb75d16d3cda8b3 Mon Sep 17 00:00:00 2001 From: barry-jin Date: Fri, 19 Mar 2021 13:15:21 -0700 Subject: [PATCH] make stack use faster API --- python/mxnet/numpy/multiarray.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/python/mxnet/numpy/multiarray.py b/python/mxnet/numpy/multiarray.py index 0f55c318d950..ab4447c90cae 100644 --- a/python/mxnet/numpy/multiarray.py +++ b/python/mxnet/numpy/multiarray.py @@ -499,7 +499,7 @@ def _get_np_advanced_indexing(self, key): if type(idcs) == NDArray: # pylint: disable=unidiomatic-typecheck idcs = idcs.as_np_ndarray() else: - idcs = _npi.stack(*[i if isinstance(i, self.__class__) else i.as_np_ndarray() for i in idcs]) + idcs = _mx_nd_np.stack([i if isinstance(i, self.__class__) else i.as_np_ndarray() for i in idcs]) sliced = _npi.gather_nd(self, idcs) # Reshape due to `None` entries in `key`. if new_axes: @@ -516,7 +516,7 @@ def _set_np_advanced_indexing(self, key, value): if type(idcs) == NDArray: # pylint: disable=unidiomatic-typecheck idcs = idcs.as_np_ndarray() else: - idcs = _npi.stack(*[i if isinstance(i, self.__class__) else i.as_np_ndarray() for i in idcs]) + idcs = _mx_nd_np.stack([i if isinstance(i, self.__class__) else i.as_np_ndarray() for i in idcs]) vshape = get_oshape_of_gather_nd_op(self.shape, idcs.shape) value_nd = self._prepare_value_nd(value, bcast_shape=vshape, squeeze_axes=new_axes) self._scatter_set_nd(value_nd, idcs) @@ -762,7 +762,7 @@ def __getitem__(self, key): all((isinstance(arr, NDArray) and _np.issubdtype(arr.dtype, _np.integer) and \ arr.ndim > 0) for arr in key): # Equivalent case in numpy/_symbol.py - return _npi.advanced_indexing_multiple(self, _npi.stack(*key)) + return _npi.advanced_indexing_multiple(self, _mx_nd_np.stack(key)) elif isinstance(key, tuple) and dc.is_deferred_compute(): # Equivalent to isinstance(key, tuple) case in numpy/_symbol.py # Only enabled in deferred compute mode, as this codepath prevents