Add interleaved_matmul_* to npx namespace (#20375)

barry-jin · web-flow · commit c68dffe11f50 · 2021-07-06T10:41:40.000-04:00
diff --git a/src/operator/contrib/transformer.cc b/src/operator/contrib/transformer.cc
@@ -648,6 +648,7 @@ void BackwardInterleavedMatMulEncDecValAttCPU(const nnvm::NodeAttrs& attrs,
 }
 
 NNVM_REGISTER_OP(_contrib_interleaved_matmul_selfatt_qk)
+.add_alias("_npx_interleaved_matmul_selfatt_qk")
 .describe(R"code(Compute the matrix multiplication between the projections of
 queries and keys in multihead attention use as self attention.
 
@@ -691,6 +692,7 @@ NNVM_REGISTER_OP(_backward_interleaved_matmul_selfatt_qk)
 .set_attr<FCompute>("FCompute<cpu>", BackwardInterleavedMatMulSelfAttQKCPU);
 
 NNVM_REGISTER_OP(_contrib_interleaved_matmul_selfatt_valatt)
+.add_alias("_npx_interleaved_matmul_selfatt_valatt")
 .describe(R"code(Compute the matrix multiplication between the projections of
 values and the attention weights in multihead attention use as self attention.
 
@@ -738,6 +740,7 @@ NNVM_REGISTER_OP(_backward_interleaved_matmul_selfatt_valatt)
 .set_attr<FCompute>("FCompute<cpu>", BackwardInterleavedMatMulSelfAttValAttCPU);
 
 NNVM_REGISTER_OP(_contrib_interleaved_matmul_encdec_qk)
+.add_alias("_npx_interleaved_matmul_encdec_qk")
 .describe(R"code(Compute the matrix multiplication between the projections of
 queries and keys in multihead attention use as encoder-decoder.
 
@@ -784,6 +787,7 @@ NNVM_REGISTER_OP(_backward_interleaved_matmul_encdec_qk)
 .set_attr<FCompute>("FCompute<cpu>", BackwardInterleavedMatMulEncDecQKCPU);
 
 NNVM_REGISTER_OP(_contrib_interleaved_matmul_encdec_valatt)
+.add_alias("_npx_interleaved_matmul_encdec_valatt")
 .describe(R"code(Compute the matrix multiplication between the projections of
 values and the attention weights in multihead attention use as encoder-decoder.
 
diff --git a/src/operator/tensor/elemwise_sum.cc b/src/operator/tensor/elemwise_sum.cc
@@ -139,6 +139,7 @@ NNVM_REGISTER_OP(add_n)
 MXNET_ADD_SPARSE_OP_ALIAS(add_n)
 MXNET_ADD_SPARSE_OP_ALIAS(ElementWiseSum)
 .add_alias("ElementWiseSum")
+.add_alias("_npx_add_n")
 .describe(R"doc(Adds all input arguments element-wise.
 
 .. math::
diff --git a/src/operator/tensor/matrix_op.cc b/src/operator/tensor/matrix_op.cc
@@ -612,6 +612,7 @@ NNVM_REGISTER_OP(_backward_slice_axis)
 .set_attr<FCompute>("FCompute<cpu>", SliceAxisGrad_<cpu>);
 
 NNVM_REGISTER_OP(slice_like)
+.add_alias("_npx_slice_like")
 .describe(R"code(Slices a region of the array like the shape of another array.
 This function is similar to ``slice``, however, the `begin` are always `0`s
 and `end` of specific axes are inferred from the second input `shape_like`.
diff --git a/tests/python/gpu/test_numpy_op.py b/tests/python/gpu/test_numpy_op.py
@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 
 import sys
-import numpy as _np
+import numpy as onp
 import pytest
 import mxnet as mx
 from mxnet import np
@@ -45,87 +45,87 @@ def dbg(name, data):
         print('{} = {}'.format(name, data))
 
     configs = [
-        ('ii', [(5, 5)], lambda *args: (_np.eye(5),)),
-        ('ii->i', [(5, 5)], lambda *args: (_np.eye(5),)),
-        ('ij->i', [(5, 5)], lambda *args: (_np.ones((5, 5)),)),
-        ('...j->...', [(5, 5)], lambda *args: (_np.ones((5, 5)),)),
-        ('ji', [(2, 3)], lambda *args: (_np.ones((2, 3)),)),
-        ('ij->ji', [(2, 3)], lambda *args: (_np.ones((2, 3)),)),
-        ('ij, jk', [(5, 0), (0, 4)], lambda *args: (_np.empty((5, 0)), _np.empty((0, 4)))),
+        ('ii', [(5, 5)], lambda *args: (onp.eye(5),)),
+        ('ii->i', [(5, 5)], lambda *args: (onp.eye(5),)),
+        ('ij->i', [(5, 5)], lambda *args: (onp.ones((5, 5)),)),
+        ('...j->...', [(5, 5)], lambda *args: (onp.ones((5, 5)),)),
+        ('ji', [(2, 3)], lambda *args: (onp.ones((2, 3)),)),
+        ('ij->ji', [(2, 3)], lambda *args: (onp.ones((2, 3)),)),
+        ('ij, jk', [(5, 0), (0, 4)], lambda *args: (onp.empty((5, 0)), onp.empty((0, 4)))),
 
         ('i, i', [(5,), (5,)], lambda *args: (args[1], args[0])),
-        ('ij, j', [(5, 5), (5,)], lambda *args: (_np.tile(args[1][None, :], [5, 1]),
+        ('ij, j', [(5, 5), (5,)], lambda *args: (onp.tile(args[1][None, :], [5, 1]),
                                                  args[0].sum(axis=0))),
-        ('...j, j', [(5, 5), (5,)], lambda *args: (_np.tile(args[1][None, :], [5, 1]),
-                                                   _np.sum(args[0], axis=0))),
-        ('..., ...', [(), (2, 3)], lambda *args: (_np.sum(args[1], axis=None),
-                                                  args[0] * _np.ones((2, 3)))),
-        (', ij', [(), (2, 3)], lambda *args: (_np.sum(args[1], axis=None),
-                                              args[0] * _np.ones((2, 3)))),
-        ('i, j', [(2,), (5, )], lambda *args: (_np.sum(args[1], axis=None) * _np.ones(2),
-                                               _np.sum(args[0], axis=None) * _np.ones(5))),
-        ('ijk, jil->kl', [(3, 4, 5), (4, 3, 2)], lambda *args: (_np.tile(_np.transpose(_np.sum(args[1],
+        ('...j, j', [(5, 5), (5,)], lambda *args: (onp.tile(args[1][None, :], [5, 1]),
+                                                   onp.sum(args[0], axis=0))),
+        ('..., ...', [(), (2, 3)], lambda *args: (onp.sum(args[1], axis=None),
+                                                  args[0] * onp.ones((2, 3)))),
+        (', ij', [(), (2, 3)], lambda *args: (onp.sum(args[1], axis=None),
+                                              args[0] * onp.ones((2, 3)))),
+        ('i, j', [(2,), (5, )], lambda *args: (onp.sum(args[1], axis=None) * onp.ones(2),
+                                               onp.sum(args[0], axis=None) * onp.ones(5))),
+        ('ijk, jil->kl', [(3, 4, 5), (4, 3, 2)], lambda *args: (onp.tile(onp.transpose(onp.sum(args[1],
                                                                   axis=-1))[:, :, None], [1, 1, 5]),
-                                                                _np.tile(_np.transpose(_np.sum(args[0],
+                                                                onp.tile(onp.transpose(onp.sum(args[0],
                                                                   axis=-1))[:, :, None], [1, 1, 2]))),
-        ('ijk, jil->kl', [(33, 44, 55), (44, 33, 22)], lambda *args: (_np.tile(_np.transpose(_np.sum(args[1],
+        ('ijk, jil->kl', [(33, 44, 55), (44, 33, 22)], lambda *args: (onp.tile(onp.transpose(onp.sum(args[1],
                                                                   axis=-1))[:, :, None], [1, 1, 55]),
-                                                                _np.tile(_np.transpose(_np.sum(args[0],
+                                                                onp.tile(onp.transpose(onp.sum(args[0],
                                                                   axis=-1))[:, :, None], [1, 1, 22]))),
-        ('ki, jk->ij', [(3, 2), (4, 3)], lambda *args: (_np.tile(args[1].sum(axis=0)[:, None], [1, 2]),
-                                                        _np.tile(args[0].sum(axis=1)[None, :], [4, 1]))),
-        ('ki, ...k->i...', [(3, 2), (4, 3)], lambda *args: (_np.tile(args[1].sum(axis=0)[:, None], [1, 2]),
-                                                            _np.tile(args[0].sum(axis=1)[None, :], [4, 1]))),
-        ('k..., jk', [(3, 2), (4, 3)], lambda *args: (_np.tile(args[1].sum(axis=0)[:, None], [1, 2]),
-                                                      _np.tile(args[0].sum(axis=1)[None, :], [4, 1]))),
+        ('ki, jk->ij', [(3, 2), (4, 3)], lambda *args: (onp.tile(args[1].sum(axis=0)[:, None], [1, 2]),
+                                                        onp.tile(args[0].sum(axis=1)[None, :], [4, 1]))),
+        ('ki, ...k->i...', [(3, 2), (4, 3)], lambda *args: (onp.tile(args[1].sum(axis=0)[:, None], [1, 2]),
+                                                            onp.tile(args[0].sum(axis=1)[None, :], [4, 1]))),
+        ('k..., jk', [(3, 2), (4, 3)], lambda *args: (onp.tile(args[1].sum(axis=0)[:, None], [1, 2]),
+                                                      onp.tile(args[0].sum(axis=1)[None, :], [4, 1]))),
         (('ij,jk'), [(2, 5), (5, 2)],
-            lambda *args: (_np.dot(_np.ones((2, 2)), args[1].T),
-            _np.dot(args[0].T, _np.ones((2, 2))))),
+            lambda *args: (onp.dot(onp.ones((2, 2)), args[1].T),
+            onp.dot(args[0].T, onp.ones((2, 2))))),
         (('ij,jk,kl'), [(2, 2), (2, 5), (5, 2)],
-            lambda *args: (_np.dot(_np.ones((2, 2)), _np.dot(args[1], args[2]).T),
-            _np.dot(args[0].T, _np.dot(_np.ones((2, 2)), args[2].T)),
-            _np.dot(_np.dot(args[0], args[1]).T, _np.ones((2, 2))))),
+            lambda *args: (onp.dot(onp.ones((2, 2)), onp.dot(args[1], args[2]).T),
+            onp.dot(args[0].T, onp.dot(onp.ones((2, 2)), args[2].T)),
+            onp.dot(onp.dot(args[0], args[1]).T, onp.ones((2, 2))))),
         (('ij,jk,kl->il'), [(2, 2), (2, 5), (5, 2)],
-            lambda *args: (_np.dot(_np.ones((2, 2)), _np.dot(args[1], args[2]).T),
-            _np.dot(args[0].T, _np.dot(_np.ones((2, 2)), args[2].T)),
-            _np.dot(_np.dot(args[0], args[1]).T, _np.ones((2, 2))))),
+            lambda *args: (onp.dot(onp.ones((2, 2)), onp.dot(args[1], args[2]).T),
+            onp.dot(args[0].T, onp.dot(onp.ones((2, 2)), args[2].T)),
+            onp.dot(onp.dot(args[0], args[1]).T, onp.ones((2, 2))))),
         (('ij,jk,kl->il'), [(67, 89), (89, 55), (55, 99)],
-            lambda *args: (_np.dot(_np.ones((67, 99)), _np.dot(args[1], args[2]).T),
-            _np.dot(args[0].T, _np.dot(_np.ones((67, 99)), args[2].T)),
-            _np.dot(_np.dot(args[0], args[1]).T, _np.ones((67, 99))))),
+            lambda *args: (onp.dot(onp.ones((67, 99)), onp.dot(args[1], args[2]).T),
+            onp.dot(args[0].T, onp.dot(onp.ones((67, 99)), args[2].T)),
+            onp.dot(onp.dot(args[0], args[1]).T, onp.ones((67, 99))))),
         (('ij,jk,kl, lm->im'), [(12, 54), (54, 32), (32, 45), (45, 67)],
-            lambda *args: (_np.dot(_np.ones((12, 67)), _np.dot(args[1], _np.dot(args[2], args[3])).T),
-            _np.dot(args[0].T, _np.dot(_np.ones((12, 67)), _np.dot(args[2], args[3]).T)),
-            _np.dot(_np.dot(args[0], args[1]).T, _np.dot(_np.ones((12, 67)), args[3].T)),
-            _np.dot(_np.dot(args[0], _np.dot(args[1], args[2])).T, _np.ones((12, 67))))),
+            lambda *args: (onp.dot(onp.ones((12, 67)), onp.dot(args[1], onp.dot(args[2], args[3])).T),
+            onp.dot(args[0].T, onp.dot(onp.ones((12, 67)), onp.dot(args[2], args[3]).T)),
+            onp.dot(onp.dot(args[0], args[1]).T, onp.dot(onp.ones((12, 67)), args[3].T)),
+            onp.dot(onp.dot(args[0], onp.dot(args[1], args[2])).T, onp.ones((12, 67))))),
 
         # broadcast axis
-        ('ij, ij -> i', [(1, 4), (2, 4)], lambda *args: (_np.sum(args[1], axis=0)[None, :],
-                                                         _np.tile(args[0], [2, 1]))),
+        ('ij, ij -> i', [(1, 4), (2, 4)], lambda *args: (onp.sum(args[1], axis=0)[None, :],
+                                                         onp.tile(args[0], [2, 1]))),
         ('...ij, ...jk -> ...ik', [(1, 4), (4, 2)], lambda *args: (args[1].sum(axis=1)[None, :],
-                                                                   _np.tile(args[0].sum(axis=0)[: ,None], [1, 2]))),
-        ('...ij, ...jk -> ...ik', [(2, 4), (4, 2)], lambda *args: (_np.tile(args[1].sum(axis=1)[None, :], [2, 1]),
-                                                                   _np.tile(args[0].sum(axis=0)[: ,None], [1, 2]))),
+                                                                   onp.tile(args[0].sum(axis=0)[: ,None], [1, 2]))),
+        ('...ij, ...jk -> ...ik', [(2, 4), (4, 2)], lambda *args: (onp.tile(args[1].sum(axis=1)[None, :], [2, 1]),
+                                                                   onp.tile(args[0].sum(axis=0)[: ,None], [1, 2]))),
         ('...ij, ...jk -> ...ik', [(3, 2, 1, 4), (3, 2, 4, 2)], lambda *args: (
                                                             args[1].sum(axis=3)[:, :, None, :],
-                                                            _np.tile(args[0].sum(axis=2)[:, :, :, None], [1, 1, 1, 2]))),
+                                                            onp.tile(args[0].sum(axis=2)[:, :, :, None], [1, 1, 1, 2]))),
         ('...ij, ...ik -> ...jk', [(1, 1, 1, 4), (1, 1, 1, 3)], lambda *args: (
-                                                            _np.tile(args[1].sum(axis=3)[:, :, :, None], [1, 1, 1, 4]),
-                                                            _np.tile(args[0].sum(axis=3)[:, :, : ,None], [1, 1, 1, 3]))),
+                                                            onp.tile(args[1].sum(axis=3)[:, :, :, None], [1, 1, 1, 4]),
+                                                            onp.tile(args[0].sum(axis=3)[:, :, : ,None], [1, 1, 1, 3]))),
         ('...ij, ...jc -> ...ic', [(1, 1, 5, 3), (1, 1, 3, 2)], lambda *args: (
-                                                            _np.tile(args[1].sum(axis=3)[:, :, None, :], [1, 1, 5, 1]),
-                                                            _np.tile(args[0].sum(axis=2)[:, :, : ,None], [1, 1, 1, 2]))),
+                                                            onp.tile(args[1].sum(axis=3)[:, :, None, :], [1, 1, 5, 1]),
+                                                            onp.tile(args[0].sum(axis=2)[:, :, : ,None], [1, 1, 1, 2]))),
         ('...ij, ...jc -> ...ic', [(1, 2, 5, 4), (1, 2, 4, 2)], lambda *args: (
-                                                            _np.tile(args[1].sum(axis=3)[:, :, None, :], [1, 1, 5, 1]),
-                                                            _np.tile(args[0].sum(axis=2)[:, :, : ,None], [1, 1, 1, 2]))),
+                                                            onp.tile(args[1].sum(axis=3)[:, :, None, :], [1, 1, 5, 1]),
+                                                            onp.tile(args[0].sum(axis=2)[:, :, : ,None], [1, 1, 1, 2]))),
         ('...ij, ...jc -> ...ic', [(2, 1, 5, 4), (2, 1, 4, 2)], lambda *args: (
-                                                            _np.tile(args[1].sum(axis=3)[:, :, None, :], [1, 1, 5, 1]),
-                                                             _np.tile(args[0].sum(axis=2)[:, :, : ,None], [1, 1, 1, 2]))),
+                                                            onp.tile(args[1].sum(axis=3)[:, :, None, :], [1, 1, 5, 1]),
+                                                             onp.tile(args[0].sum(axis=2)[:, :, : ,None], [1, 1, 1, 2]))),
         # test with cuTensor using workspace
         (('ij,jk,kl->il'), [(64, 200), (200, 64), (64, 64)],
-            lambda *args: (_np.dot(_np.ones((64, 64)), _np.dot(args[1], args[2]).T),
-            _np.dot(args[0].T, _np.dot(_np.ones((64, 64)), args[2].T)),
-            _np.dot(_np.dot(args[0], args[1]).T, _np.ones((64, 64)))))
+            lambda *args: (onp.dot(onp.ones((64, 64)), onp.dot(args[1], args[2]).T),
+            onp.dot(args[0].T, onp.dot(onp.ones((64, 64)), args[2].T)),
+            onp.dot(onp.dot(args[0], args[1]).T, onp.ones((64, 64)))))
     ]
 
     dtypes = ['float16', 'float32', 'float64', 'int32']
@@ -144,11 +144,11 @@ def dbg(name, data):
                             x = []
                             x_np = []
                             for shape in operands:
-                                tmp = _np.array(_np.random.uniform(-0.3, 0.3, shape), dtype=dtype)
+                                tmp = onp.array(onp.random.uniform(-0.3, 0.3, shape), dtype=dtype)
                                 x_np.append(tmp)
                                 x.append(np.array(tmp, dtype=dtype))
                                 x[-1].attach_grad()
-                            expected_np = _np.einsum(subscripts, *x_np, optimize=False, dtype=dtype).astype(dtype)
+                            expected_np = onp.einsum(subscripts, *x_np, optimize=False, dtype=dtype).astype(dtype)
                             with mx.autograd.record():
                                 out_mx = test_einsum(*x)
                             assert out_mx.shape == expected_np.shape
diff --git a/tests/python/unittest/test_numpy_op.py b/tests/python/unittest/test_numpy_op.py
diff --git a/tests/python/unittest/test_subgraph_op.py b/tests/python/unittest/test_subgraph_op.py
