@samskalicky The base is equal to the size of the diagonal. The output is reshaped to `oshape + (base,)`, and the input is reshaped to `ishape`. `idx` is therefore the part corresponding to `oshape` and `ishape`. And it corresponds to the mth element in the diagonal, where `m = i % base = I - idx * base`. We need to multiply it by `stride`, which is basically the sum of the strides of `axis1` and `axis2` in order compute the index of the element we want. However, because the diagonal we want is not necessarily the main diagonal (only when `k = 0`), we need to add to the result an offset.
[ Full content available at: https://github.com/apache/incubator-mxnet/pull/12430 ] This message was relayed via gitbox.apache.org for [email protected]
