add numpy op moveaxis

fix code style

python/mxnet/_numpy_op_doc.py

@@ -103,3 +103,50 @@ def _np_cumsum(a, axis=None, dtype=None, out=None):
 
     """
     pass
+
+
+def moveaxis(a, source, destination):
+    """Move axes of an array to new positions.
+
+    Other axes remain in their original order.
+
+    Parameters
+    ----------
+    a : ndarray
+    The array whose axes should be reordered.
+    source : int or sequence of int
+    Original positions of the axes to move. These must be unique.
+    destination : int or sequence of int
+    Destination positions for each of the original axes. These must also be
+    unique.
+
+    Returns
+    -------
+    result : ndarray
+    Array with moved axes. This array is a view of the input array.
+
+    See Also
+    --------
+    transpose: Permute the dimensions of an array.
+    swapaxes: Interchange two axes of an array.
+    Examples
+    --------
+
+    >>> x = np.zeros((3, 4, 5))
+    >>> np.moveaxis(x, 0, -1).shape
+    (4, 5, 3)
+    >>> np.moveaxis(x, -1, 0).shape
+    (5, 3, 4)
+
+    These all achieve the same result:
+
+    >>> np.transpose(x).shape
+    (5, 4, 3)
+    >>> np.swapaxes(x, 0, -1).shape
+    (5, 4, 3)
+    >>> np.moveaxis(x, [0, 1], [-1, -2]).shape
+    (5, 4, 3)
+    >>> np.moveaxis(x, [0, 1, 2], [-1, -2, -3]).shape
+    (5, 4, 3)
+    """
+    pass

src/operator/numpy/np_matrix_op-inl.h

@@ -60,6 +60,82 @@ void NumpyTranspose(const nnvm::NodeAttrs& attrs,
   }
 }
 
+struct NumpyMoveaxisParam : public dmlc::Parameter<NumpyMoveaxisParam> {
+  mxnet::TShape source;
+  mxnet::TShape destination;
+  DMLC_DECLARE_PARAMETER(NumpyMoveaxisParam) {
+    DMLC_DECLARE_FIELD(source)
+    .describe("Original positions of the axes to move. These must be unique.");
+    DMLC_DECLARE_FIELD(destination)
+    .describe("Destination positions for each of the original axes. "
+              "These must also be unique.");
+  }
+};
+
+template<typename xpu>
+void NumpyMoveaxisCompute(const nnvm::NodeAttrs& attrs,
+                          const OpContext& ctx,
+                          const std::vector<TBlob>& inputs,
+                          const std::vector<OpReqType>& req,
+                          const std::vector<TBlob>& outputs) {
+  using namespace mshadow;
+  using namespace mshadow::expr;
+  const NumpyMoveaxisParam& param = nnvm::get<NumpyMoveaxisParam>(attrs.parsed);
+  CHECK_EQ(inputs.size(), 1U);
+  CHECK_EQ(outputs.size(), 1U);
+  CHECK_EQ(req[0], kWriteTo) << "Moveaxis does not support inplace";
+  mxnet::TShape axes(inputs[0].ndim(), -1);
+  mxnet::TShape real_src(param.source.ndim(), -1);
+  mxnet::TShape real_des(param.destination.ndim(), -1);
+  std::vector<bool> state_axes(inputs[0].ndim(), false);
+  CHECK_EQ(param.source.ndim(), param.destination.ndim())
+    << "source and destination not equal.";
+  for (int i = 0; i < param.source.ndim(); ++i) {
+    if (param.source[i] >= 0) {
+      CHECK_LT(static_cast<size_t>(param.source[i]), inputs[0].ndim());
+      real_src[i] = param.source[i];
+    } else {
+      CHECK_LT(param.source[i] + inputs[0].ndim(), inputs[0].ndim());
+      real_src[i] = param.source[i] + inputs[0].ndim();
+    }
+    if (param.destination[i] >= 0) {
+      CHECK_LT(static_cast<size_t>(param.destination[i]), inputs[0].ndim());
+      real_des[i] = param.destination[i];
+    } else {
+      CHECK_LT(param.destination[i] + inputs[0].ndim(), inputs[0].ndim());
+      real_des[i] = param.destination[i] + inputs[0].ndim();
+    }
+  }
+  if (inputs[0].ndim() > 1) {
+    for (int i = 0; i < param.source.ndim() - 1; ++i) {
+      for (int j = i + 1; j < param.source.ndim(); ++j) {
+        CHECK_NE(real_src[i], real_src[j])
+          << "repeated axis in `source` argument";
+        CHECK_NE(real_des[i], real_des[j])
+          << "repeated axis in `destination` argument";
+      }
+    }
+  }
+  for (int i = 0; i < param.source.ndim(); ++i) {
+    axes[real_des[i]] = real_src[i];
+    state_axes[real_src[i]] = true;
+  }
+  for (int i = 0; i < axes.ndim(); ++i) {
+    if (axes[i] < 0) {
+      for (int j = 0; j < axes.ndim(); ++j) {
+        if (state_axes[j] == false) {
+          axes[i] = j;
+          state_axes[j] = true;
+          break;
+        }
+      }
+    }
+  }
+  MSHADOW_TYPE_SWITCH(outputs[0].type_flag_, Dtype, {
+    TransposeImpl<xpu>(ctx.run_ctx, inputs[0], outputs[0], axes);
+  })
+}
+
 }  // namespace op
 }  // namespace mxnet
 

src/operator/numpy/np_matrix_op.cc

@@ -30,6 +30,7 @@ namespace mxnet {
 namespace op {
 
 DMLC_REGISTER_PARAMETER(NumpyTransposeParam);
+DMLC_REGISTER_PARAMETER(NumpyMoveaxisParam);
 
 bool NumpyTransposeShape(const nnvm::NodeAttrs& attrs,
                          mxnet::ShapeVector *in_attrs,
@@ -345,5 +346,95 @@ Examples::
 .add_argument("data", "NDArray-or-Symbol[]", "List of arrays to stack")
 .add_arguments(StackParam::__FIELDS__());
 
+bool NumpyMoveaxisShape(const nnvm::NodeAttrs& attrs,
+                        mxnet::ShapeVector *in_attrs,
+                        mxnet::ShapeVector *out_attrs) {
+  const NumpyMoveaxisParam& param = nnvm::get<NumpyMoveaxisParam>(attrs.parsed);
+  CHECK_EQ(in_attrs->size(), 1U);
+  CHECK_EQ(out_attrs->size(), 1U);
+  mxnet::TShape& shp = (*in_attrs)[0];
+  CHECK_LE(shp.ndim(), 6) << "Transpose support at most 6 dimensions";
+  CHECK_EQ(param.source.ndim(), param.destination.ndim())
+    << "source and destination not equal.";
+  mxnet::TShape ret(shp.ndim(), -1);
+  mxnet::TShape axes(shp.ndim(), -1);
+  std::vector<bool> state_axes(shp.ndim(), false);
+  mxnet::TShape real_src(param.source.ndim(), -1);
+  mxnet::TShape real_des(param.destination.ndim(), -1);
+  for (int i = 0; i < param.source.ndim(); ++i) {
+    if (param.source[i] >= 0) {
+      CHECK_LT(static_cast<size_t>(param.source[i]), shp.ndim());
+      real_src[i] = param.source[i];
+    } else {
+      CHECK_LT(param.source[i] + shp.ndim(), shp.ndim());
+      real_src[i] = param.source[i] + shp.ndim();
+    }
+    if (param.destination[i] >= 0) {
+      CHECK_LT(static_cast<size_t>(param.destination[i]), shp.ndim());
+      real_des[i] = param.destination[i];
+    } else {
+      CHECK_LT(param.destination[i] + shp.ndim(), shp.ndim());
+      real_des[i] = param.destination[i] + shp.ndim();
+    }
+  }
+  if (shp.ndim() > 1) {
+    for (int i = 0; i < param.source.ndim() - 1; ++i) {
+      for (int j = i + 1; j < param.source.ndim(); ++j) {
+        CHECK_NE(real_src[i], real_src[j])
+          << "repeated axis in `source` argument";
+        CHECK_NE(real_des[i], real_des[j])
+          << "repeated axis in `destination` argument";
+      }
+    }
+  }
+  for (int i = 0; i < param.source.ndim(); ++i) {
+    axes[real_des[i]] = real_src[i];
+    state_axes[real_src[i]] = true;
+  }
+  for (int i = 0; i < axes.ndim(); ++i) {
+    if (axes[i] < 0) {
+      for (int j = 0; j < axes.ndim(); ++j) {
+        if (state_axes[j] == false) {
+          axes[i] = j;
+          state_axes[j] = true;
+          break;
+        }
+      }
+    }
+  }
+  for (int i = 0; i < shp.ndim(); ++i) {
+    CHECK(axes[i] < static_cast<int64_t>(shp.ndim()));
+    ret[i] = shp[axes[i]];
+  }
+  SHAPE_ASSIGN_CHECK(*out_attrs, 0, ret);
+  return shape_is_known(ret);
+}
+
+NNVM_REGISTER_OP(_np_moveaxis)
+.describe(R"code(Move axes of an array to new positions.
+Other axes remain in their original order.
+)code" ADD_FILELINE)
+.set_num_inputs(1)
+.set_num_outputs(1)
+.set_attr_parser(ParamParser<NumpyMoveaxisParam>)
+.set_attr<mxnet::FInferShape>("FInferShape", NumpyMoveaxisShape)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<1, 1>)
+.set_attr<nnvm::FGradient>("FGradient",
+  [](const nnvm::NodePtr& n, const std::vector<nnvm::NodeEntry>& ograds) {
+     const NumpyMoveaxisParam& param = nnvm::get<NumpyMoveaxisParam>(n->attrs.parsed);
+     std::ostringstream os1;
+     os1 << param.source;
+     std::ostringstream os2;
+     os2 << param.destination;
+     return MakeNonlossGradNode("_np_moveaxis", n, ograds, {},
+                                {{"source", os2.str()}, {"destination", os1.str()}});
+})
+.set_attr<FCompute>("FCompute<cpu>", NumpyMoveaxisCompute<cpu>)
+.set_attr<nnvm::FListInputNames>("FListInputNames",
+  [](const NodeAttrs& attrs) { return std::vector<std::string>{"a"};
+})
+.add_argument("a", "NDArray-or-Symbol", "Source input")
+.add_arguments(NumpyMoveaxisParam::__FIELDS__());
+
 }  // namespace op
 }  // namespace mxnet

src/operator/numpy/np_matrix_op.cu

@@ -46,5 +46,8 @@ NNVM_REGISTER_OP(_backward_np_concat)
 NNVM_REGISTER_OP(_npi_stack)
 .set_attr<FCompute>("FCompute<gpu>", StackOpForward<gpu>);
 
+NNVM_REGISTER_OP(_np_moveaxis)
+.set_attr<FCompute>("FCompute<gpu>", NumpyMoveaxisCompute<gpu>);
+
 }  // namespace op
 }  // namespace mxnet

tests/python/unittest/test_numpy_op.py

@@ -1647,6 +1647,50 @@ def hybrid_forward(self, F, a):
                         assert_almost_equal(mx_out.asnumpy(), np_out, rtol=1e-3, atol=1e-5)
 
 
+@with_seed()
+@use_np
+def test_np_moveaxis():
+    class TestMoveaxis(HybridBlock):
+        def __init__(self, source=None, destination=None):
+            super(TestMoveaxis, self).__init__()
+            self._source = source
+            self._destination= destination
+
+        def hybrid_forward(self, F, x):
+            return F.np.moveaxis(x, source=self._source, destination=self._destination)
+
+    dtypes = ['int32', 'int64', 'float16', 'float32', 'float64']
+    for hybridize in [False, True]:
+        for dtype in dtypes:
+            for ndim in [0, 1, 2, 3, 4, 5, 6]:
+                shape = rand_shape_nd(ndim, dim=5, allow_zero_size=True)
+                np_data = _np.random.uniform(low=-100, high=100, size=shape).astype(dtype)
+                mx_data = np.array(np_data, dtype=dtype)
+                axis = [i for i in range(ndim)]
+                random.shuffle(axis)
+                for i in range(ndim):
+                    source = random.sample(axis, i)
+                    destination = random.sample(axis, i)
+
+                    # test gluon
+                    test_moveaxis = TestMoveaxis(source,destination)
+                    if hybridize:
+                        test_moveaxis.hybridize()
+                    np_out = _np.moveaxis(np_data, source=source, destination=destination)
+                    mx_data.attach_grad()
+                    with mx.autograd.record():
+                        mx_out = test_moveaxis(mx_data)
+                    assert mx_out.shape == np_out.shape
+                    mx_out.backward()
+                    assert same(mx_data.grad.shape, mx_data.shape)
+                    assert same(mx_data.grad.asnumpy(), _np.ones(shape))
+                    # test imperative
+                    np_out = _np.moveaxis(np_data, source=source, destination=destination)
+                    mx_out = np.moveaxis(mx_data, source=source, destination= destination)
+                    assert np_out.dtype == mx_out.dtype
+                    assert same(mx_out.asnumpy(), np_out)
+
+
 if __name__ == '__main__':
     import nose
     nose.runmodule()