haojin2 commented on a change in pull request #17553: Operator
diag_indices_from [Numpy]
URL: https://github.com/apache/incubator-mxnet/pull/17553#discussion_r376742177
##########
File path: python/mxnet/symbol/numpy/_symbol.py
##########
@@ -4410,6 +4410,45 @@ def unravel_index(indices, shape, order='C'): # pylint:
disable=redefined-outer-
raise NotImplementedError('Don not support column-major
(Fortran-style) order at this moment')
+def diag_indices_from(arr):
+ """
+ This returns a tuple of indices that can be used to access the main
diagonal of an array
+ a with a.ndim >= 2 dimensions and shape (n, n, …, n). For a.ndim = 2 this
is
+ the usual diagonal, for a.ndim > 2 this is the set of indices to access
+ a[i, i, ..., i] for i = [0..n-1].
+
+ Parameters:
+ -------------
+ arr : _Symbol
+ Input array for acessing the main diagonal. All dimensions
+ should have equal length.
+
+ Return:
+ -------------
+ diag: _Symbol
+ indices of the main diagonal.
+
+ Examples:
+ -------------
+ >>> a = np.arange(16).reshape(4, 4)
+ >>> a
+ array([[ 0, 1, 2, 3],
+ [ 4, 5, 6, 7],
+ [ 8, 9, 10, 11],
+ [12, 13, 14, 15]])
+ >>> idx = np.diag_indices_from(a)
+ >>> idx
+ (array([0, 1, 2, 3]), array([0, 1, 2, 3]))
+ >>> a[idx] = 100
+ >>> a
+ array([[100, 1, 2, 3],
+ [ 4, 100, 6, 7],
+ [ 8, 9, 100, 11],
+ [ 12, 13, 14, 100]])
+ """
+ return _npi.diag_indices_from(arr)
Review comment:
@reminisce is there any way to specify the output to be a tuple here?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
