[FRONTEND][Mxnet][nuymp] Adding _npi_advanced_indexing_multiple

[insop]

Adding _npi_advanced_indexing_multiple as discussed in #7186

Need to find a proper mx.sym.np and wanted ask for the reviewers help to find it, so that test case can be valid

@sxjscience , @junrushao1994

Details:

_npi_advanced_indexing_multiple, BART model. This is triggered when we call a[idx1, idx2]. Also see the MXNet-side implementation.

• Example from (BART)[https://raw.githubusercontent.com/dmlc/gluon-nlp/master/src/gluonnlp/models/bart.py]

            batch_indices = mx.npx.arange_like(sequence, axis=0).astype(mx.np.int32)
            outputs = sequence[batch_indices, valid_length - 1]

• Standalone example

  import mxnet as mx 
  from mxnet import use_np 
  from mxnet.gluon import nn  
  from mxnet import np, npx
  
  sequence = np.array([[1, 2,11,12], [3, 4,23,24], [5, 6,35,36]])    
  
  Out[30]: 
  array([[ 1.,  2., 11., 12.],
         [ 3.,  4., 23., 24.],
         [ 5.,  6., 35., 36.]])
  
  
  batch_indices = mx.npx.arange_like(sequence, axis=0).astype(mx.np.int32)
  
  Out[32]: array([0, 1, 2], dtype=int32)
  
  valid_length=2 
  
  outputs = sequence[batch_indices, valid_length - 1]
  Out[35]: array([2., 4., 6.])
  
  
  

• Pytorch advanced indxing example

import torch

    In [44]: a = torch.randn(5, 7, dtype=torch.double)                                                                                                                               

In [45]: a                                                                                                                                                                       
Out[45]: 
tensor([[-1.2230,  0.7823,  0.6655, -0.8564, -0.2611, -0.0423, -0.6728],
        [ 1.6607,  0.9779, -0.2754, -0.7090, -0.3243,  2.2017, -1.7534],
        [-1.9319,  0.5544,  2.0244, -0.8144, -0.2657,  0.7849, -0.4825],
        [ 0.0085,  1.0663,  0.1695, -0.3458, -0.4960,  1.2339,  0.6244],
        [ 0.5265, -2.0689, -0.4739,  0.5544,  0.8612,  0.2270, -2.0888]],
       dtype=torch.float64)

In [46]: t = a[(0,1,2,3,4),1]                                                                                                                                                    

In [47]: t                                                                                                                                                                       
Out[47]: tensor([ 0.7823,  0.9779,  0.5544,  1.0663, -2.0689], dtype=torch.float64)

In [48]: t = a[(0,1,2,3,3),1]                                                                                                                                                    In [49]: t                                                                                                                                                                       
Out[49]: tensor([0.7823, 0.9779, 0.5544, 1.0663, 1.0663], dtype=torch.float64)

[insop]

@junrushao1994, @sxjscience

Any idea what to use in mx.sym.np for the advanced indexing ?

[sxjscience]

You may use this:

import mxnet as mx
mx.npx.set_np()
row_sel = mx.sym.var('row_sel').as_np_ndarray()
col = mx.sym.var('col').as_np_ndarray()
data = mx.sym.var('data').as_np_ndarray()
out = data[row_sel, col]
print(out)

[insop]

Hi @sxjscience

Thank you for the suggestion; however, I still not clear how to create mx_sym so that I can use it for relay.frontend.from_mxnet.

Any suggestion would appreciated.
Thank you.

[sxjscience]

Is it possible to do the following?

row_sel_sym = mx.sym.var('row_sel').as_np_ndarray()
col_sym = mx.sym.var('col').as_np_ndarray()
data_sym = mx.sym.var('data').as_np_ndarray()
mx_sym = data_sym[row_sel_sym, col_sym]

Then, you may pass mx_sym to the relay.frontend.from_mxnet.

[insop]

Hi @sxjscience

Thank you for the suggestion. I think I had tried that as well.
So with the following, I got this exception IndexError: Only integer, slice, or tuple of these types are supported! Received key=(<_Symbol row_sel>, <_Symbol col>) from mxnet/symbol/numpy/_symbol.py:135: (link)

I will dig more, but if you have any suggestion, please let me know.

def test_forward_npi_advanced_indexing_multiple(data_shape, row_sel, col, dtype, target, ctx, kind):
    data_np = np.random.uniform(size=data_shape).astype(dtype)
    ref_res = mx.np.array(data_np)[row_sel, col]

    row_sel_sym = mx.sym.var("row_sel").as_np_ndarray()
    data_sym = mx.sym.var("data").as_np_ndarray()
    col_sym = mx.sym.var("col").as_np_ndarray()
    mx_sym = data_sym[row_sel_sym, col_sym]

    mod, _ = relay.frontend.from_mxnet(
        mx_sym, shape={"data": data_shape, "row_sel": row_sel, "col": col}, dtype=dtype
    )
    intrp = relay.create_executor(kind, mod=mod, ctx=ctx, target=target)
    op_res = intrp.evaluate()(data_np)
    tvm.testing.assert_allclose(op_res.asnumpy(), ref_res.asnumpy(), rtol=1e-5)
    tvm.testing.assert_allclose(ref_res.asnumpy(), ref_res.asnumpy(), rtol=1e-5)

[sxjscience]

Also, I think you will need to install the nightly version of MXNet 2.0. You may follow the guide in https://github.com/dmlc/gluon-nlp.

# Install the version with CUDA 10.1
python3 -m pip install -U --pre "mxnet-cu101>=2.0.0b20201206" -f https://dist.mxnet.io/python

# Install the version with CUDA 10.2
python3 -m pip install -U --pre "mxnet-cu102>=2.0.0b20201206" -f https://dist.mxnet.io/python

# Install the version with CUDA 11
python3 -m pip install -U --pre "mxnet-cu110>=2.0.0b20201206" -f https://dist.mxnet.io/python

# Install the cpu-only version
python3 -m pip install -U --pre "mxnet>=2.0.0b20201206" -f https://dist.mxnet.io/python

[insop]

Thank you, I have updated the MXNet; however, something had happened to my workspace.
I kept getting this, let me try few more things, and get back on the test case.

..
Check failed: it != type_key2index_.end() == false: Cannot find type auto_scheduler.PythonBasedMeasureCallback. Did you forget to register the node by TVM_REGISTER_NODE_TYPE ?

[jroesch]

This PR appears to be out of date, please feel free to reopen it if this is not the case.

As part of the new year we are attempting to triage the project's open pull requests to ensure that code which
is ready for review and/or merging receives adequate attention.

Thanks again for your contribution, and feel free to reach out to discuss these changes.