[MXNET-1382] Add the index_array operator

[nickguletskii]

Description

This pull request implements index_array, an operator that returns an array of indexes of the input array.

For an input array with shape (d_1, d_2, ..., d_n), index_array returns a (d_1, d_2, ..., d_n, n) array idx, where idx[i_1, i_2, ..., i_n, :] = [i_1, i_2, ..., i_n].

Additionally, when the parameter axes is specified, idx will be a (d_1, d_2, ..., d_n, m) array where m is the length of axes, and the following
equality will hold: idx[i_1, i_2, ..., i_n, j] = i_{axes[j]}.

Examples:

x = mx.nd.ones((3, 2))

mx.nd.contrib.index_array(x) = [[[0 0]
                                 [0 1]]

                                [[1 0]
                                 [1 1]]

                                [[2 0]
                                 [2 1]]]

x = mx.nd.ones((3, 2, 2))

mx.nd.contrib.index_array(x, axes=(1, 0)) = [[[[0 0]
                                               [0 0]]

                                              [[1 0]
                                               [1 0]]]


                                             [[[0 1]
                                               [0 1]]

                                              [[1 1]
                                               [1 1]]]

Motivation

This operator can be used to generate meshgrids for tensors without knowing their exact shapes during construction. For instance, this operator can be used to make a makeshift prior box generator for anchor-based computer vision models:

feature_map = F.ones((8, 128, 128, 256)) # N x H x W x C, no shape information when using the Symbol API.
prior_box_stride = 16
box_size=[8, 8]

template = F.squeeze(F.slice_axis(feature_map, begin=0, end=1, axis=-1), axis=-1) # N x H x W
box_centres = F.contrib.index_array(template, axes=(-2, -1, -2, -1)).astype("float32") # N x H x W x 4
box_centres = F.broadcast_mul(box_centres, F.array([prior_box_stride]).reshape((1, 1, 1, 1))) # N x H x W x 4
corner_offsets = F.array(box_size).reshape((1, 1, 1, 2))
corner_offsets = F.concat(-corner_offsets/2, corner_offsets/2, dim=-1)
box_corners = F.broadcast_plus(box_centres, corner_offsets)

Also, this operator can be applied to implement positional encodings for sequence processing, e.g.:

sequence_embeddings = F.ones((65, 8, 256)) # T x N x C, no shape information when using the Symbol API.
template = sequence_embeddings.reshape((0, 0, -1, 2)) # T x N x C -> T x N x (C/2) x 2
pos, i = F.split(
    F.contrib.index_array(template, axes=(0, 2)).astype("float32"), # T x N x (C/2) x 2 x 2
    axis=-1,
    num_outputs=2,
    squeeze_axis=True
) # T x N x (C/2) x 2 and T x N x (C/2) x 2
base = F.ones((1, 1, 1, 1)) * 10000
dmodel = F.slice_axis(F.shape_array(sequence_embeddings), begin=-1, end=None, axis=0)
dmodel = dmodel.reshape((1, 1, 1, 1)).astype("float32")
tmp = F.broadcast_div(pos, F.broadcast_power(base, F.broadcast_div(2 * i,  dmodel))) # T x N x (C/2) x 2
sin_input, cos_input = F.split(tmp, axis=-1, num_outputs=2, squeeze_axis=True) # T x N x (C/2) and T x N x (C/2)
positional_encoding = F.stack(F.sin(sin_input), F.cos(cos_input), axis=-1).reshape((0, 0, -3)) # T x N x C

I've also encountered situations where this operator would have been useful for some indexing tricks.

Checklist

Essentials

Please feel free to remove inapplicable items for your PR.

• The PR title starts with [MXNET-$JIRA_ID], where $JIRA_ID refers to the relevant JIRA issue created (except PRs with tiny changes)
• Changes are complete (i.e. I finished coding on this PR)
• All changes have test coverage:
• Unit tests are added for small changes to verify correctness (e.g. adding a new operator)
• Nightly tests are added for complicated/long-running ones (e.g. changing distributed kvstore)
• Build tests will be added for build configuration changes (e.g. adding a new build option with NCCL)
• Code is well-documented:
• For user-facing API changes, API doc string has been updated.
• For new C++ functions in header files, their functionalities and arguments are documented.
• For new examples, README.md is added to explain the what the example does, the source of the dataset, expected performance on test set and reference to the original paper if applicable
• Check the API doc at http:https://mxnet-ci-doc.s3-accelerate.dualstack.amazonaws.com/PR-$PR_ID/$BUILD_ID/index.html
• To the my best knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change

Changes

• The IndexArray operator (for the CPU)
• The IndexArray operator (for the GPU)
• Tests for the CPU implementation
• Tests for the GPU implementation
• Entries on the Python API reference pages

Comments

• This operator always returns kInt64.
• The new tests in test_operator_gpu.py are exactly the same as in test_operator.py. I couldn't find any evidence that test_operator.py ever gets called with a GPU default_context, so I copied the tests into test_operator_gpu.py to make sure that the GPU implementation works too.

[piyushghai]

Thanks for your contributions @nickguletskii. Can you also look into the CI failures ?
@mxnet-label-bot Add [Opeartor, pr-awaiting-review]

[piyushghai]

@mxnet-label-bot Update [Operator, pr-awaiting-review]

[nickguletskii]

The CI failure seems to be unrelated. Here's the relevant part of the log:

[2019-04-03 23:57:30   ERROR] Cuda initialization failure with error 35. Please check cuda installation:  http:https://docs.nvidia.com/cuda/cuda-installation-guide-linux/index.html.

terminate called after throwing an instance of 'std::runtime_error'

  what():  Failed to create object

If I am not mistaken, this error code corresponds to the following error message: "CUDA driver version is insufficient for CUDA runtime version". Could this be caused by the recent update to the nvidia/cuda:10.0-cudnn7-devel Docker tag? The TensorRT job uses nvidia/cuda:10.0-cudnn7-devel, while other jobs use nvidia/cuda:9.1-cudnn7-devel (e.g. the ubuntu_gpu build), which may be why the other jobs are not failing.

[fhieber]

Thanks for implementing this! This seems similar to what I requested in #14574, correct?

Any reason this should be in the contrib namespace?

[nickguletskii]

@fhieber In a way, yes. The behaviour you have described in #14574 can be emulated by applying reshape, index_array, broadcast_add, and finally, reshape_list. I was thinking of implementing arange_like instead of index_array, but I don't find the arange operator itself versatile enough and I didn't want to make arange_like too different from arange.

There's no particular reason why I've placed index_array into the contrib namespace other than the fact that it's a new operator. I'd be happy to update this pull request and move index_array out of the contrib namespace if the community deems it desirable.

[fhieber]

Thanks @nickguletskii , I agree, your operator seems more versatile, thanks for working on it!

contrib is fine with me, I just asked out of curiosity on whether the preferred or standard way of introducing new operators is to always place in contrib first.

[nickguletskii]

@piyushghai Sorry to bother you, but should I rebase this pull request on master now that the failing job (test_unix_python3_tensorrt_gpu) has been disabled (#14642), or should I just leave it as it is and wait for the review?

[piyushghai]

@nickguletskii You should rebase with the latest master branch. That'll enable you to push again, allowing the CI to pass on your PR.
The reviews will happen on it in the meanwhile :)

[nickguletskii]

Huh, this is weird, I've updated this PR to be compatible with the recently merged support for zero-dim and zero-shape arrays, and now the CI for Windows CPU tests are failing. However, when I try to run the tests locally (on a 64 bit Windows machine), the tests pass just fine. I'll make further attempts to reproduce this issue on my local machine tomorrow. I suspect that I've screwed up with memory access somewhere because the test runner crashes without providing any useful debugging information.

[nickguletskii]

The previous CI failures on Windows were caused by indexing C-style arrays using ints. I'm not sure whether this was just undefined behaviour or a bug in the older versions of MSVC++, but it did not happen with VC++ 16.0. I fixed the issue by using ptrdiff_t to index C-style arrays.

[nickguletskii]

@szha @reminisce This PR should be ready for a review now. The CI failure seems to be unrelated (the failing test is test_operator_gpu.test_convolution_multiple_streams, described in #14329).

[vandanavk]

@szha @reminisce @fhieber Is this PR good to go?

[reminisce]

Thank you for your contribution and sorry for the late review. Have some comments on code cosmetics.

[reminisce]

The tests defined in test_operator.py will be run with ctx=mx.gpu() in CI. There is no need to duplicate test code in this file.

[nickguletskii]

Sorry, I've removed these redundant tests in 3cdfcf5.

[reminisce]

If the whole function is expected to be numpy compatible, you can use a decorator mx.use_np_compat to wrap the whole function. It would be easier for us to just delete the decorator in the next major release than putting additional efforts of adjusting indentation.

[nickguletskii]

Thanks, I've updated the subtests to use @mx.use_np_compat in 139567b.

[reminisce]

Same as above, please use the decorator mx.use_np_compat.

[reminisce]

Same as above.

[karan6181]

@nickguletskii Could you please re-trigger the CI again?
@szha Could you please review this PR? Thanks!