text: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *   http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   * KIND, either express or implied.  See the License for the
   * specific language governing permissions and limitations
   * under the License.
   */
  
  /*!
   *  Copyright (c) 2020 by Contributors
   * \file np_cross-inl.h
   * \brief Function definition of cross product of two (arrays of) vectors
   */
  
  #ifndef MXNET_OPERATOR_NUMPY_NP_CROSS_INL_H_
  #define MXNET_OPERATOR_NUMPY_NP_CROSS_INL_H_
  
  #include <mxnet/operator_util.h>
  #include <vector>
  #include <algorithm>
  #include <string>
  #include "../mshadow_op.h"
  #include "../operator_common.h"
  #include "../elemwise_op_common.h"
  #include "../tensor/broadcast_reduce_op.h"
  #include "../tensor/elemwise_binary_broadcast_op.h"
  #include "../tensor/matrix_op-inl.h"
  
  namespace mxnet {
  namespace op {
  
  using namespace mshadow;
  
  struct NumpyCrossParam : public dmlc::Parameter<NumpyCrossParam> {
    int axisa, axisb, axisc;
    DMLC_DECLARE_PARAMETER(NumpyCrossParam) {
      DMLC_DECLARE_FIELD(axisa)
      .set_default(-1)
      .describe("Axis of `a` that defines the vector(s).  By default, the last axis.");
      DMLC_DECLARE_FIELD(axisb)
      .set_default(-1)
      .describe("Axis of `b` that defines the vector(s).  By default, the last axis.");
      DMLC_DECLARE_FIELD(axisc)
      .set_default(-1)
      .describe("Axis of `c` containing the cross product vector(s)."
                "Ignored if both input vectors have dimension 2, as the return is scalar."
                "By default, the last axis.");
    }
    void SetAttrDict(std::unordered_map<std::string, std::string>* dict) {
      std::ostringstream axisa_s, axisb_s, axisc_s;
      axisa_s << axisa;
      axisb_s << axisb;
      axisc_s << axisc;
      (*dict)["axisa"] = axisa_s.str();
      (*dict)["axisb"] = axisb_s.str();
      (*dict)["axisc"] = axisc_s.str();
    }
  };
  
  #define SUM_NDIM_SWITCH(ndim, NDim, ...)  \
    if (ndim == 1) {                    \
      const int NDim = 1;               \
      {__VA_ARGS__}                     \
    } else if (ndim == 2) {             \
      const int NDim = 2;               \
      {__VA_ARGS__}                     \
    } else if (ndim == 3) {             \
      const int NDim = 3;               \
      {__VA_ARGS__}                     \
    } else if (ndim == 4) {             \
      const int NDim = 4;               \
      {__VA_ARGS__}                     \
    } else if (ndim <= broadcast::MAX_DIM) {  \
      const int NDim = broadcast::MAX_DIM;    \
      {__VA_ARGS__}                     \
    } else {                            \
      LOG(FATAL) << "NDim too large ";  \
    }
  
  struct CrossInAssign {
    template<typename DType>
    MSHADOW_XINLINE static void Map(int i, const DType *in_ptr, DType *out_ptr,
                                    const int stride, const int index, const int msize) {
      if (index < stride && i * stride + index < msize) {
        out_ptr[i] = in_ptr[i * stride + index];
      }
    }
  };
  
  template<int req>
  struct CrossOutAssign {
    template<typename DType>
    MSHADOW_XINLINE static void Map(int i, const DType *in_ptr, DType *out_ptr,
                                    const int positive, const int stride,
                                    const int index, const int msize) {
      if (index < stride && i * stride + index < msize) {
        KERNEL_ASSIGN(out_ptr[i * stride + index], req, positive == 1 ? in_ptr[i] : -in_ptr[i]);
      }
    }
  };
  
  template<int req>
  struct ResAssign {
    template<typename DType>
    MSHADOW_XINLINE static void Map(int i, const DType *in_data, DType *out_data) {
      KERNEL_ASSIGN(out_data[i], req, in_data[i]);
    }
  };
  
  struct DeleteAssign {
    template<typename DType>
    MSHADOW_XINLINE static void Map(int i, const DType *in_data, DType *out_data,
                                    const int in_stride, const int out_stride) {
      const DType *in_ptr = in_data + i * in_stride;
      DType *out_ptr = out_data + i * out_stride;
      if (in_stride == out_stride + 1) {
        for (int idx = 0; idx < out_stride; ++idx) {
          out_ptr[idx] = in_ptr[idx];
        }
      }
    }
  };
  
  // Get moveaxis index.
  inline mxnet::Tuple<int> GetMoveaxisIndex(const int& source,
                                            const int& destination,
                                            const mxnet::TShape& shape) {
    const int ndim = shape.ndim();
    const int src_axis = CheckAxis(source, ndim);
    const int dest_axis = CheckAxis(destination, ndim);
    std::vector<int> moveaxis_index_vec;
    for (int i = 0; i < ndim; ++i) {
      if (i != src_axis) { moveaxis_index_vec.push_back(i); }
    }
    moveaxis_index_vec.insert(moveaxis_index_vec.begin() + dest_axis, src_axis);
    return mxnet::Tuple<int>(moveaxis_index_vec);
  }
  
  // Get moveaxis shape.
  inline mxnet::TShape GetMoveaxisShape(const Tuple<int>& moveaxis_index,
                                        const mxnet::TShape& org_shape) {
    const int ndim = org_shape.ndim();
    if (ndim == 0) { return mxnet::TShape(0, 0); }
    CHECK_EQ(moveaxis_index.ndim(), org_shape.ndim()) << "moveaxis index dismatch original shape.";
    std::vector<int> moveaxis_shape_vec(ndim, -1);
    for (int i = 0; i < ndim; ++i) {
      moveaxis_shape_vec[i] = org_shape[moveaxis_index[i]];
    }
    return mxnet::TShape(moveaxis_shape_vec.begin(), moveaxis_shape_vec.end());
  }
  
  // Get or check broadcast shape for cross product.
  inline void GetOrCheckLRShape(const nnvm::NodeAttrs& attrs,
                                const mxnet::TShape& a_moveaxis_shape,
                                const mxnet::TShape& b_moveaxis_shape,
                                mxnet::TShape *c_shape_ptr = nullptr) {
    const int a_ndim = a_moveaxis_shape.ndim();
    const int b_ndim = b_moveaxis_shape.ndim();
    mxnet::TShape a_cutoff_shape(a_ndim - 1, -1);
    mxnet::TShape b_cutoff_shape(b_ndim - 1, -1);
    for (int i = 0; i < a_ndim - 1; ++i) {
      a_cutoff_shape[i] = a_moveaxis_shape[i];
    }
    for (int i = 0; i < b_ndim - 1; ++i) {
      b_cutoff_shape[i] = b_moveaxis_shape[i];
    }
    mxnet::ShapeVector in_shape_vec({ a_cutoff_shape, b_cutoff_shape});
    mxnet::ShapeVector out_shape_vec({ mxnet::TShape() });
    mxnet::op::BinaryBroadcastShape(attrs, &in_shape_vec, &out_shape_vec);
    if (c_shape_ptr && (a_moveaxis_shape[a_ndim - 1] == 3 || b_moveaxis_shape[b_ndim - 1] == 3)) {
      mxnet::TShape c_shape(out_shape_vec[0].ndim() + 1, -1);
      for (int i = 0; i < c_shape.ndim() - 1; ++i) {
        c_shape[i] = out_shape_vec[0][i];
      }
      c_shape[c_shape.ndim() - 1] = 3;
      *c_shape_ptr = c_shape;
    } else {
      *c_shape_ptr = out_shape_vec[0];
    }
  }
  
  // Get data[..., 0] shape.
  inline mxnet::TShape GetCutoffShape(const mxnet::TShape& shape) {
    if (shape.ndim() == 0 || !ndim_is_known(shape)) { return mxnet::TShape(0, 0); }
    mxnet::TShape cutoff_shape(shape.ndim() - 1, -1);
    for (int i = 0; i < shape.ndim() - 1; ++i) { cutoff_shape[i] = shape[i]; }
    return cutoff_shape;
  }
  
  template<typename DType>
  inline size_t AligndWorkspaceSize(const size_t& offset,
                                    const size_t& add_size) {
    size_t m32 = 32;
    size_t wks = offset + add_size;
    return ((wks * sizeof(DType) + m32 - 1) / m32 * m32 + sizeof(DType) - 1) / sizeof(DType);
  }
  
  // Calculate workspace size for numpy cross forward op.
  template<typename xpu, typename DType>
  inline size_t NumpyCrossWorkspaceSize(const mxnet::TShape& a_moveaxis_shape,
                                        const mxnet::TShape& b_moveaxis_shape,
                                        const mxnet::TShape& c_moveaxis_shape,
                                        const int& a_axis,
                                        const int& b_axis,
                                        const OpContext& ctx,
                                        const std::vector<OpReqType>& req) {
    if (kNullOp == req[0]) { return 0U; }
    // Zero-size input, no need to launch kernel
    if (0U == a_moveaxis_shape.Size() || 0U == b_moveaxis_shape.Size()) { return 0U; }
  
    size_t workspace_size = 0;
    const int a_ndim = a_moveaxis_shape.ndim();
    const int b_ndim = b_moveaxis_shape.ndim();
    const int c_ndim = c_moveaxis_shape.ndim();
  
    if (ctx.run_ctx.get_ctx().dev_mask() == cpu::kDevMask) {
      if (a_moveaxis_shape[a_ndim - 1] == 2 && b_moveaxis_shape[b_ndim - 1] == 2) {
        // Case 1: a.shape[-1] == 2 and b.shape[-1] == 2, param.axisc is ignored.
        workspace_size += a_moveaxis_shape.ProdShape(0, a_ndim - 1);
        workspace_size += b_moveaxis_shape.ProdShape(0, b_ndim - 1);
        workspace_size += c_moveaxis_shape.Size();
      } else {
        // Case 2, 3, 4: a.shape[-1] == 3 or b.shape[-1] == 3, param.axisc is not ignored.
        workspace_size += a_moveaxis_shape.ProdShape(0, a_ndim - 1);
        workspace_size += b_moveaxis_shape.ProdShape(0, b_ndim - 1);
        workspace_size += c_moveaxis_shape.Size();
        workspace_size += 3 * c_moveaxis_shape.ProdShape(0, c_ndim - 1);
      }
      if (a_axis != a_ndim -1 || b_axis != b_ndim - 1) {
        workspace_size += a_moveaxis_shape.Size();
        workspace_size += b_moveaxis_shape.Size();
      }
    } else {
      if (a_moveaxis_shape[a_ndim - 1] == 2 && b_moveaxis_shape[b_ndim - 1] == 2) {
        // Case 1: a.shape[-1] == 2 and b.shape[-1] == 2, param.axisc is ignored.
        workspace_size = AligndWorkspaceSize<DType>(workspace_size,
                                                    a_moveaxis_shape.ProdShape(0, a_ndim - 1));
        workspace_size = AligndWorkspaceSize<DType>(workspace_size,
                                                    b_moveaxis_shape.ProdShape(0, b_ndim - 1));
        workspace_size = AligndWorkspaceSize<DType>(workspace_size,
                                                    c_moveaxis_shape.Size());
      } else {
        // Case 2, 3, 4: a.shape[-1] == 3 or b.shape[-1] == 3, param.axisc is not ignored.
        workspace_size = AligndWorkspaceSize<DType>(workspace_size,
                                                    a_moveaxis_shape.ProdShape(0, a_ndim - 1));
        workspace_size = AligndWorkspaceSize<DType>(workspace_size,
                                                    b_moveaxis_shape.ProdShape(0, b_ndim - 1));
        for (int i = 0; i < 3; ++i) {
          workspace_size = AligndWorkspaceSize<DType>(workspace_size,
                                                      c_moveaxis_shape.ProdShape(0, c_ndim - 1));
        }
        workspace_size = AligndWorkspaceSize<DType>(workspace_size,
                                                    c_moveaxis_shape.Size());
      }
      if (a_axis != a_ndim -1 || b_axis != b_ndim - 1) {
        workspace_size = AligndWorkspaceSize<DType>(workspace_size, a_moveaxis_shape.Size());
        workspace_size = AligndWorkspaceSize<DType>(workspace_size, b_moveaxis_shape.Size());
      }
    }
    return workspace_size;
  }
  
  template<typename xpu, typename DType, int a_dim, int b_dim>
  struct NumpyCrossForwardImpl {
    static void op(const TBlob& a, const TBlob& b, const TBlob& c,
                   const std::vector<Tuple<int> >& moveaxis_index_vec,
                   const std::vector<mxnet::TShape>& moveaxis_shape_vec,
                   const int a_axis, const int b_axis, const int c_axis,
                   const nnvm::NodeAttrs& attrs,
                   const OpContext& ctx,
                   const std::vector<OpReqType>& req,
                   const Tensor<xpu, 1, DType>& workspace) {
      CHECK(a_dim == 3 || b_dim == 3)
        << "no specialized NumpyCrossOp defined for template parameters.";
      Stream<xpu> *s = ctx.get_stream<xpu>();
      const Tuple<int>& a_moveaxis_index = moveaxis_index_vec[0];
      const Tuple<int>& b_moveaxis_index = moveaxis_index_vec[1];
      const mxnet::TShape& a_moveaxis_shape = moveaxis_shape_vec[0];
      const mxnet::TShape& b_moveaxis_shape = moveaxis_shape_vec[1];
      const mxnet::TShape& c_moveaxis_shape = moveaxis_shape_vec[2];
      const int a_ndim = a_moveaxis_shape.ndim();
      const int b_ndim = b_moveaxis_shape.ndim();
      const int c_ndim = b_moveaxis_shape.ndim();
      CHECK_EQ(c_moveaxis_shape[c_ndim - 1], 3)
        << "no specialized NumpyCrossOp defined for template parameters.";
  
      TBlob aw_data, bw_data, c_data, cw_data, a_data, b_data;
      std::vector<TBlob> cw_data_vec;
      if (ctx.run_ctx.get_ctx().dev_mask() == cpu::kDevMask) {
        // Allocate workspace in cpu, no need to align address.
        DType *aw_ptr = workspace.dptr_;
        DType *bw_ptr = aw_ptr + a_moveaxis_shape.ProdShape(0, a_ndim - 1);
        DType *cw_ptr = bw_ptr + b_moveaxis_shape.ProdShape(0, b_ndim - 1);
        DType *c_ptr = cw_ptr + 3 * c_moveaxis_shape.ProdShape(0, c_ndim - 1);
        a_data = a;
        b_data = b;
        if (a_axis != a_ndim -1 || b_axis != b_ndim - 1) {
          DType *a_ptr = c_ptr + c_moveaxis_shape.Size();
          DType *b_ptr = a_ptr + a_moveaxis_shape.Size();
          a_data = TBlob(a_ptr, a_moveaxis_shape, a.dev_mask(), a.dev_id());
          b_data = TBlob(b_ptr, b_moveaxis_shape, b.dev_mask(), b.dev_id());
          TransposeImpl<xpu>(ctx.run_ctx, a, a_data,
                             mxnet::TShape(a_moveaxis_index.begin(), a_moveaxis_index.end()));
          TransposeImpl<xpu>(ctx.run_ctx, b, b_data,
                             mxnet::TShape(b_moveaxis_index.begin(), b_moveaxis_index.end()));
        }
        aw_data = TBlob(aw_ptr, GetCutoffShape(a_moveaxis_shape), a.dev_mask(), a.dev_id());
        bw_data = TBlob(bw_ptr, GetCutoffShape(b_moveaxis_shape), b.dev_mask(), b.dev_id());
        cw_data = TBlob(cw_ptr, c_moveaxis_shape, c.dev_mask(), c.dev_id());
        c_data = TBlob(c_ptr, c_moveaxis_shape, c.dev_mask(), c.dev_id());
        for (int i = 0; i < 3; ++i) {
          cw_data_vec.push_back(TBlob(cw_ptr + i * c_moveaxis_shape.ProdShape(0, c_ndim - 1),
                                      GetCutoffShape(c_moveaxis_shape), c.dev_mask(), c.dev_id()));
        }
      } else {
        // Allocate workspace in gpu, need to align address.
        size_t offset = 0;
        aw_data = TBlob(workspace.dptr_ + offset, GetCutoffShape(a_moveaxis_shape),
                        a.dev_mask(), a.dev_id());
        offset = AligndWorkspaceSize<DType>(offset, aw_data.shape_.Size());
  
        bw_data = TBlob(workspace.dptr_ + offset, GetCutoffShape(b_moveaxis_shape),
                        b.dev_mask(), b.dev_id());
        offset = AligndWorkspaceSize<DType>(offset, bw_data.shape_.Size());
  
        cw_data = TBlob(workspace.dptr_ + offset, c_moveaxis_shape, c.dev_mask(), c.dev_id());
        for (int i = 0; i < 3; ++i) {
          cw_data_vec.push_back(TBlob(workspace.dptr_ + offset,
                                      GetCutoffShape(c_moveaxis_shape), c.dev_mask(), c.dev_id()));
          offset = AligndWorkspaceSize<DType>(offset, cw_data_vec[i].shape_.Size());
        }
        c_data = TBlob(workspace.dptr_ + offset, c_moveaxis_shape, c.dev_mask(), c.dev_id());
        offset = AligndWorkspaceSize<DType>(offset, c_data.shape_.Size());
  
        a_data = a;
        b_data = b;
        if (a_axis != a_ndim -1 || b_axis != b_ndim - 1) {
          a_data = TBlob(workspace.dptr_ + offset, a_moveaxis_shape,
                         a.dev_mask(), a.dev_id());
          offset = AligndWorkspaceSize<DType>(offset, a_data.shape_.Size());
  
          b_data = TBlob(workspace.dptr_ + offset, b_moveaxis_shape,
                         b.dev_mask(), b.dev_id());
          offset = AligndWorkspaceSize<DType>(offset, b_data.shape_.Size());
  
          TransposeImpl<xpu>(ctx.run_ctx, a, a_data,
                            mxnet::TShape(a_moveaxis_index.begin(), a_moveaxis_index.end()));
          TransposeImpl<xpu>(ctx.run_ctx, b, b_data,
                            mxnet::TShape(b_moveaxis_index.begin(), b_moveaxis_index.end()));
        }
      }
      std::vector<int> positive_vec;
      std::vector<int> a_index_vec, b_index_vec, c_index_vec;
      std::vector<OpReqType> req_vec;
      if (a_dim == 2 && b_dim == 3) {
        a_index_vec = {1, 0, 0, 1};
        b_index_vec = {2, 2, 1, 0};
        c_index_vec = {0, 1, 2, 2};
        positive_vec = {1, 0, 1, 0};
        req_vec = { kWriteTo, kWriteTo, kWriteTo, kAddTo };
      } else if (a_dim == 3 && b_dim == 2) {
        a_index_vec = {2, 2, 0, 1};
        b_index_vec = {1, 0, 1, 0};
        c_index_vec = {0, 1, 2, 2};
        positive_vec = {0, 1, 1, 0};
        req_vec = { kWriteTo, kWriteTo, kWriteTo, kAddTo };
      } else {
        a_index_vec = {1, 2, 2, 0, 0, 1};
        b_index_vec = {2, 1, 0, 2, 1, 0};
        c_index_vec = {0, 0, 1, 1, 2, 2};
        positive_vec = {1, 0, 1, 0, 1, 0};
        req_vec = { kWriteTo, kAddTo, kWriteTo, kAddTo, kWriteTo, kAddTo};
      }
      for (size_t i = 0; i < a_index_vec.size(); ++i) {
        int idx = c_index_vec[i];
        mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, aw_data.Size(), a_data.dptr<DType>(),
                                                     aw_data.dptr<DType>(), a_data.size(a_ndim - 1),
                                                     a_index_vec[i], a_data.Size());
        mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, bw_data.Size(), b_data.dptr<DType>(),
                                                     bw_data.dptr<DType>(), b_data.size(b_ndim - 1),
                                                     b_index_vec[i], b_data.Size());
  #if !defined(__CUDACC__)
        BinaryBroadcastCompute<xpu, op::mshadow_op::mul>(attrs, ctx, { aw_data, bw_data },
                                                         { kWriteTo }, { cw_data_vec[idx] });
  #else
        BinaryBroadcastRTCCompute {"mul"}(attrs, ctx, { aw_data, bw_data },
                                          { kWriteTo }, { cw_data_vec[idx] });
  #endif  // !defined(__CUDACC__)
        MXNET_ASSIGN_REQ_SWITCH(req_vec[i], req_type, {
          mxnet_op::Kernel<CrossOutAssign<req_type>, xpu>::Launch(s, cw_data_vec[idx].Size(),
                                                                  cw_data_vec[idx].dptr<DType>(),
                                                                  c_data.dptr<DType>(),
                                                                  positive_vec[i],
                                                                  c_data.size(c_ndim - 1),
                                                                  idx, c_data.Size());
        });
      }
      cw_data = cw_data.reshape(c.shape_);
      const DType *res_ptr = c_data.dptr<DType>();
      if (c_axis != c_ndim -1) {
        const Tuple<int> c_axis_index = GetMoveaxisIndex(-1, c_axis, c_moveaxis_shape);
        TransposeImpl<xpu>(ctx.run_ctx, c_data, cw_data,
                           mxnet::TShape(c_axis_index.begin(), c_axis_index.end()));
        res_ptr = cw_data.dptr<DType>();
      }
      MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
        mxnet_op::Kernel<ResAssign<req_type>, xpu>::Launch(
          s, c.Size(), res_ptr, c.dptr<DType>());
      });
    }
  };
  
  template<typename xpu, typename DType>
  struct NumpyCrossForwardImpl<xpu, DType, 2, 2> {
    static void op(const TBlob& a, const TBlob& b, const TBlob& c,
                   const std::vector<Tuple<int> >& moveaxis_index_vec,
                   const std::vector<mxnet::TShape>& moveaxis_shape_vec,
                   const int a_axis, const int b_axis, const int c_axis,
                   const nnvm::NodeAttrs& attrs,
                   const OpContext& ctx,
                   const std::vector<OpReqType>& req,
                   const Tensor<xpu, 1, DType>& workspace) {
      Stream<xpu> *s = ctx.get_stream<xpu>();
      const Tuple<int>& a_moveaxis_index = moveaxis_index_vec[0];
      const Tuple<int>& b_moveaxis_index = moveaxis_index_vec[1];
      const mxnet::TShape& a_moveaxis_shape = moveaxis_shape_vec[0];
      const mxnet::TShape& b_moveaxis_shape = moveaxis_shape_vec[1];
      const mxnet::TShape& c_shape = c.shape_;
      const int a_ndim = a_moveaxis_shape.ndim();
      const int b_ndim = b_moveaxis_shape.ndim();
  
      TBlob aw_data, bw_data, cw_data, a_data, b_data;
      if (ctx.run_ctx.get_ctx().dev_mask() == cpu::kDevMask) {
        // Allocate workspace in cpu, no need to align address.
        DType *aw_ptr = workspace.dptr_;
        DType *bw_ptr = aw_ptr + a_moveaxis_shape.ProdShape(0, a_ndim - 1);
        DType *cw_ptr = bw_ptr + b_moveaxis_shape.ProdShape(0, b_ndim - 1);
        aw_data = TBlob(aw_ptr, GetCutoffShape(a_moveaxis_shape), a.dev_mask(), a.dev_id());
        bw_data = TBlob(bw_ptr, GetCutoffShape(b_moveaxis_shape), b.dev_mask(), b.dev_id());
        cw_data = TBlob(cw_ptr, c_shape, c.dev_mask(), c.dev_id());
        a_data = a;
        b_data = b;
        if (a_axis != a_ndim -1 || b_axis != b_ndim - 1) {
          DType *a_ptr = cw_ptr + c_shape.Size();
          DType *b_ptr = a_ptr + a_moveaxis_shape.Size();
          a_data = TBlob(a_ptr, a_moveaxis_shape, a.dev_mask(), a.dev_id());
          b_data = TBlob(b_ptr, b_moveaxis_shape, b.dev_mask(), b.dev_id());
          TransposeImpl<xpu>(ctx.run_ctx, a, a_data,
                             mxnet::TShape(a_moveaxis_index.begin(), a_moveaxis_index.end()));
          TransposeImpl<xpu>(ctx.run_ctx, b, b_data,
                             mxnet::TShape(b_moveaxis_index.begin(), b_moveaxis_index.end()));
        }
      } else {
        // Allocate workspace in cpu, need to align address.
        size_t offset = 0;
        aw_data = TBlob(workspace.dptr_ + offset, GetCutoffShape(a_moveaxis_shape),
                        a.dev_mask(), a.dev_id());
        offset = AligndWorkspaceSize<DType>(offset, aw_data.shape_.Size());
  
        bw_data = TBlob(workspace.dptr_ + offset, GetCutoffShape(b_moveaxis_shape),
                        b.dev_mask(), b.dev_id());
        offset = AligndWorkspaceSize<DType>(offset, bw_data.shape_.Size());
  
        cw_data = TBlob(workspace.dptr_ + offset, c_shape,
                        c.dev_mask(), c.dev_id());
        offset = AligndWorkspaceSize<DType>(offset, cw_data.shape_.Size());
        a_data = a;
        b_data = b;
        if (a_axis != a_ndim -1 || b_axis != b_ndim - 1) {
          a_data = TBlob(workspace.dptr_ + offset, a_moveaxis_shape,
                         a.dev_mask(), a.dev_id());
          offset = AligndWorkspaceSize<DType>(offset, a_data.shape_.Size());
  
          b_data = TBlob(workspace.dptr_ + offset, b_moveaxis_shape,
                         b.dev_mask(), b.dev_id());
          offset = AligndWorkspaceSize<DType>(offset, b_data.shape_.Size());
  
          TransposeImpl<xpu>(ctx.run_ctx, a, a_data,
                             mxnet::TShape(a_moveaxis_index.begin(), a_moveaxis_index.end()));
          TransposeImpl<xpu>(ctx.run_ctx, b, b_data,
                             mxnet::TShape(b_moveaxis_index.begin(), b_moveaxis_index.end()));
        }
      }
      mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, aw_data.Size(), a_data.dptr<DType>(),
                                                   aw_data.dptr<DType>(), a_data.size(a_ndim - 1),
                                                   0, a_data.Size());
      mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, bw_data.Size(), b_data.dptr<DType>(),
                                                   bw_data.dptr<DType>(), b_data.size(b_ndim - 1),
                                                   1, b_data.Size());
  #if !defined(__CUDACC__)
      BinaryBroadcastCompute<xpu, op::mshadow_op::mul>(attrs, ctx, { aw_data, bw_data },
                                                       { req[0] }, { c });
  #else
      BinaryBroadcastRTCCompute {"mul"}(attrs, ctx, { aw_data, bw_data },
                                        { req[0] }, { c });
  #endif  // !defined(__CUDACC__)
      mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, aw_data.Size(), a_data.dptr<DType>(),
                                                   aw_data.dptr<DType>(), a_data.size(a_ndim - 1),
                                                   1, a_data.Size());
      mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, bw_data.Size(), b_data.dptr<DType>(),
                                                   bw_data.dptr<DType>(), b_data.size(b_ndim - 1),
                                                   0, b_data.Size());
  #if !defined(__CUDACC__)
      BinaryBroadcastCompute<xpu, op::mshadow_op::mul>(attrs, ctx, { aw_data, bw_data },
                                                       { kWriteTo }, { cw_data });
      BinaryBroadcastCompute<xpu, op::mshadow_op::minus>(attrs, ctx, { c, cw_data },
                                                         { kWriteTo }, { c });
  #else
      BinaryBroadcastRTCCompute {"mul"}(attrs, ctx, { aw_data, bw_data },
                                        { kWriteTo }, { cw_data });
      BinaryBroadcastRTCCompute {"sub"}(attrs, ctx, { c, cw_data },
                                        { kWriteTo }, { c });
  #endif  // !defined(__CUDACC__)
    }
  };
  
  template<typename xpu>
  void NumpyCrossForward(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;
    CHECK_EQ(inputs.size(), 2U);
    CHECK_EQ(outputs.size(), 1U);
    CHECK_EQ(req.size(), 1U);
  
    Stream<xpu> *s = ctx.get_stream<xpu>();
    const TBlob& a = inputs[0];
    const TBlob& b = inputs[1];
    const TBlob& c = outputs[0];
  
    if (kNullOp == req[0]) { return; }
    // Zero-size output, no need to launch kernel
    if (0U == a.Size() || 0U == b.Size()) { return; }
  
    const NumpyCrossParam& param = nnvm::get<NumpyCrossParam>(attrs.parsed);
    const mxnet::TShape& a_shape = a.shape_;
    const mxnet::TShape& b_shape = b.shape_;
    const mxnet::TShape& c_shape = c.shape_;
    const int a_ndim = a_shape.ndim();
    const int b_ndim = b_shape.ndim();
    const int c_ndim = c_shape.ndim();
    Tuple<int> a_moveaxis_index = GetMoveaxisIndex(param.axisa, -1, a_shape);
    Tuple<int> b_moveaxis_index = GetMoveaxisIndex(param.axisb, -1, b_shape);
    Tuple<int> c_moveaxis_index = GetMoveaxisIndex(param.axisc, -1, c_shape);
    mxnet::TShape a_moveaxis_shape = GetMoveaxisShape(a_moveaxis_index, a_shape);
    mxnet::TShape b_moveaxis_shape = GetMoveaxisShape(b_moveaxis_index, b_shape);
    mxnet::TShape c_moveaxis_shape = GetMoveaxisShape(c_moveaxis_index, c_shape);
    const int a_axis = CheckAxis(param.axisa, a_ndim);
    const int b_axis = CheckAxis(param.axisb, b_ndim);
    const int c_axis = CheckAxis(param.axisc, c_ndim);
    const std::vector<mxnet::TShape > shape_vec({ a_moveaxis_shape, b_moveaxis_shape,
                                                  c_moveaxis_shape });
    const std::vector<Tuple<int> > index_vec({ a_moveaxis_index, b_moveaxis_index,
                                               c_moveaxis_index });
  
    MSHADOW_SGL_DBL_TYPE_SWITCH(c.type_flag_, DType, {
      // Calculate workspace.
      size_t workspace_size = NumpyCrossWorkspaceSize<xpu, DType>(a_moveaxis_shape,
                                                                  b_moveaxis_shape,
                                                                  c_moveaxis_shape,
                                                                  a_axis, b_axis, ctx, req);
      Tensor<xpu, 1, DType> workspace = ctx.requested[0].get_space_typed<xpu, 1, DType>(
        Shape1(workspace_size), s);
  
      if (a_moveaxis_shape[a_ndim - 1] == 2) {
        if (b_moveaxis_shape[b_ndim - 1] == 2) {
          // Case 1: a.shape[-1] == 2 and b.shape[-1] == 2, param.axisc is ignored.
          NumpyCrossForwardImpl<xpu, DType, 2, 2>::op(a, b, c, index_vec, shape_vec,
                                                      a_axis, b_axis, c_axis,
                                                      attrs, ctx, req, workspace);
        } else {
          // Case 2: a.shape[-1] == 2 and b.shape[-1] == 3, param.axisc is not ignored.
          NumpyCrossForwardImpl<xpu, DType, 2, 3>::op(a, b, c, index_vec, shape_vec,
                                                      a_axis, b_axis, c_axis,
                                                      attrs, ctx, req, workspace);
        }
      } else {
        if (b_moveaxis_shape[b_ndim - 1] == 2) {
          // Case 3: a.shape[-1] == 3 and b.shape[-1] == 2, param.axisc is not ignored.
          NumpyCrossForwardImpl<xpu, DType, 3, 2>::op(a, b, c, index_vec, shape_vec,
                                                      a_axis, b_axis, c_axis,
                                                      attrs, ctx, req, workspace);
        } else {
          // Case 4: a.shape[-1] == 3 and b.shape[-1] == 3, param.axisc is not ignored.
          NumpyCrossForwardImpl<xpu, DType, 3, 3>::op(a, b, c, index_vec, shape_vec,
                                                      a_axis, b_axis, c_axis,
                                                      attrs, ctx, req, workspace);
        }
      }
    });
  }
  
  inline bool CheckUseBroadcast(const mxnet::TShape& a_move_shape,
                                const mxnet::TShape& b_move_shape) {
    return !(GetCutoffShape(a_move_shape) == GetCutoffShape(b_move_shape));
  }
  
  inline mxnet::TShape GetOriShape(const mxnet::TShape& move_shape,
                                   const int axis) {
    Tuple<int> origin_index = GetMoveaxisIndex(-1, axis, move_shape);
    return GetMoveaxisShape(origin_index, move_shape);
  }
  
  inline std::vector<int> GetReduceAxis(const mxnet::TShape& move_shape,
                                        const mxnet::TShape& broad_move_shape) {
    std::vector<int> axis_idx;
    if (move_shape.ndim() == broad_move_shape.ndim() ||
        move_shape.ndim() == broad_move_shape.ndim() + 1) {
      for (int i = 0; i < move_shape.ndim() - 1; ++i) {
        if (move_shape[i] != broad_move_shape[i]) { axis_idx.push_back(i); }
      }
    }
    return axis_idx;
  }
  
  template<typename xpu, typename DType, int a_dim, int b_dim>
  inline void CrossImplWrap(const std::vector<TBlob>& inputs,
                            const std::vector<TBlob>& outputs,
                            const std::vector<int>& axises,
                            const nnvm::NodeAttrs& attrs,
                            const OpContext& ctx,
                            const OpReqType& req,
                            const Tensor<xpu, 1, DType>& workspace) {
    using namespace mshadow;
    CHECK_EQ(inputs.size(), 2U);
    CHECK_EQ(outputs.size(), 1U);
    CHECK_EQ(axises.size(), 3U);
    const TBlob& a = inputs[0];
    const TBlob& b = inputs[1];
    const TBlob& c = outputs[0];
    if (kNullOp == req) { return; }
    if (0U == a.Size() || 0U == b.Size()) { return; }
  
    const int a_axis = CheckAxis(axises[0], a.ndim());
    const int b_axis = CheckAxis(axises[1], b.ndim());
    const int c_axis = CheckAxis(axises[2], c.ndim());
    Tuple<int> a_move_index = GetMoveaxisIndex(a_axis, -1, a.shape_);
    Tuple<int> b_move_index = GetMoveaxisIndex(b_axis, -1, b.shape_);
    Tuple<int> c_move_index = GetMoveaxisIndex(c_axis, -1, c.shape_);
    mxnet::TShape a_move_shape = GetMoveaxisShape(a_move_index, a.shape_);
    mxnet::TShape b_move_shape = GetMoveaxisShape(b_move_index, b.shape_);
    mxnet::TShape c_move_shape = GetMoveaxisShape(c_move_index, c.shape_);
    // Check workspace size.
    size_t workspace_size = NumpyCrossWorkspaceSize<xpu, DType>(a_move_shape, b_move_shape,
                                                                c_move_shape, a_axis, b_axis,
                                                                ctx, { req });
    CHECK_GE(workspace.MSize(), workspace_size)
      << "Not enough working space size for cross product(should >= " << workspace_size << ")";
    NumpyCrossForwardImpl<xpu, DType, a_dim, b_dim>::op(a, b, c,
                                                        { a_move_index, b_move_index},
                                                        { a_move_shape, b_move_shape, c_move_shape},
                                                        a_axis, b_axis, c_axis,
                                                        attrs, ctx, { req }, workspace);
  }
  
  template<typename xpu, typename DType>
  struct ReduceImplWrap {
    static size_t wks(const mxnet::TShape& out_shape,
                      const mxnet::TShape& out_move_shape,
                      const mxnet::TShape& in_shape,
                      const mxnet::TShape& in_move_shape,
                      const OpContext& ctx, const OpReqType& req) {
      size_t ws_reduce = 0U;
      std::vector<int> reduce_axis = GetReduceAxis(out_move_shape, in_move_shape);
      if (reduce_axis.empty() || req == kNullOp) { return 0U; }
      ws_reduce = broadcast::ReduceWorkspaceSize(ctx.get_stream<xpu>(),
                                                 out_shape, req, in_shape,
                                                 sizeof(DType));
      return ws_reduce;
    }
  
    static void op(const TBlob& work_in,
                   const TBlob& work_out,
                   const TBlob& out_data,
                   const OpContext& ctx,
                   const OpReqType& out_req,
                   const Tensor<xpu, 1, char> workspace_tensor) {
      Stream<xpu> *s = ctx.get_stream<xpu>();
      // Reduce work_in to work_out.
      SUM_NDIM_SWITCH(work_out.ndim(), NDim, {
        op::broadcast::Reduce<mshadow_op::sum, NDim, DType, op::mshadow_op::identity, false>(
          s, work_out, kWriteTo, workspace_tensor, work_in);
      });
      // Copy work_out to out_data.
      MXNET_ASSIGN_REQ_SWITCH(out_req, req_type, {
        mxnet_op::Kernel<ResAssign<req_type>, xpu>::Launch(
          s, out_data.Size(), work_out.dptr<DType>(), out_data.dptr<DType>());
      });
    }
  };
  
  template<typename xpu, typename DType, int a_dim, int b_dim>
  struct NumpyCrossBackwardImpl {
    static std::vector<size_t>
    CrossBwkWorkspaceSize(const bool use_broadcast,
                          const mxnet::TShape& a_shape,
                          const mxnet::TShape& b_shape,
                          const mxnet::TShape& c_shape,
                          const mxnet::TShape& a_move_shape,
                          const mxnet::TShape& b_move_shape,
                          const mxnet::TShape& c_move_shape,
                          const int& a_axis, const int& b_axis, const int& c_axis,
                          const OpContext& ctx,
                          const std::vector<OpReqType>& req) {
      CHECK((a_dim == 2 && b_dim == 3) || (a_dim == 3 && b_dim == 2))
        << "no specialized NumpyCrossOp defined for template parameters.";
      std::vector<size_t> workspace_size(3, 0U);
      if (use_broadcast) {
        size_t ws_a = 0U, ws_b = 0U, rws_a = 0U, rws_b = 0U;
        size_t ws1 = NumpyCrossWorkspaceSize<xpu, DType>(b_move_shape, c_move_shape, c_move_shape,
                                                         b_axis, c_axis, ctx, { kWriteTo });
        size_t ws2 = NumpyCrossWorkspaceSize<xpu, DType>(b_move_shape, c_move_shape, a_move_shape,
                                                         b_axis, c_axis, ctx, { kWriteTo });
        ws_a = std::max(ws1, ws2);
        size_t ws3 = NumpyCrossWorkspaceSize<xpu, DType>(c_move_shape, a_move_shape, c_move_shape,
                                                         c_axis, a_axis, ctx, { kWriteTo });
        size_t ws4 = NumpyCrossWorkspaceSize<xpu, DType>(c_move_shape, a_move_shape, b_move_shape,
                                                         c_axis, a_axis, ctx, { kWriteTo });
        ws_b = std::max(ws3, ws4);
        // Get delete result shape.
        mxnet::TShape c_move_dshape = c_move_shape;
        c_move_dshape[c_move_shape.ndim() - 1] = 2;
        if (a_dim == 2) {
          mxnet::TShape c_dshape = GetOriShape(c_move_dshape, a_axis);
          // Calculate workspace size used in sum(grad_a).
          size_t rws1 = ReduceImplWrap<xpu, DType>::wks(c_dshape, c_move_dshape,
                                                        c_shape, c_move_shape, ctx, req[0]);
          // Calculate workspace size used in sum(grad_b).
          size_t rws2 = ReduceImplWrap<xpu, DType>::wks(b_shape, b_move_shape,
                                                        c_shape, c_move_shape, ctx, req[1]);
          rws_a = std::max(rws1, rws2);
        }
        if (b_dim == 2) {
          mxnet::TShape c_dshape = GetOriShape(c_move_dshape, b_axis);
          // Calculate workspace size used in sum(grad_a).
          size_t rws1 = ReduceImplWrap<xpu, DType>::wks(a_shape, a_move_shape,
                                                        c_shape, c_move_shape, ctx, req[0]);
          // Calculate workspace size used in sum(grad_b).
          size_t rws2 = ReduceImplWrap<xpu, DType>::wks(c_dshape, c_move_dshape,
                                                        c_shape, c_move_shape, ctx, req[1]);
          rws_b = std::max(rws1, rws2);
        }
        size_t rws = (std::max(rws_a, rws_b) + sizeof(DType) - 1) / sizeof(DType);
        workspace_size[0] += std::max(ws_a, ws_b);  // For cross product workspace.
        workspace_size[1] += c_move_shape.Size();   // For cross product result.
        workspace_size[1] += c_move_shape.Size();   // For delete result shape.
        workspace_size[2] += rws;                   // For reduce workspace.
      } else {
        mxnet::TShape a_moveaxis_shape = (a_dim == 2 ? c_move_shape : a_move_shape);
        mxnet::TShape b_moveaxis_shape = (b_dim == 2 ? c_move_shape : b_move_shape);
        size_t ws1 = NumpyCrossWorkspaceSize<xpu, DType>(b_moveaxis_shape, c_move_shape,
                                                         a_moveaxis_shape, b_axis, c_axis, ctx,
                                                         { kWriteTo });
        size_t ws2 = NumpyCrossWorkspaceSize<xpu, DType>(c_move_shape, a_moveaxis_shape,
                                                         b_moveaxis_shape, c_axis, a_axis, ctx,
                                                         { kWriteTo });
        workspace_size[0] += std::max(ws1, ws2);  // For cross product workspace.
        if (a_dim == 2 && b_dim == 3) {
          workspace_size[1] += a_moveaxis_shape.Size();  // For cross product result.
          workspace_size[1] += a_move_shape.Size();      // For delete kernel result.
        }
        if (a_dim == 3 && b_dim == 2) {
          workspace_size[1] += b_moveaxis_shape.Size();  // For cross product result.
          workspace_size[1] += b_move_shape.Size();      // For delete kernel result.
        }
      }
      return workspace_size;
    }
  
    static void op(const bool use_broadcast,
                   const TBlob& grad_c, const TBlob& a, const TBlob& b,
                   const TBlob& grad_a, const TBlob& grad_b,
                   const std::vector<mxnet::TShape>& moveaxis_shape_vec,
                   const int a_axis, const int b_axis, const int c_axis,
                   const nnvm::NodeAttrs &attrs,
                   const OpContext& ctx,
                   const std::vector<OpReqType>& req) {
      CHECK((a_dim == 2 && b_dim == 3) || (a_dim == 3 && b_dim == 2))
        << "no specialized NumpyCrossOp defined for template parameters.";
      Stream<xpu> *s = ctx.get_stream<xpu>();
      const mxnet::TShape& a_move_shp = moveaxis_shape_vec[0];
      const mxnet::TShape& b_move_shp = moveaxis_shape_vec[1];
      const mxnet::TShape& c_move_shp = moveaxis_shape_vec[2];
      std::vector<int> a_reduce_axis = GetReduceAxis(a_move_shp, c_move_shp);
      std::vector<int> b_reduce_axis = GetReduceAxis(b_move_shp, c_move_shp);
      const int c_ndim = c_move_shp.ndim();
      // Get delete result shape.
      mxnet::TShape c_move_dshp = c_move_shp;
      c_move_dshp[c_move_shp.ndim() - 1] = 2;
      std::vector<size_t> wk_size = CrossBwkWorkspaceSize(use_broadcast,
                                                          a.shape_, b.shape_, grad_c.shape_,
                                                          a_move_shp, b_move_shp, c_move_shp,
                                                          a_axis, b_axis, c_axis, ctx, req);
      Tensor<xpu, 1, DType> workspace = ctx.requested[0].get_space_typed<xpu, 1, DType>(
        Shape1(wk_size[0] + wk_size[1] + wk_size[2]), s);
      if (use_broadcast) {
        // Use broadcast in forward, need reduce in backward.
        DType *w0_ptr = workspace.dptr_;
        DType *w1_ptr = w0_ptr + wk_size[0];
        DType *w2_ptr = w1_ptr + c_move_shp.Size();
        char *w3_ptr = reinterpret_cast<char*>(w2_ptr + c_move_shp.Size());
        TBlob w0_data(w0_ptr, Shape1(wk_size[0]), grad_c.dev_mask(), grad_c.dev_id());
        Tensor<xpu, 1, char> w3_tensor(w3_ptr, Shape1(wk_size[2] * sizeof(DType)), s);
        if (a_dim == 2) {  // a_dim == 2, b_dim == 3
          TBlob w1_data(w1_ptr, c_move_shp, grad_c.dev_mask(), grad_c.dev_id());
          TBlob w2_data(w2_ptr, c_move_dshp, grad_c.dev_mask(), grad_c.dev_id());
          // Calculate grad_a = cross(b, grad_c).
          CrossImplWrap<xpu, DType, 3, 3>({ b, grad_c }, { w1_data }, { b_axis, c_axis, -1 },
                                          attrs, ctx, kWriteTo, w0_data.get<xpu, 1, DType>(s));
          // Copy w1_data to w2_data with delete.
          mxnet_op::Kernel<DeleteAssign, xpu>::Launch(s, c_move_dshp.ProdShape(0, c_ndim - 1),
                                                      w1_data.dptr<DType>(),
                                                      w2_data.dptr<DType>(), 3, 2);
          // Transpose w2_data to w1_data.
          if (a_axis != grad_a.ndim() - 1) {
            const Tuple<int> axis_idx = GetMoveaxisIndex(-1, a_axis, c_move_dshp);
            mxnet::TShape c_dshp = GetMoveaxisShape(axis_idx, c_move_dshp);
            w1_data = TBlob(w1_ptr, c_dshp, grad_c.dev_mask(), grad_c.dev_id());
            TransposeImpl<xpu>(ctx.run_ctx, w2_data, w1_data,
                               mxnet::TShape(axis_idx.begin(), axis_idx.end()));
            w2_data = TBlob(w2_ptr, grad_a.shape_, grad_c.dev_mask(), grad_c.dev_id());
          } else {
            // If no transpose, exchange the pointer.
            w1_data = TBlob(w2_ptr, c_move_dshp, grad_c.dev_mask(), grad_c.dev_id());
            w2_data = TBlob(w1_ptr, grad_a.shape_, grad_c.dev_mask(), grad_c.dev_id());
          }
          // Reduce w1_data to w2_data.
          if (a_reduce_axis.empty()) {
            // No need Reduce w1_data, Copy w1_data to grad_a.
            MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
              mxnet_op::Kernel<ResAssign<req_type>, xpu>::Launch(
                s, grad_a.Size(), w1_data.dptr<DType>(), grad_a.dptr<DType>());
            });
          } else {
            // Need Reduce w1_data to w2_data and Copy w2_data to grad_a.
            ReduceImplWrap<xpu, DType>::op(w1_data, w2_data, grad_a, ctx, req[0], w3_tensor);
          }
          // Calculate grad_b = cross(grad_c, a).
          if (b_reduce_axis.empty()) {
            CrossImplWrap<xpu, DType, 3, 2>({ grad_c, a }, { grad_b }, { c_axis, a_axis, b_axis },
                                            attrs, ctx, req[1], w0_data.get<xpu, 1, DType>(s));
          } else {
            mxnet::TShape c_shp = GetOriShape(c_move_shp, b_axis);
            w1_data = TBlob(w1_ptr, c_shp, grad_c.dev_mask(), grad_c.dev_id());
            w2_data = TBlob(w2_ptr, grad_b.shape_, grad_c.dev_mask(), grad_c.dev_id());
            CrossImplWrap<xpu, DType, 3, 2>({ grad_c, a }, { w1_data }, { c_axis, a_axis, b_axis },
                                            attrs, ctx, req[1], w0_data.get<xpu, 1, DType>(s));
            // Need Reduce w1_data to w2_data and Copy w2_data to grad_b.
            ReduceImplWrap<xpu, DType>::op(w1_data, w2_data, grad_b, ctx, req[1], w3_tensor);
          }
        }  // End of a_dim == 2
        if (b_dim == 2) {  // a_dim == 3, b_dim == 2
          TBlob w1_data(w1_ptr, c_move_shp, grad_c.dev_mask(), grad_c.dev_id());
          TBlob w2_data(w2_ptr, c_move_dshp, grad_c.dev_mask(), grad_c.dev_id());
          // Calculate grad_b = cross(grad_c, a).
          CrossImplWrap<xpu, DType, 3, 3>({ grad_c, a }, { w1_data }, { c_axis, a_axis, -1 },
                                          attrs, ctx, kWriteTo, w0_data.get<xpu, 1, DType>(s));
          // Copy w1_data to w2_data with delete.
          mxnet_op::Kernel<DeleteAssign, xpu>::Launch(s, c_move_dshp.ProdShape(0, c_ndim - 1),
                                                      w1_data.dptr<DType>(),
                                                      w2_data.dptr<DType>(), 3, 2);
          // Transpose w2_data to w1_data.
          if (b_axis != grad_b.ndim() - 1) {
            const Tuple<int> axis_idx = GetMoveaxisIndex(-1, b_axis, c_move_dshp);
            mxnet::TShape c_dshp = GetMoveaxisShape(axis_idx, c_move_dshp);
            w1_data = TBlob(w1_ptr, c_dshp, grad_c.dev_mask(), grad_c.dev_id());
            TransposeImpl<xpu>(ctx.run_ctx, w2_data, w1_data,
                               mxnet::TShape(axis_idx.begin(), axis_idx.end()));
            w2_data = TBlob(w2_ptr, grad_b.shape_, grad_c.dev_mask(), grad_c.dev_id());
          } else {
            // If no transpose, exchange the pointer.
            w1_data = TBlob(w2_ptr, c_move_dshp, grad_c.dev_mask(), grad_c.dev_id());
            w2_data = TBlob(w1_ptr, grad_b.shape_, grad_c.dev_mask(), grad_c.dev_id());
          }
          // Reduce w1_data to w2_data.
          if (b_reduce_axis.empty()) {
            // No need Reduce w1_data, Copy w1_data to grad_b.
            MXNET_ASSIGN_REQ_SWITCH(req[1], req_type, {
              mxnet_op::Kernel<ResAssign<req_type>, xpu>::Launch(
                s, grad_b.Size(), w1_data.dptr<DType>(), grad_b.dptr<DType>());
            });
          } else {
            // Need Reduce w1_data to w2_data and Copy w2_data to grad_a.
            ReduceImplWrap<xpu, DType>::op(w1_data, w2_data, grad_b, ctx, req[1], w3_tensor);
          }
          // Calculate grad_a = cross(b, grad_c).
          if (a_reduce_axis.empty()) {
            CrossImplWrap<xpu, DType, 2, 3>({ b, grad_c }, { grad_a }, { b_axis, c_axis, a_axis },
                                            attrs, ctx, req[0], workspace);
          } else {
            mxnet::TShape c_shp = GetOriShape(c_move_shp, a_axis);
            w1_data = TBlob(w1_ptr, c_shp, grad_c.dev_mask(), grad_c.dev_id());
            w2_data = TBlob(w2_ptr, grad_a.shape_, grad_c.dev_mask(), grad_c.dev_id());
            CrossImplWrap<xpu, DType, 2, 3>({ b, grad_c }, { w1_data }, { b_axis, c_axis, a_axis },
                                            attrs, ctx, req[0], w0_data.get<xpu, 1, DType>(s));
            // Need Reduce w1_data to w2_data and Copy w2_data to grad_b.
            ReduceImplWrap<xpu, DType>::op(w1_data, w2_data, grad_a, ctx, req[0], w3_tensor);
          }
        }  // End of b_dim == 3
      } else {
        // No use broadcast in forward, not need reduce in backward.
        DType *w0_ptr = workspace.dptr_;
        DType *w1_ptr = w0_ptr + wk_size[0];
        DType *w2_ptr = w1_ptr + c_move_shp.Size();
        TBlob w0_data(w0_ptr, Shape1(wk_size[0]), grad_c.dev_mask(), grad_c.dev_id());
        if (a_dim == 2) {  // a_dim == 2, b_dim == 3
          TBlob w1_data(w1_ptr, c_move_shp, grad_c.dev_mask(), grad_c.dev_id());
          TBlob w2_data(w2_ptr, a_move_shp, grad_c.dev_mask(), grad_c.dev_id());
          // Calculate w1_data = cross(b, grad_c).
          CrossImplWrap<xpu, DType, 3, 3>({ b, grad_c }, { w1_data }, { b_axis, c_axis, -1 },
                                          attrs, ctx, kWriteTo, w0_data.get<xpu, 1, DType>(s));
          // Calculate grad_b = cross(grad_c, a).
          CrossImplWrap<xpu, DType, 3, 2>({ grad_c, a }, { grad_b }, { c_axis, a_axis, b_axis },
                                          attrs, ctx, req[1], w0_data.get<xpu, 1, DType>(s));
          // Copy w1_data to w2_data with delete.
          mxnet_op::Kernel<DeleteAssign, xpu>::Launch(s, a_move_shp.ProdShape(0, a.ndim() - 1),
                                                      w1_data.dptr<DType>(),
                                                      w2_data.dptr<DType>(), 3, 2);
          DType *res_ptr = w2_data.dptr<DType>();
          if (a_axis != grad_a.ndim() - 1) {
            // Transpose w2_data to w1_data.
            const Tuple<int> grad_a_axis_idx = GetMoveaxisIndex(-1, a_axis, a_move_shp);
            w1_data = TBlob(w1_ptr, grad_a.shape_, grad_c.dev_mask(), grad_c.dev_id());
            TransposeImpl<xpu>(ctx.run_ctx, w2_data, w1_data,
                               mxnet::TShape(grad_a_axis_idx.begin(), grad_a_axis_idx.end()));
            res_ptr = w1_data.dptr<DType>();
          }
          // Copy w1_data to grad_a.
          MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
            mxnet_op::Kernel<ResAssign<req_type>, xpu>::Launch(
              s, grad_a.Size(), res_ptr, grad_a.dptr<DType>());
          });
        }  // End of a_dim == 2
        if (b_dim == 2) {  // a_dim == 3, b_dim == 2
          TBlob w1_data(w1_ptr, c_move_shp, grad_c.dev_mask(), grad_c.dev_id());
          TBlob w2_data(w2_ptr, b_move_shp, grad_c.dev_mask(), grad_c.dev_id());
          // Calculate grad_a = cross(b, grad_c).
          CrossImplWrap<xpu, DType, 2, 3>({ b, grad_c }, { grad_a }, { b_axis, c_axis, a_axis },
                                          attrs, ctx, req[0], w0_data.get<xpu, 1, DType>(s));
          // Calculate w1_data = cross(grad_c, a).
          CrossImplWrap<xpu, DType, 3, 3>({ grad_c, a }, { w1_data }, { c_axis, a_axis, -1 },
                                          attrs, ctx, kWriteTo, w0_data.get<xpu, 1, DType>(s));
          // Copy w1_data to w2_data with delete.
          mxnet_op::Kernel<DeleteAssign, xpu>::Launch(s, b_move_shp.ProdShape(0, b.ndim() - 1),
                                                      w1_data.dptr<DType>(),
                                                      w2_data.dptr<DType>(), 3, 2);
          DType *res_ptr = w2_data.dptr<DType>();
          if (b_axis != grad_b.ndim() - 1) {
            // Transpose w2_data to w1_data.
            const Tuple<int> grad_b_axis_idx = GetMoveaxisIndex(-1, b_axis, b_move_shp);
            w1_data = TBlob(w1_ptr, grad_b.shape_, grad_c.dev_mask(), grad_c.dev_id());
            TransposeImpl<xpu>(ctx.run_ctx, w2_data, w1_data,
                               mxnet::TShape(grad_b_axis_idx.begin(), grad_b_axis_idx.end()));
            res_ptr = w1_data.dptr<DType>();
          }
          // Copy w1_data to grad_b.
          MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
            mxnet_op::Kernel<ResAssign<req_type>, xpu>::Launch(
              s, grad_b.Size(), res_ptr, grad_b.dptr<DType>());
          });
        }  // End of b_dim == 2
      }  // End of use_broadcast
    }
  };
  
  template<typename xpu, typename DType>
  struct NumpyCrossBackwardImpl<xpu, DType, 3, 3> {
    static std::vector<size_t>
    CrossBwkWorkspaceSize(const bool use_broadcast,
                          const mxnet::TShape& a_shape,
                          const mxnet::TShape& b_shape,
                          const mxnet::TShape& c_shape,
                          const mxnet::TShape& a_move_shape,
                          const mxnet::TShape& b_move_shape,
                          const mxnet::TShape& c_move_shape,
                          const int& a_axis, const int& b_axis, const int& c_axis,
                          const OpContext& ctx,
                          const std::vector<OpReqType>& req) {
      std::vector<size_t> workspace_size(3, 0U);
      if (use_broadcast) {
        // Calculate workspace size used in cross(b_move, grad_c_move).
        size_t ws1 = NumpyCrossWorkspaceSize<xpu, DType>(b_move_shape, c_move_shape, c_move_shape,
                                                         b_axis, c_axis, ctx, { kWriteTo });
        // Calculate workspace size used in cross(grad_c_move, a_move).
        size_t ws2 = NumpyCrossWorkspaceSize<xpu, DType>(c_move_shape, a_move_shape, c_move_shape,
                                                         c_axis, a_axis, ctx, { kWriteTo });
        // Calculate workspace size used in sum(grad_a).
        size_t rws1 = ReduceImplWrap<xpu, DType>::wks(a_shape, a_move_shape,
                                                      c_shape, c_move_shape, ctx, req[0]);
        // Calculate workspace size used in sum(grad_b).
        size_t rws2 = ReduceImplWrap<xpu, DType>::wks(b_shape, b_move_shape,
                                                      c_shape, c_move_shape, ctx, req[1]);
        // For cross product workspace.
        workspace_size[0] += std::max(ws1, ws2);
        // For reduce workspace.
        workspace_size[1] += (std::max(rws1, rws2) + sizeof(DType) - 1) / sizeof(DType);
        // For cross result and reduce result.
        workspace_size[2] += c_move_shape.Size();
        workspace_size[2] += std::max(a_move_shape.Size(), b_move_shape.Size());
      } else {
        size_t ws1 = NumpyCrossWorkspaceSize<xpu, DType>(b_move_shape, c_move_shape, a_move_shape,
                                                         b_axis, c_axis, ctx, { req[0] });
        size_t ws2 = NumpyCrossWorkspaceSize<xpu, DType>(c_move_shape, a_move_shape, b_move_shape,
                                                         c_axis, a_axis, ctx, { req[1] });
        workspace_size[0] += std::max(ws1, ws2);  // For cross product workspace.
      }
      return workspace_size;
    }
  
    static void op(const bool use_broadcast,
                   const TBlob& grad_c, const TBlob& a, const TBlob& b,
                   const TBlob& grad_a, const TBlob& grad_b,
                   const std::vector<mxnet::TShape>& moveaxis_shape_vec,
                   const int a_axis, const int b_axis, const int c_axis,
                   const nnvm::NodeAttrs &attrs,
                   const OpContext& ctx,
                   const std::vector<OpReqType>& req) {
      Stream<xpu> *s = ctx.get_stream<xpu>();
      const mxnet::TShape& a_move_shp = moveaxis_shape_vec[0];
      const mxnet::TShape& b_move_shp = moveaxis_shape_vec[1];
      const mxnet::TShape& c_move_shp = moveaxis_shape_vec[2];
      std::vector<size_t> wk_size = CrossBwkWorkspaceSize(use_broadcast,
                                                          a.shape_, b.shape_, grad_c.shape_,
                                                          a_move_shp, b_move_shp, c_move_shp,
                                                          a_axis, b_axis, c_axis, ctx, req);
      Tensor<xpu, 1, DType> workspace = ctx.requested[0].get_space_typed<xpu, 1, DType>(
        Shape1(wk_size[0] + wk_size[1] + wk_size[2]), s);
      if (use_broadcast) {
        // Use broadcast in forward, need reduce in backward.
        std::vector<int> a_reduce_axis = GetReduceAxis(a_move_shp, c_move_shp);
        std::vector<int> b_reduce_axis = GetReduceAxis(b_move_shp, c_move_shp);
        // Allocate workspace.
        DType *w0_ptr = workspace.dptr_;
        char *w1_ptr = reinterpret_cast<char*>(w0_ptr + wk_size[0]);
        TBlob w0_data(w0_ptr, Shape1(wk_size[0]), grad_c.dev_mask(), grad_c.dev_id());
        Tensor<xpu, 1, char> w1_tensor(w1_ptr, Shape1(wk_size[1] * sizeof(DType)), s);
        if (a_reduce_axis.empty()) {
          // Calculate grad_a = cross(b, grad_c).
          CrossImplWrap<xpu, DType, 3, 3>({ b, grad_c }, { grad_a }, { b_axis, c_axis, a_axis },
                                          attrs, ctx, req[0], w0_data.get<xpu, 1, DType>(s));
        } else {
          mxnet::TShape c_shp = GetOriShape(c_move_shp, a_axis);
          DType *w2_ptr = w0_ptr + wk_size[0] + wk_size[1];
          DType *w3_ptr = w2_ptr + c_move_shp.Size();
          TBlob w2_data(w2_ptr, c_shp, grad_c.dev_mask(), grad_c.dev_id());
          TBlob w3_data(w3_ptr, grad_a.shape_, grad_c.dev_mask(), grad_c.dev_id());
          // Calculate w2_data = cross(b, grad_c).
          CrossImplWrap<xpu, DType, 3, 3>({ b, grad_c }, { w2_data }, { b_axis, c_axis, a_axis },
                                          attrs, ctx, kWriteTo, w0_data.get<xpu, 1, DType>(s));
          // Reduce w2_data to w3_data and Copy w3_data to grad_a.
          ReduceImplWrap<xpu, DType>::op(w2_data, w3_data, grad_a, ctx, req[0], w1_tensor);
        }
        if (b_reduce_axis.empty()) {
          // Calculate grad_b = cross(grad_c, a).
          CrossImplWrap<xpu, DType, 3, 3>({ grad_c, a }, { grad_b }, { c_axis, a_axis, b_axis },
                                          attrs, ctx, req[1], w0_data.get<xpu, 1, DType>(s));
        } else {
          mxnet::TShape c_shp = GetOriShape(c_move_shp, b_axis);
          DType *w2_ptr = w0_ptr + wk_size[0] + wk_size[1];
          DType *w3_ptr = w2_ptr + c_move_shp.Size();
          TBlob w2_data(w2_ptr, c_shp, grad_c.dev_mask(), grad_c.dev_id());
          TBlob w3_data(w3_ptr, grad_b.shape_, grad_c.dev_mask(), grad_c.dev_id());
          // Calculate w2_data = cross(grad_c, a).
          CrossImplWrap<xpu, DType, 3, 3>({ grad_c, a }, { w2_data }, { c_axis, a_axis, b_axis },
                                          attrs, ctx, kWriteTo, w0_data.get<xpu, 1, DType>(s));
          // Reduce w2_data to w3_data and Copy w3_data to grad_b.
          ReduceImplWrap<xpu, DType>::op(w2_data, w3_data, grad_b, ctx, req[1], w1_tensor);
        }
      } else {
        CrossImplWrap<xpu, DType, 3, 3>({ b, grad_c }, { grad_a }, { b_axis, c_axis, a_axis },
                                        attrs, ctx, req[0], workspace);
        CrossImplWrap<xpu, DType, 3, 3>({ grad_c, a }, { grad_b }, { c_axis, a_axis, b_axis },
                                        attrs, ctx, req[1], workspace);
      }
    }
  };
  
  template<typename xpu, typename DType>
  struct NumpyCrossBackwardImpl<xpu, DType, 2, 2> {
    static std::vector<size_t>
    CrossBwkWorkspaceSize(const bool use_broadcast,
                          const mxnet::TShape& a_shape,
                          const mxnet::TShape& b_shape,
                          const mxnet::TShape& c_shape,
                          const mxnet::TShape& a_move_shape,
                          const mxnet::TShape& b_move_shape,
                          const mxnet::TShape& c_move_shape,
                          const int& a_axis, const int& b_axis, const int& c_axis,
                          const OpContext& ctx,
                          const std::vector<OpReqType>& req) {
      std::vector<size_t> workspace_size(3, 0U);
      const int a_ndim = a_move_shape.ndim();
      const int b_ndim = b_move_shape.ndim();
      const int c_ndim = c_move_shape.ndim();
      mxnet::TShape grad_move_shape(c_ndim + 1, 2);
      for (int i = 0; i < c_ndim; ++i) { grad_move_shape[i] = c_move_shape[i]; }
  
      workspace_size[0] += grad_move_shape.Size();                 // For grad_a_move or grad_b_move.
      workspace_size[0] += a_move_shape.ProdShape(0, a_ndim - 1);  // For a_move work data.
      workspace_size[0] += b_move_shape.ProdShape(0, b_ndim - 1);  // For b_move work data.
      workspace_size[0] =                                          // For c_move work data.
        AligndWorkspaceSize<DType>(workspace_size[0], c_move_shape.Size());
  
      if (a_axis != a_ndim -1 || b_axis != b_ndim - 1) {
        if (ctx.run_ctx.get_ctx().dev_mask() == cpu::kDevMask) {
          workspace_size[1] += a_move_shape.Size();     // For a_move size.
          workspace_size[1] += b_move_shape.Size();     // For b_move size.
          workspace_size[1] += grad_move_shape.Size();  // For grad_a_move or grad_b_move trans.
        } else {
          workspace_size[1] = AligndWorkspaceSize<DType>(workspace_size[1], a_move_shape.Size());
          workspace_size[1] = AligndWorkspaceSize<DType>(workspace_size[1], b_move_shape.Size());
          workspace_size[1] = AligndWorkspaceSize<DType>(workspace_size[1], grad_move_shape.Size());
        }
      }
      if (use_broadcast) {
        mxnet::TShape grad_a_dshape = GetOriShape(grad_move_shape, a_axis);
        mxnet::TShape grad_b_dshape = GetOriShape(grad_move_shape, b_axis);
        size_t rws1 = ReduceImplWrap<xpu, DType>::wks(a_shape, a_move_shape,
                                                      grad_a_dshape, grad_move_shape, ctx, req[0]);
        size_t rws2 = ReduceImplWrap<xpu, DType>::wks(b_shape, b_move_shape,
                                                      grad_b_dshape, grad_move_shape, ctx, req[1]);
        size_t rws = (std::max(rws1, rws2) + sizeof(DType) - 1) / sizeof(DType);
        workspace_size[2] += std::max(a_shape.Size(), b_shape.Size());  // For reduce result.
        workspace_size[2] += rws;                                       // For reduce workspace.
      }
      return workspace_size;
    }
  
    static void op(const bool use_broadcast,
                   const TBlob& grad_c, const TBlob& a, const TBlob& b,
                   const TBlob& grad_a, const TBlob& grad_b,
                   const std::vector<mxnet::TShape>& moveaxis_shape_vec,
                   const int a_axis, const int b_axis, const int c_axis,
                   const nnvm::NodeAttrs &attrs,
                   const OpContext& ctx,
                   const std::vector<OpReqType>& req) {
      Stream<xpu> *s = ctx.get_stream<xpu>();
      Tuple<int> a_move_idx = GetMoveaxisIndex(a_axis, -1, a.shape_);
      Tuple<int> b_move_idx = GetMoveaxisIndex(b_axis, -1, b.shape_);
      const mxnet::TShape& a_move_shp = moveaxis_shape_vec[0];
      const mxnet::TShape& b_move_shp = moveaxis_shape_vec[1];
      const mxnet::TShape& c_move_shp = grad_c.shape_;
      std::vector<int> a_reduce_axis = GetReduceAxis(a_move_shp, c_move_shp);
      std::vector<int> b_reduce_axis = GetReduceAxis(b_move_shp, c_move_shp);
      const int a_ndim = a_move_shp.ndim();
      const int b_ndim = b_move_shp.ndim();
      const int c_ndim = c_move_shp.ndim();
      mxnet::TShape grad_move_shp(c_ndim + 1, 2);
      for (int i = 0; i < c_ndim; ++i) { grad_move_shp[i] = c_move_shp[i]; }
      // Calculate workspace size.
      std::vector<size_t> wk_size = CrossBwkWorkspaceSize(use_broadcast,
                                                          a.shape_, b.shape_, grad_c.shape_,
                                                          a_move_shp, b_move_shp, c_move_shp,
                                                          a_axis, b_axis, c_axis, ctx, req);
      Tensor<xpu, 1, DType> workspace = ctx.requested[0].get_space_typed<xpu, 1, DType>(
        Shape1(wk_size[0] + wk_size[1] + wk_size[2]), s);
  
      // Allocate workspace in cpu, no need to align address.
      DType *grad_ptr = workspace.dptr_;
      DType *aw_ptr = grad_ptr + grad_move_shp.Size();
      DType *bw_ptr = aw_ptr + a_move_shp.ProdShape(0, a_ndim - 1);
      DType *cw_ptr = bw_ptr + b_move_shp.ProdShape(0, b_ndim - 1);
      TBlob grad_move_data(grad_ptr, grad_move_shp, grad_c.dev_mask(), grad_c.dev_id());
      TBlob aw_data(aw_ptr, GetCutoffShape(a_move_shp), grad_c.dev_mask(), grad_c.dev_id());
      TBlob bw_data(bw_ptr, GetCutoffShape(b_move_shp), grad_c.dev_mask(), grad_c.dev_id());
      TBlob cw_data(cw_ptr, c_move_shp, grad_c.dev_mask(), grad_c.dev_id());
      TBlob a_move_data = a;
      TBlob b_move_data = b;
      TBlob grad_data = grad_move_data;
      size_t offset = 0;
      if (a_axis != a_ndim -1 || b_axis != b_ndim - 1) {
        if (ctx.run_ctx.get_ctx().dev_mask() == cpu::kDevMask) {
          DType *a_ptr = workspace.dptr_ + wk_size[0];
          DType *b_ptr = a_ptr + a_move_shp.Size();
          a_move_data = TBlob(a_ptr, a_move_shp, a.dev_mask(), a.dev_id());
          b_move_data = TBlob(b_ptr, b_move_shp, b.dev_mask(), b.dev_id());
          TransposeImpl<xpu>(ctx.run_ctx, a, a_move_data,
                             mxnet::TShape(a_move_idx.begin(), a_move_idx.end()));
          TransposeImpl<xpu>(ctx.run_ctx, b, b_move_data,
                             mxnet::TShape(b_move_idx.begin(), b_move_idx.end()));
        } else {
          DType *w1_ptr = workspace.dptr_ + wk_size[0];
          a_move_data = TBlob(w1_ptr + offset, a_move_shp, a.dev_mask(), a.dev_id());
          offset = AligndWorkspaceSize<DType>(offset, a_move_shp.Size());
  
          b_move_data = TBlob(w1_ptr + offset, b_move_shp, a.dev_mask(), a.dev_id());
          offset = AligndWorkspaceSize<DType>(offset, a_move_shp.Size());
  
          TransposeImpl<xpu>(ctx.run_ctx, a, a_move_data,
                             mxnet::TShape(a_move_idx.begin(), a_move_idx.end()));
          TransposeImpl<xpu>(ctx.run_ctx, b, b_move_data,
                             mxnet::TShape(b_move_idx.begin(), b_move_idx.end()));
        }
      }
      // Copy b_move_data[..., 1] to bw_data.
      mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, bw_data.Size(),
                                                   b_move_data.dptr<DType>(),
                                                   bw_data.dptr<DType>(),
                                                   b_move_data.size(b_ndim - 1),
                                                   1, b_move_data.Size());
      // cw_data = grad_c_move * b_move_data[..., 1].
  #if !defined(__CUDACC__)
      BinaryBroadcastCompute<xpu, op::mshadow_op::mul>(attrs, ctx, { grad_c, bw_data },
                                                       { kWriteTo }, { cw_data });
  #else
      BinaryBroadcastRTCCompute {"mul"}(attrs, ctx, { grad_c, bw_data },
                                        { kWriteTo }, { cw_data });
  #endif  // !defined(__CUDACC__)
      // Copy cw_data to grad_move_data[..., 0].
      mxnet_op::Kernel<CrossOutAssign<kWriteTo>, xpu>::Launch(s, cw_data.Size(),
                                                              cw_data.dptr<DType>(),
                                                              grad_move_data.dptr<DType>(),
                                                              true, grad_move_data.size(c_ndim),
                                                              0, grad_move_data.Size());
      // Copy b_move_data[..., 0] to bw_data.
      mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, bw_data.Size(),
                                                   b_move_data.dptr<DType>(),
                                                   bw_data.dptr<DType>(),
                                                   b_move_data.size(b_ndim - 1),
                                                   0, b_move_data.Size());
      // cw_data = grad_c_move * b_move_data[..., 0].
  #if !defined(__CUDACC__)
      BinaryBroadcastCompute<xpu, op::mshadow_op::mul>(attrs, ctx, { grad_c, bw_data },
                                                       { kWriteTo }, { cw_data });
  #else
      BinaryBroadcastRTCCompute {"mul"}(attrs, ctx, { grad_c, bw_data },
                                        { kWriteTo }, { cw_data });
  #endif  // !defined(__CUDACC__)
      // Copy -cw_data to grad_move_data[..., 1].
      mxnet_op::Kernel<CrossOutAssign<kWriteTo>, xpu>::Launch(s, cw_data.Size(),
                                                              cw_data.dptr<DType>(),
                                                              grad_move_data.dptr<DType>(),
                                                              false, grad_move_data.size(c_ndim),
                                                              1, grad_move_data.Size());
      // Transpose grad_move_data according to a_axis.
      grad_data = grad_move_data;
      if (a_axis != a_ndim - 1) {
        mxnet::TShape grad_shp = GetOriShape(grad_move_shp, a_axis);
        if (ctx.run_ctx.get_ctx().dev_mask() == cpu::kDevMask) {
          DType *grad_ptr = workspace.dptr_ + wk_size[0] + a_move_shp.Size() + b_move_shp.Size();
          grad_data = TBlob(grad_ptr, grad_shp, grad_c.dev_mask(), grad_c.dev_id());
        } else {
          DType *w1_ptr = workspace.dptr_ + wk_size[0];
          grad_data = TBlob(w1_ptr + offset, grad_shp, grad_c.dev_mask(), grad_c.dev_id());
          offset = AligndWorkspaceSize<DType>(offset, grad_shp.Size());
        }
        const Tuple<int> axis_idx = GetMoveaxisIndex(-1, a_axis, grad_move_shp);
        TransposeImpl<xpu>(ctx.run_ctx, grad_move_data, grad_data,
                           mxnet::TShape(axis_idx.begin(), axis_idx.end()));
      }
      if (!a_reduce_axis.empty()) {
        size_t interval = std::max(grad_a.Size(), grad_b.Size());
        DType *grad_delete_ptr = workspace.dptr_ + wk_size[0] + wk_size[1];
        char *dw_ptr = reinterpret_cast<char*>(grad_delete_ptr + interval);
        TBlob grad_delete_data(grad_delete_ptr, grad_a.shape_, grad_c.dev_mask(), grad_c.dev_id());
        Tensor<xpu, 1, char> dw_tensor(dw_ptr, Shape1((wk_size[2] - interval) * sizeof(DType)), s);
        // Reduce grad_data to grad_delete_data and copy to grad_a.
        ReduceImplWrap<xpu, DType>::op(grad_data, grad_delete_data, grad_a, ctx, req[0], dw_tensor);
      } else {
        MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
          mxnet_op::Kernel<ResAssign<req_type>, xpu>::Launch(
            s, grad_a.Size(), grad_data.dptr<DType>(), grad_a.dptr<DType>());
        });
      }
  
      // Copy a_move_data[..., 1] to aw_data.
      mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, aw_data.Size(),
                                                   a_move_data.dptr<DType>(),
                                                   aw_data.dptr<DType>(),
                                                   a_move_data.size(a_ndim - 1),
                                                   1, a_move_data.Size());
      // cw_data = grad_c_move * a_move_data[..., 1].
  #if !defined(__CUDACC__)
      BinaryBroadcastCompute<xpu, op::mshadow_op::mul>(attrs, ctx, { grad_c, aw_data },
                                                       { kWriteTo }, { cw_data });
  #else
      BinaryBroadcastRTCCompute {"mul"}(attrs, ctx, { grad_c, aw_data },
                                        { kWriteTo }, { cw_data });
  #endif  // !defined(__CUDACC__)
      // Copy -cw_data to grad_move_data[..., 0].
      mxnet_op::Kernel<CrossOutAssign<kWriteTo>, xpu>::Launch(s, cw_data.Size(),
                                                              cw_data.dptr<DType>(),
                                                              grad_move_data.dptr<DType>(),
                                                              false, grad_move_data.size(c_ndim),
                                                              0, grad_move_data.Size());
      // Copy a_move_data[..., 0] to aw_data.
      mxnet_op::Kernel<CrossInAssign, xpu>::Launch(s, aw_data.Size(),
                                                   a_move_data.dptr<DType>(),
                                                   aw_data.dptr<DType>(),
                                                   a_move_data.size(a_ndim - 1),
                                                   0, a_move_data.Size());
      // cw_data = grad_c_move * a_move_data[..., 0].
  #if !defined(__CUDACC__)
      BinaryBroadcastCompute<xpu, op::mshadow_op::mul>(attrs, ctx, { grad_c, aw_data },
                                                       { kWriteTo }, { cw_data });
  #else
      BinaryBroadcastRTCCompute {"mul"}(attrs, ctx, { grad_c, aw_data },
                                        { kWriteTo }, { cw_data });
  #endif  // !defined(__CUDACC__)
      // Copy cw_data to grad_move_data[..., 1].
      mxnet_op::Kernel<CrossOutAssign<kWriteTo>, xpu>::Launch(s, cw_data.Size(),
                                                              cw_data.dptr<DType>(),
                                                              grad_move_data.dptr<DType>(),
                                                              true, grad_move_data.size(c_ndim),
                                                              1, grad_move_data.Size());
      // Transpose grad_move_data according to b_axis.
      grad_data = grad_move_data;
      if (b_axis != b_ndim - 1) {
        mxnet::TShape grad_shp = GetOriShape(grad_move_shp, b_axis);
        if (ctx.run_ctx.get_ctx().dev_mask() == cpu::kDevMask) {
          DType *grad_ptr = workspace.dptr_ + wk_size[0] + a_move_shp.Size() + b_move_shp.Size();
          grad_data = TBlob(grad_ptr, grad_shp, grad_c.dev_mask(), grad_c.dev_id());
        } else {
          DType *w1_ptr = workspace.dptr_ + wk_size[0];
          grad_data = TBlob(w1_ptr + offset, grad_shp, grad_c.dev_mask(), grad_c.dev_id());
          offset = AligndWorkspaceSize<DType>(offset, grad_shp.Size());
        }
        const Tuple<int> axis_idx = GetMoveaxisIndex(-1, b_axis, grad_move_shp);
        TransposeImpl<xpu>(ctx.run_ctx, grad_move_data, grad_data,
                           mxnet::TShape(axis_idx.begin(), axis_idx.end()));
      }
      if (!b_reduce_axis.empty()) {
        size_t interval = std::max(grad_a.Size(), grad_b.Size());
        DType *grad_delete_ptr = workspace.dptr_ + wk_size[0] + wk_size[1];
        char *dw_ptr = reinterpret_cast<char*>(grad_delete_ptr + interval);
        TBlob grad_delete_data(grad_delete_ptr, grad_b.shape_, grad_c.dev_mask(), grad_c.dev_id());
        Tensor<xpu, 1, char> dw_tensor(dw_ptr, Shape1((wk_size[2] - interval) * sizeof(DType)), s);
        // Reduce grad_data to grad_delete_data and copy to grad_a.
        ReduceImplWrap<xpu, DType>::op(grad_data, grad_delete_data, grad_b, ctx, req[1], dw_tensor);
      } else {
        MXNET_ASSIGN_REQ_SWITCH(req[0], req_type, {
          mxnet_op::Kernel<ResAssign<req_type>, xpu>::Launch(
            s, grad_b.Size(), grad_data.dptr<DType>(), grad_b.dptr<DType>());
        });
      }
    }
  };
  
  template <typename xpu>
  void NumpyCrossBackward(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;
    CHECK_EQ(inputs.size(), 3U);
    CHECK_EQ(outputs.size(), 2U);
    CHECK_EQ(req.size(), 2U);
    const TBlob& grad_c = inputs[0];
    const TBlob& a = inputs[1];
    const TBlob& b = inputs[2];
    const TBlob& grad_a = outputs[0];
    const TBlob& grad_b = outputs[1];
  
    if (kNullOp == req[0] && kNullOp == req[1]) { return; }
    // Zero-size output, no need to launch kernel
    if (0U == grad_c.Size()) { return; }
  
    const mxnet::TShape& a_shape = a.shape_;
    const mxnet::TShape& b_shape = b.shape_;
    const mxnet::TShape& c_shape = grad_c.shape_;
    const int a_ndim = a_shape.ndim();
    const int b_ndim = b_shape.ndim();
    const int c_ndim = c_shape.ndim();
    const NumpyCrossParam& param = nnvm::get<NumpyCrossParam>(attrs.parsed);
    Tuple<int> a_moveaxis_index = GetMoveaxisIndex(param.axisa, -1, a_shape);
    Tuple<int> b_moveaxis_index = GetMoveaxisIndex(param.axisb, -1, b_shape);
    Tuple<int> c_moveaxis_index = GetMoveaxisIndex(param.axisc, -1, c_shape);
    mxnet::TShape a_moveaxis_shape = GetMoveaxisShape(a_moveaxis_index, a_shape);
    mxnet::TShape b_moveaxis_shape = GetMoveaxisShape(b_moveaxis_index, b_shape);
    mxnet::TShape c_moveaxis_shape = GetMoveaxisShape(c_moveaxis_index, c_shape);
    const int a_axis = CheckAxis(param.axisa, a_ndim);
    const int b_axis = CheckAxis(param.axisb, b_ndim);
    const int c_axis = CheckAxis(param.axisc, c_ndim);
    std::vector<mxnet::TShape> move_shp_vec({ a_moveaxis_shape, b_moveaxis_shape, c_moveaxis_shape });
  
    MSHADOW_SGL_DBL_TYPE_SWITCH(grad_c.type_flag_, DType, {
      bool use_broadcast = CheckUseBroadcast(a_moveaxis_shape, b_moveaxis_shape);
      if (a_moveaxis_shape[a_ndim - 1] == 2) {
        if (b_moveaxis_shape[b_ndim - 1] == 2) {
          // Case 1: a.shape[-1] == 2 and b.shape[-1] == 2, param.axisc is ignored.
          NumpyCrossBackwardImpl<xpu, DType, 2, 2>::op(use_broadcast, grad_c, a, b, grad_a, grad_b,
                                                       move_shp_vec, a_axis, b_axis, c_axis,
                                                       attrs, ctx, req);
        } else {
          // Case 2: a.shape[-1] == 2 and b.shape[-1] == 3, param.axisc is not ignored.
          NumpyCrossBackwardImpl<xpu, DType, 2, 3>::op(use_broadcast, grad_c, a, b, grad_a, grad_b,
                                                       move_shp_vec, a_axis, b_axis, c_axis,
                                                       attrs, ctx, req);
        }
      } else {
        if (b_moveaxis_shape[b_ndim - 1] == 2) {
          // Case 3: a.shape[-1] == 3 and b.shape[-1] == 2, param.axisc is not ignored.
          NumpyCrossBackwardImpl<xpu, DType, 3, 2>::op(use_broadcast, grad_c, a, b, grad_a, grad_b,
                                                       move_shp_vec, a_axis, b_axis, c_axis,
                                                       attrs, ctx, req);
        } else {
          // Case 4: a.shape[-1] == 3 and b.shape[-1] == 3, param.axisc is not ignored.
          NumpyCrossBackwardImpl<xpu, DType, 3, 3>::op(use_broadcast, grad_c, a, b, grad_a, grad_b,
                                                       move_shp_vec, a_axis, b_axis, c_axis,
                                                       attrs, ctx, req);
        }
      }
    });
  }
  
  }  // namespace op
  }  // namespace mxnet
  
  #endif  // MXNET_OPERATOR_NUMPY_NP_CROSS_INL_H_
  