!1567 lstm&transpose_r0.3

Merge pull request !1567 from baihuawei/lstm_r0.3

mindspore/ccsrc/kernel/cpu/mkldnn/lstm_cpu_kernel.cc

@@ -0,0 +1,120 @@
+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed 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.
+ */
+#include "kernel/cpu/mkldnn/lstm_cpu_kernel.h"
+#include <string>
+#include "common/utils.h"
+#include "kernel/cpu/mkldnn/mkl_kernel_engine.h"
+#include "device/cpu/cpu_device_address.h"
+
+namespace mindspore {
+namespace kernel {
+void LstmCPUKernel::InitKernel(const CNodePtr &kernel_node) {
+  MS_EXCEPTION_IF_NULL(kernel_node);
+  std::vector<size_t> src_shape = AnfAlgo::GetInputDeviceShape(kernel_node, 0);
+  bidirectional_ = AnfAlgo::GetNodeAttr<bool>(kernel_node, "bidirectional");
+  input_size_ = AnfAlgo::GetNodeAttr<int>(kernel_node, "input_size");
+  hidden_size_ = AnfAlgo::GetNodeAttr<int>(kernel_node, "hidden_size");
+  num_layers_ = AnfAlgo::GetNodeAttr<int>(kernel_node, "num_layers");
+  batch_size_ = SizeToInt(src_shape[1]);
+  seq_len_ = SizeToInt(src_shape[0]);
+  num_directions_ = 1;
+  if (bidirectional_) {
+    num_directions_ = 2;
+  }
+  int gate_size = 4 * hidden_size_;
+  for (int i = 0; i < num_layers_; ++i) {
+    weight_size_ += gate_size * (i == 0 ? input_size_ : hidden_size_ * num_directions_);
+    weight_h_size_ += gate_size * hidden_size_;
+  }
+  weight_size_ = weight_size_ * num_directions_;
+  weight_h_size_ = weight_h_size_ * num_directions_;
+}
+
+bool LstmCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
+                           const std::vector<kernel::AddressPtr> & /*workspace*/,
+                           const std::vector<kernel::AddressPtr> &outputs) {
+  using dt = dnnl::memory::data_type;
+  using tag = dnnl::memory::format_tag;
+  using dim = dnnl::memory::dims;
+  auto eng = MKLKernelEngine::Get().engine();
+  dnnl::stream s(eng);
+  auto formatted_md = [](dim dimensions, tag layout) { return dnnl::memory::desc{{dimensions}, dt::f32, layout}; };
+  auto generic_md = [](dim dimensions) { return dnnl::memory::desc{{dimensions}, dt::f32, tag::any}; };
+  dnnl::rnn_direction direction = dnnl::rnn_direction::unidirectional;
+  if (bidirectional_) {
+    direction = dnnl::rnn_direction::bidirectional_concat;
+  }
+  dim src_dims = {seq_len_, batch_size_, input_size_};
+  dim src_h_dims = {num_layers_, num_directions_, batch_size_, hidden_size_};
+  dim src_c_dims = {num_layers_, num_directions_, batch_size_, hidden_size_};
+  dim weights_dims = {num_layers_, num_directions_, input_size_, 4, hidden_size_};
+  dim weights_h_dims = {num_layers_, num_directions_, hidden_size_, 4, hidden_size_};
+  dim bias_dims = {num_layers_, num_directions_, 4, hidden_size_};
+  dim dst_dims = {seq_len_, batch_size_, hidden_size_ * num_directions_};
+  dim dst_h_dims = {num_layers_, num_directions_, batch_size_, hidden_size_};
+  dim dst_c_dims = {num_layers_, num_directions_, batch_size_, hidden_size_};
+  dnnl::memory::desc src_desc = formatted_md(src_dims, tag::tnc);
+  dnnl::memory::desc src_h_desc = formatted_md(src_h_dims, tag::ldnc);
+  dnnl::memory::desc src_c_desc = formatted_md(src_c_dims, tag::ldnc);
+  dnnl::memory::desc bias_desc = formatted_md(bias_dims, tag::ldgo);
+  dnnl::memory::desc dst_desc = formatted_md(dst_dims, tag::tnc);
+  dnnl::memory::desc dst_h_desc = formatted_md(dst_h_dims, tag::ldnc);
+  dnnl::memory::desc dst_c_desc = formatted_md(dst_c_dims, tag::ldnc);
+  dnnl::lstm_forward::desc desc = dnnl::lstm_forward::desc(
+    dnnl::prop_kind::forward_training, direction, src_desc, src_h_desc, src_c_desc, generic_md(weights_dims),
+    generic_md(weights_h_dims), bias_desc, dst_desc, dst_h_desc, dst_c_desc);
+  auto prim_desc = dnnl::lstm_forward::primitive_desc(desc, MKLKernelEngine::Get().engine());
+
+  auto workspace_memory = dnnl::memory(prim_desc.workspace_desc(), eng);
+  auto src_memory = dnnl::memory(formatted_md(src_dims, tag::tnc), eng);
+  src_memory.set_data_handle(inputs[0]->addr);
+  auto src_h_memory = dnnl::memory(formatted_md(src_h_dims, tag::ldnc), eng);
+  auto src_c_memory = dnnl::memory(formatted_md(src_c_dims, tag::ldnc), eng);
+  src_h_memory.set_data_handle(inputs[1]->addr);
+  src_c_memory.set_data_handle(inputs[2]->addr);
+  auto user_weights_memory = dnnl::memory(formatted_md(weights_dims, tag::ldgoi), eng);
+  auto user_weights_h_memory = dnnl::memory(formatted_md(weights_h_dims, tag::ldgoi), eng);
+  auto bias_memory = dnnl::memory(formatted_md(bias_dims, tag::ldgo), eng);
+  user_weights_memory.set_data_handle(inputs[3]->addr);
+  user_weights_h_memory.set_data_handle(reinterpret_cast<float *>(inputs[3]->addr) + weight_size_);
+  bias_memory.set_data_handle(reinterpret_cast<float *>(inputs[3]->addr) + weight_size_ + weight_h_size_);
+  auto weights_memory = dnnl::memory(prim_desc.weights_layer_desc(), eng);
+  auto weights_h_memory = dnnl::memory(prim_desc.weights_iter_desc(), eng);
+  dnnl::reorder(user_weights_memory, weights_memory).execute(s, user_weights_memory, weights_memory);
+  dnnl::reorder(user_weights_h_memory, weights_h_memory).execute(s, user_weights_h_memory, weights_h_memory);
+  auto dst_memory = dnnl::memory(formatted_md(dst_dims, tag::tnc), eng);
+  auto dst_h_memory = dnnl::memory(formatted_md(dst_h_dims, tag::ldnc), eng);
+  auto dst_c_memory = dnnl::memory(formatted_md(dst_c_dims, tag::ldnc), eng);
+  dnnl::lstm_forward fw_layer(prim_desc);
+  workspace_memory.set_data_handle(outputs[3]->addr);
+  dst_memory.set_data_handle(outputs[0]->addr);
+  dst_h_memory.set_data_handle(outputs[1]->addr);
+  dst_c_memory.set_data_handle(outputs[2]->addr);
+  fw_layer.execute(s, {{DNNL_ARG_SRC_LAYER, src_memory},
+                       {DNNL_ARG_SRC_ITER, src_h_memory},
+                       {DNNL_ARG_SRC_ITER_C, src_c_memory},
+                       {DNNL_ARG_WEIGHTS_LAYER, weights_memory},
+                       {DNNL_ARG_WEIGHTS_ITER, weights_h_memory},
+                       {DNNL_ARG_BIAS, bias_memory},
+                       {DNNL_ARG_DST_LAYER, dst_memory},
+                       {DNNL_ARG_DST_ITER, dst_h_memory},
+                       {DNNL_ARG_DST_ITER_C, dst_c_memory},
+                       {DNNL_ARG_WORKSPACE, workspace_memory}});
+  s.wait();
+  return true;
+}
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/kernel/cpu/mkldnn/lstm_cpu_kernel.h

@@ -0,0 +1,59 @@
+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed 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.
+ */
+
+#ifndef MINDSPORE_CCSRC_KERNEL_CPU_LSTM_CPU_KERNEL_H
+#define MINDSPORE_CCSRC_KERNEL_CPU_LSTM_CPU_KERNEL_H
+#include <vector>
+#include <memory>
+#include "kernel/cpu/mkldnn/mkl_cpu_kernel.h"
+namespace mindspore {
+namespace kernel {
+class LstmCPUKernel : public MKLCPUKernel {
+ public:
+  LstmCPUKernel() = default;
+  ~LstmCPUKernel() override = default;
+  void InitKernel(const CNodePtr &kernel_node) override;
+
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+              const std::vector<AddressPtr> &outputs) override;
+
+ private:
+  int weight_size_ = 0;
+  int weight_h_size_ = 0;
+  int input_size_;
+  int hidden_size_;
+  int num_layers_;
+  int batch_size_;
+  int seq_len_;
+  int num_directions_;
+  bool bidirectional_;
+};
+
+MS_REG_CPU_KERNEL(LSTM,
+                  KernelAttr()
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32),
+                  LstmCPUKernel);
+}  // namespace kernel
+}  // namespace mindspore
+#endif  // MINDSPORE_CCSRC_KERNEL_CPU_LSTM_CPU_KERNEL_H

mindspore/ccsrc/kernel/cpu/mkldnn/lstm_grad_cpu_kernel.cc

@@ -0,0 +1,174 @@
+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed 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.
+ */
+#include "kernel/cpu/mkldnn/lstm_grad_cpu_kernel.h"
+#include <cstring>
+#include <cmath>
+#include <numeric>
+#include <string>
+#include "common/utils.h"
+#include "kernel/cpu/mkldnn/mkl_kernel_engine.h"
+#include "device/cpu/cpu_device_address.h"
+
+namespace mindspore {
+namespace kernel {
+
+void LSTMGradCPUKernel::InitKernel(const CNodePtr &kernel_node) {
+  MS_EXCEPTION_IF_NULL(kernel_node);
+  std::vector<size_t> src_shape = AnfAlgo::GetInputDeviceShape(kernel_node, 0);
+  bidirectional_ = AnfAlgo::GetNodeAttr<bool>(kernel_node, "bidirectional");
+  input_size_ = AnfAlgo::GetNodeAttr<int>(kernel_node, "input_size");
+  hidden_size_ = AnfAlgo::GetNodeAttr<int>(kernel_node, "hidden_size");
+  num_layers_ = AnfAlgo::GetNodeAttr<int>(kernel_node, "num_layers");
+  batch_size_ = SizeToInt(src_shape[1]);
+  seq_len_ = SizeToInt(src_shape[0]);
+  num_directions_ = 1;
+  if (bidirectional_) {
+    num_directions_ = 2;
+  }
+  int gate_size = 4 * hidden_size_;
+  for (int i = 0; i < num_layers_; ++i) {
+    weight_size_ += gate_size * (i == 0 ? input_size_ : hidden_size_ * num_directions_);
+    weight_h_size_ += gate_size * hidden_size_;
+  }
+  weight_size_ = weight_size_ * num_directions_;
+  weight_h_size_ = weight_h_size_ * num_directions_;
+}
+
+bool LSTMGradCPUKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
+                               const std::vector<kernel::AddressPtr> &workspace /*workspace*/,
+                               const std::vector<kernel::AddressPtr> &outputs) {
+  using tag = dnnl::memory::format_tag;
+  using dt = dnnl::memory::data_type;
+  using dim = dnnl::memory::dims;
+  auto eng = MKLKernelEngine::Get().engine();
+  dnnl::stream s(eng);
+  auto formatted_md = [](dim dimensions, tag layout) { return dnnl::memory::desc{{dimensions}, dt::f32, layout}; };
+  auto generic_md = [](dim dimensions) { return dnnl::memory::desc{{dimensions}, dt::f32, tag::any}; };
+  dnnl::rnn_direction direction = dnnl::rnn_direction::unidirectional;
+  if (bidirectional_) {
+    direction = dnnl::rnn_direction::bidirectional_concat;
+  }
+  dim src_dims = {seq_len_, batch_size_, input_size_};
+  dim src_h_dims = {num_layers_, num_directions_, batch_size_, hidden_size_};
+  dim src_c_dims = {num_layers_, num_directions_, batch_size_, hidden_size_};
+  dim weights_dims = {num_layers_, num_directions_, input_size_, 4, hidden_size_};
+  dim weights_h_dims = {num_layers_, num_directions_, hidden_size_, 4, hidden_size_};
+  dim bias_dims = {num_layers_, num_directions_, 4, hidden_size_};
+  dim dst_dims = {seq_len_, batch_size_, hidden_size_ * num_directions_};
+  dim dst_h_dims = {num_layers_, num_directions_, batch_size_, hidden_size_};
+  dim dst_c_dims = {num_layers_, num_directions_, batch_size_, hidden_size_};
+  dnnl::memory::desc src_desc = formatted_md(src_dims, tag::tnc);
+  dnnl::memory::desc src_h_desc = formatted_md(src_h_dims, tag::ldnc);
+  dnnl::memory::desc src_c_desc = formatted_md(src_c_dims, tag::ldnc);
+  dnnl::memory::desc bias_desc = formatted_md(bias_dims, tag::ldgo);
+  dnnl::memory::desc dst_desc = formatted_md(dst_dims, tag::tnc);
+  dnnl::memory::desc dst_h_desc = formatted_md(dst_h_dims, tag::ldnc);
+  dnnl::memory::desc dst_c_desc = formatted_md(dst_c_dims, tag::ldnc);
+  dnnl::lstm_forward::desc forward_desc = dnnl::lstm_forward::desc(
+    dnnl::prop_kind::forward_training, direction, src_desc, src_h_desc, src_c_desc, generic_md(weights_dims),
+    generic_md(weights_h_dims), bias_desc, dst_desc, dst_h_desc, dst_c_desc);
+  auto prim_forward_desc = dnnl::lstm_forward::primitive_desc(forward_desc, eng);
+  dnnl::lstm_backward::desc backward_desc = dnnl::lstm_backward::desc(
+    dnnl::prop_kind::backward, direction, src_desc, src_h_desc, src_c_desc, generic_md(weights_dims),
+    generic_md(weights_h_dims), bias_desc, dst_desc, dst_h_desc, dst_c_desc, src_desc, src_h_desc, src_c_desc,
+    generic_md(weights_dims), generic_md(weights_h_dims), bias_desc, dst_desc, dst_h_desc, dst_c_desc);
+  auto prim_backward_desc = dnnl::lstm_backward::primitive_desc(backward_desc, eng, prim_forward_desc);
+
+  // construct fw memory
+  auto src_memory = dnnl::memory(formatted_md(src_dims, tag::tnc), eng);
+  src_memory.set_data_handle(inputs[0]->addr);
+  auto src_h_memory = dnnl::memory(formatted_md(src_h_dims, tag::ldnc), eng);
+  auto src_c_memory = dnnl::memory(formatted_md(src_c_dims, tag::ldnc), eng);
+  src_h_memory.set_data_handle(inputs[1]->addr);
+  src_c_memory.set_data_handle(inputs[2]->addr);
+  auto user_weights_memory = dnnl::memory(formatted_md(weights_dims, tag::ldgoi), eng);
+  auto user_weights_h_memory = dnnl::memory(formatted_md(weights_h_dims, tag::ldgoi), eng);
+  auto bias_memory = dnnl::memory(formatted_md(bias_dims, tag::ldgo), eng);
+  user_weights_memory.set_data_handle(inputs[3]->addr);
+  user_weights_h_memory.set_data_handle(reinterpret_cast<float *>(inputs[3]->addr) + weight_size_);
+  bias_memory.set_data_handle(reinterpret_cast<float *>(inputs[3]->addr) + weight_size_ + weight_h_size_);
+  auto weights_memory = dnnl::memory(prim_backward_desc.weights_layer_desc(), eng);
+  auto weights_h_memory = dnnl::memory(prim_backward_desc.weights_iter_desc(), eng);
+  std::vector<float> net_x(src_memory.get_desc().get_size(), 0.0f);
+  std::vector<float> net_w(weights_memory.get_desc().get_size(), 0.0f);
+  std::vector<float> net_wh(weights_h_memory.get_desc().get_size(), 0.0f);
+  dnnl::reorder(user_weights_memory, weights_memory).execute(s, user_weights_memory, weights_memory);
+  dnnl::reorder(user_weights_h_memory, weights_h_memory).execute(s, user_weights_h_memory, weights_h_memory);
+  auto dst_memory = dnnl::memory(formatted_md(dst_dims, tag::tnc), eng);
+  dst_memory.set_data_handle(inputs[4]->addr);
+  auto dst_h_memory = dnnl::memory(formatted_md(dst_h_dims, tag::ldnc), eng);
+  auto dst_c_memory = dnnl::memory(formatted_md(dst_c_dims, tag::ldnc), eng);
+  dst_h_memory.set_data_handle(inputs[5]->addr);
+  dst_c_memory.set_data_handle(inputs[6]->addr);
+  auto workspace_memory = dnnl::memory(prim_forward_desc.workspace_desc(), eng);
+  workspace_memory.set_data_handle(inputs[10]->addr);
+  auto diff_src_memory = dnnl::memory(formatted_md(src_dims, tag::tnc), eng);
+  auto diff_src_h_memory = dnnl::memory(formatted_md(src_h_dims, tag::ldnc), eng);
+  auto diff_src_c_memory = dnnl::memory(formatted_md(src_c_dims, tag::ldnc), eng);
+  auto user_diff_weights_memory = dnnl::memory(formatted_md(weights_dims, tag::ldgoi), eng);
+  auto user_diff_weights_h_memory = dnnl::memory(formatted_md(weights_h_dims, tag::ldgoi), eng);
+  auto diff_weights_memory = dnnl::memory(prim_backward_desc.diff_weights_layer_desc(), eng);
+  auto diff_weights_h_memory = dnnl::memory(prim_backward_desc.diff_weights_iter_desc(), eng);
+  write_to_dnnl_memory(net_w.data(), diff_weights_memory);
+  write_to_dnnl_memory(net_wh.data(), diff_weights_h_memory);
+  auto diff_bias_memory = dnnl::memory(formatted_md(bias_dims, tag::ldgo), eng);
+  auto diff_dst_memory = dnnl::memory(formatted_md(dst_dims, tag::tnc), eng);
+  diff_dst_memory.set_data_handle(inputs[7]->addr);
+  auto diff_dst_h_memory = dnnl::memory(formatted_md(dst_h_dims, tag::ldnc), eng);
+  diff_dst_h_memory.set_data_handle(inputs[8]->addr);
+  auto diff_dst_c_memory = dnnl::memory(formatted_md(dst_c_dims, tag::ldnc), eng);
+  diff_dst_c_memory.set_data_handle(inputs[9]->addr);
+  diff_src_memory.set_data_handle(outputs[0]->addr);
+  diff_src_h_memory.set_data_handle(outputs[1]->addr);
+  diff_src_c_memory.set_data_handle(outputs[2]->addr);
+  user_diff_weights_memory.set_data_handle(outputs[3]->addr);
+  user_diff_weights_h_memory.set_data_handle(reinterpret_cast<float *>(outputs[3]->addr) + weight_size_);
+  diff_bias_memory.set_data_handle(reinterpret_cast<float *>(outputs[3]->addr) + weight_size_ + weight_h_size_);
+  write_to_dnnl_memory(net_w.data(), user_diff_weights_memory);
+  write_to_dnnl_memory(net_wh.data(), user_diff_weights_h_memory);
+  write_to_dnnl_memory(net_w.data(), diff_bias_memory);
+  write_to_dnnl_memory(net_x.data(), diff_src_memory);
+  write_to_dnnl_memory(net_x.data(), diff_src_h_memory);
+  write_to_dnnl_memory(net_x.data(), diff_src_c_memory);
+  dnnl::lstm_backward bwd_layer(prim_backward_desc);
+  bwd_layer.execute(s, {{DNNL_ARG_SRC_LAYER, src_memory},
+                        {DNNL_ARG_SRC_ITER, src_h_memory},
+                        {DNNL_ARG_SRC_ITER_C, src_c_memory},
+                        {DNNL_ARG_WEIGHTS_LAYER, weights_memory},
+                        {DNNL_ARG_WEIGHTS_ITER, weights_h_memory},
+                        {DNNL_ARG_BIAS, bias_memory},
+                        {DNNL_ARG_DST_LAYER, dst_memory},
+                        {DNNL_ARG_DST_ITER, dst_h_memory},
+                        {DNNL_ARG_DST_ITER_C, dst_c_memory},
+                        {DNNL_ARG_DIFF_SRC_LAYER, diff_src_memory},
+                        {DNNL_ARG_DIFF_SRC_ITER, diff_src_h_memory},
+                        {DNNL_ARG_DIFF_SRC_ITER_C, diff_src_c_memory},
+                        {DNNL_ARG_DIFF_WEIGHTS_LAYER, diff_weights_memory},
+                        {DNNL_ARG_DIFF_WEIGHTS_ITER, diff_weights_h_memory},
+                        {DNNL_ARG_DIFF_BIAS, diff_bias_memory},
+                        {DNNL_ARG_DIFF_DST_LAYER, diff_dst_memory},
+                        {DNNL_ARG_DIFF_DST_ITER, diff_dst_h_memory},
+                        {DNNL_ARG_DIFF_DST_ITER_C, diff_dst_c_memory},
+                        {DNNL_ARG_WORKSPACE, workspace_memory}});
+  dnnl::reorder(diff_weights_memory, user_diff_weights_memory)
+    .execute(s, diff_weights_memory, user_diff_weights_memory);
+  dnnl::reorder(diff_weights_h_memory, user_diff_weights_h_memory)
+    .execute(s, diff_weights_h_memory, user_diff_weights_h_memory);
+  s.wait();
+  return true;
+}
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/kernel/cpu/mkldnn/lstm_grad_cpu_kernel.h

@@ -0,0 +1,67 @@
+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed 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.
+ */
+#ifndef MINDSPORE_CCSRC_KERNEL_CPU_LSTM_GRAD_CPU_KERNEL_H_
+#define MINDSPORE_CCSRC_KERNEL_CPU_LSTM_GRAD_CPU_KERNEL_H_
+
+#include <vector>
+#include <memory>
+#include "kernel/cpu/mkldnn/mkl_cpu_kernel.h"
+
+namespace mindspore {
+namespace kernel {
+class LSTMGradCPUKernel : public MKLCPUKernel {
+ public:
+  LSTMGradCPUKernel() = default;
+  ~LSTMGradCPUKernel() override = default;
+
+  void InitKernel(const CNodePtr &kernel_node) override;
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+              const std::vector<AddressPtr> &outputs) override;
+
+ private:
+  int weight_size_ = 0;
+  int weight_h_size_ = 0;
+  int input_size_;
+  int hidden_size_;
+  int num_layers_;
+  int batch_size_;
+  int seq_len_;
+  int num_directions_;
+  bool bidirectional_;
+};
+
+MS_REG_CPU_KERNEL(LSTMGrad,
+                  KernelAttr()
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddInputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32)
+                    .AddOutputAttr(kNumberTypeFloat32),
+                  LSTMGradCPUKernel);
+}  // namespace kernel
+}  // namespace mindspore
+
+#endif  // MINDSPORE_CCSRC_KERNEL_CPU_LSTM_GRAD_CPU_KERNEL_H_

mindspore/ccsrc/kernel/cpu/mkldnn/mkl_cpu_kernel.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2019 Huawei Technologies Co., Ltd
+ * Copyright 2020 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -98,5 +98,11 @@ void MKLCPUKernel::SetArgumentHandle(int arg_key, void *ptr) {
 }
 
 void MKLCPUKernel::ExecutePrimitive() { MKLKernelEngine::Get().Execute(primitive_, arguments_); }
+void MKLCPUKernel::write_to_dnnl_memory(void *handle, const dnnl::memory &mem) {
+  MKLKernelEngine::Get().write_to_dnnl_memory(handle, mem);
+}
+void MKLCPUKernel::read_from_dnnl_memory(void *handle, const dnnl::memory &mem) {
+  MKLKernelEngine::Get().read_from_dnnl_memory(handle, mem);
+}
 }  // namespace kernel
 }  // namespace mindspore

mindspore/ccsrc/kernel/cpu/mkldnn/mkl_cpu_kernel.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2019 Huawei Technologies Co., Ltd
+ * Copyright 2020 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -39,6 +39,8 @@ class MKLCPUKernel : public CPUKernel {
   dnnl::memory::format_tag GetDefaultFormatTag(const dnnl::memory::dims &dims) const;
   dnnl::memory::desc GetDefaultMemDesc(const std::vector<size_t> &shape);
   void ExecutePrimitive();
+  void write_to_dnnl_memory(void *handle, const dnnl::memory &mem);
+  void read_from_dnnl_memory(void *handle, const dnnl::memory &mem);
   std::unordered_map<int, dnnl::memory> arguments_;
   std::shared_ptr<dnnl::primitive> primitive_{nullptr};
 };

mindspore/ccsrc/kernel/cpu/mkldnn/mkl_kernel_engine.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2019 Huawei Technologies Co., Ltd
+ * Copyright 2020 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.

mindspore/ccsrc/kernel/cpu/mkldnn/mkl_kernel_engine.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2019 Huawei Technologies Co., Ltd
+ * Copyright 2020 Huawei Technologies Co., Ltd
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -15,7 +15,10 @@
  */
 #ifndef MINDSPORE_MKL_KERNEL_ENGINE_H_
 #define MINDSPORE_MKL_KERNEL_ENGINE_H_
-
+#include <cstdlib>
+#include <algorithm>
+#include <iostream>
+#include <string>
 #include <unordered_map>
 #include <vector>
 #include <memory>
@@ -39,6 +42,30 @@ class MKLKernelEngine {
   void Execute(const std::shared_ptr<dnnl::primitive> &primitive,
                const std::unordered_map<int, dnnl::memory> &arguments);
 
+  inline void read_from_dnnl_memory(void *handle, const dnnl::memory &mem) {
+    dnnl::engine eng = mem.get_engine();
+    size_t bytes = mem.get_desc().get_size();
+    if (eng.get_kind() == dnnl::engine::kind::cpu) {
+      auto dst = reinterpret_cast<uint8_t *>(handle);
+      uint8_t *src = reinterpret_cast<uint8_t *>(mem.get_data_handle());
+      for (size_t i = 0; i < bytes; ++i) {
+        dst[i] = src[i];
+      }
+    }
+  }
+  // Read from handle, write to memory
+  inline void write_to_dnnl_memory(void *handle, const dnnl::memory &mem) {
+    dnnl::engine eng = mem.get_engine();
+    size_t bytes = mem.get_desc().get_size();
+    if (eng.get_kind() == dnnl::engine::kind::cpu) {
+      auto src = reinterpret_cast<uint8_t *>(handle);
+      uint8_t *dst = reinterpret_cast<uint8_t *>(mem.get_data_handle());
+      for (size_t i = 0; i < bytes; ++i) {
+        dst[i] = src[i];
+      }
+    }
+  }
+
  private:
   MKLKernelEngine() : engine_(dnnl::engine::kind::cpu, 0), stream_(engine_) {}
   ~MKLKernelEngine() = default;

mindspore/ccsrc/kernel/cpu/transpose_cpu_kernel.cc

@@ -0,0 +1,64 @@
+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed 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.
+ */
+
+#include "kernel/cpu/transpose_cpu_kernel.h"
+#include "device/cpu/cpu_device_address.h"
+namespace mindspore {
+namespace kernel {
+const size_t kMaxDim = 100;
+void TransposeCPUFwdKernel::InitKernel(const CNodePtr &kernel_node) {
+  MS_EXCEPTION_IF_NULL(kernel_node);
+  shape_ = AnfAlgo::GetInputDeviceShape(kernel_node, 0);
+  axis_ = AnfAlgo::GetNodeAttr<std::vector<int>>(kernel_node, "perm");
+  if (shape_.size() != axis_.size()) {
+    MS_LOG(EXCEPTION) << "The size of input shape and transpose axis shape must be equal.";
+  }
+}
+bool TransposeCPUFwdKernel::Launch(const std::vector<kernel::AddressPtr> &inputs,
+                                   const std::vector<kernel::AddressPtr> & /*workspace*/,
+                                   const std::vector<kernel::AddressPtr> &outputs) {
+  auto input = reinterpret_cast<float *>(inputs[0]->addr);
+  auto output = reinterpret_cast<float *>(outputs[0]->addr);
+  size_t size = IntToSize(inputs[0]->size / sizeof(float));
+  size_t shape_size = IntToSize(shape_.size());
+  if (shape_size > kMaxDim) {
+    MS_LOG(EXCEPTION) << "Input is " << shape_size << "-D, but transpose supports max " << kMaxDim << "-D inputs.";
+  }
+  size_t pos_array[kMaxDim];
+  size_t size_offset[kMaxDim];
+  size_offset[0] = size / shape_[0];
+  for (size_t i = 1; i < shape_size; i++) {
+    size_offset[i] = size_offset[SizeToInt(i) - 1] / shape_[i];
+  }
+  for (size_t position = 0; position < size; position += 1) {
+    size_t temp_position = position;
+    pos_array[0] = temp_position / size_offset[0];
+    for (size_t i = 1; i < shape_size; i++) {
+      temp_position -= pos_array[SizeToInt(i) - 1] * size_offset[i - 1];
+      pos_array[i] = temp_position / size_offset[i];
+    }
+    size_t new_position = pos_array[axis_[SizeToInt(shape_size) - 1]];
+    size_t new_position_size = 1;
+    for (int j = shape_size - 2; j >= 0; j--) {
+      new_position_size *= shape_[axis_[j + 1]];
+      new_position += pos_array[axis_[j]] * new_position_size;
+    }
+    output[new_position] = input[position];
+  }
+  return true;
+}
+}  // namespace kernel
+}  // namespace mindspore

mindspore/ccsrc/kernel/cpu/transpose_cpu_kernel.h

@@ -0,0 +1,44 @@
+/**
+ * Copyright 2020 Huawei Technologies Co., Ltd
+ *
+ * Licensed 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.
+ */
+#ifndef MINDSPORE_CCSRC_KERNEL_CPU_TRANSPOSE_CPU_KERNEL_H_
+#define MINDSPORE_CCSRC_KERNEL_CPU_TRANSPOSE_CPU_KERNEL_H_
+#include <vector>
+#include <memory>
+#include <string>
+#include "kernel/cpu/cpu_kernel.h"
+#include "kernel/cpu/cpu_kernel_factory.h"
+namespace mindspore {
+namespace kernel {
+class TransposeCPUFwdKernel : public CPUKernel {
+ public:
+  TransposeCPUFwdKernel() = default;
+  ~TransposeCPUFwdKernel() override = default;
+
+  void InitKernel(const CNodePtr &kernel_node) override;
+
+  bool Launch(const std::vector<AddressPtr> &inputs, const std::vector<AddressPtr> &workspace,
+              const std::vector<AddressPtr> &outputs) override;
+
+ private:
+  std::vector<size_t> shape_;
+  std::vector<int> axis_;
+};
+
+MS_REG_CPU_KERNEL(Transpose, KernelAttr().AddInputAttr(kNumberTypeFloat32).AddOutputAttr(kNumberTypeFloat32),
+                  TransposeCPUFwdKernel);
+}  // namespace kernel
+}  // namespace mindspore
+#endif  // MINDSPORE_CCSRC_KERNEL_CPU_TRANSPOSE_CPU_KERNEL_H_

mindspore/nn/layer/lstm.py

@@ -18,8 +18,13 @@ from mindspore.nn.cell import Cell
 from mindspore.common.parameter import Parameter
 from mindspore.common.initializer import initializer
 from mindspore._checkparam import Validator as validator
+from mindspore import context
+import mindspore.nn as nn
+from mindspore.common.tensor import Tensor
+import numpy as np
+
+__all__ = ['LSTM', 'LSTMCell']
 
-__all__ = ['LSTM']
 
 class LSTM(Cell):
     r"""
@@ -102,6 +107,7 @@ class LSTM(Cell):
         >>> c0 = Tensor(np.ones([1 * 2, 3, 12]).astype(np.float32))
         >>> output, (hn, cn) = net(input, h0, c0)
     """
+
     def __init__(self,
                  input_size,
                  hidden_size,
@@ -118,39 +124,197 @@ class LSTM(Cell):
         self.batch_first = validator.check_value_type("batch_first", batch_first, [bool], self.cls_name)
         self.dropout = float(dropout)
         self.bidirectional = bidirectional
+        if self.batch_first:
+            self.transpose1 = P.Transpose()
+            self.transpose2 = P.Transpose()
+        num_directions = 2 if self.bidirectional else 1
+        self.cpu_target = False
+        if context.get_context("device_target") == "CPU":
+            self.cpu_target = True
+        if not self.cpu_target:
+            self.lstm = P.LSTM(input_size=self.input_size,
+                               hidden_size=self.hidden_size,
+                               num_layers=self.num_layers,
+                               has_bias=self.has_bias,
+                               bidirectional=self.bidirectional,
+                               dropout=self.dropout)
+            weight_size = 0
+            gate_size = 4 * self.hidden_size
+            for layer in range(self.num_layers):
+                input_layer_size = self.input_size if layer == 0 else self.hidden_size * num_directions
+                increment_size = gate_size * input_layer_size
+                increment_size += gate_size * self.hidden_size
+                if self.has_bias:
+                    increment_size += 2 * gate_size
+                weight_size += increment_size * num_directions
+            self.weight = Parameter(initializer(0.0, [weight_size, 1, 1]), name='weight')
+        else:
+            layer = []
+            layer.append(nn.LSTMCell(input_size=self.input_size,
+                                     hidden_size=self.hidden_size,
+                                     layer_index=0,
+                                     has_bias=self.has_bias,
+                                     bidirectional=self.bidirectional,
+                                     dropout=self.dropout))
+            for i in range(num_layers - 1):
+                layer.append(nn.LSTMCell(input_size=self.hidden_size * num_directions,
+                                         hidden_size=self.hidden_size,
+                                         layer_index=i + 1,
+                                         has_bias=self.has_bias,
+                                         bidirectional=self.bidirectional,
+                                         dropout=self.dropout))
+            self.lstms = layer
+        self.fill = P.Fill()
+        self.shape = P.Shape()
 
+    def construct(self, x, hx):
+        if self.batch_first:
+            x = self.transpose1(x, (1, 0, 2))
+        if not self.cpu_target:
+            h, c = hx
+            output, h, c, _, _ = self.lstm(x, h, c, self.weight)
+            if self.batch_first:
+                output = self.transpose2(output, (1, 0, 2))
+            return (output, (h, c))
+        h, c = hx
+        output, hn, cn, _, _ = self.lstms[0](x, h[0], c[0])
+        for i in range(1, self.num_layers):
+            output, hn, cn, _, _ = self.lstms[i](output, h[i], c[i])
+        if self.batch_first:
+            output = self.transpose2(output, (1, 0, 2))
+        return output, hn, cn, _, _
+
+
+class LSTMCell(Cell):
+    r"""
+    LSTM (Long Short-Term Memory) layer.
+
+    Applies a LSTM layer to the input.
+
+    There are two pipelines connecting two consecutive cells in a LSTM model; one is cell state pipeline
+    and another is hidden state pipeline. Denote two consecutive time nodes as :math:`t-1` and :math:`t`.
+    Given an input :math:`x_t` at time :math:`t`, an hidden state :math:`h_{t-1}` and an cell
+    state :math:`c_{t-1}` of the layer at time :math:`{t-1}`, the cell state and hidden state at
+    time :math:`t` is computed using an gating mechanism. Input gate :math:`i_t` is designed to protect the cell
+    from perturbation by irrelevant inputs. Forget gate :math:`f_t` affords protection of the cell by forgetting
+    some information in the past, which is stored in :math:`h_{t-1}`. Output gate :math:`o_t` protects other
+    units from perturbation by currently irrelevant memory contents. Candidate cell state :math:`\tilde{c}_t` is
+    calculated with the current input, on which the input gate will be applied. Finally, current cell state
+    :math:`c_{t}` and hidden state :math:`h_{t}` are computed with the calculated gates and cell states. The complete
+    formulation is as follows.
+
+    .. math::
+        \begin{array}{ll} \\
+            i_t = \sigma(W_{ix} x_t + b_{ix} + W_{ih} h_{(t-1)} + b_{ih}) \\
+            f_t = \sigma(W_{fx} x_t + b_{fx} + W_{fh} h_{(t-1)} + b_{fh}) \\
+            \tilde{c}_t = \tanh(W_{cx} x_t + b_{cx} + W_{ch} h_{(t-1)} + b_{ch}) \\
+            o_t = \sigma(W_{ox} x_t + b_{ox} + W_{oh} h_{(t-1)} + b_{oh}) \\
+            c_t = f_t * c_{(t-1)} + i_t * \tilde{c}_t \\
+            h_t = o_t * \tanh(c_t) \\
+        \end{array}
+
+    Here :math:`\sigma` is the sigmoid function, and :math:`*` is the Hadamard product. :math:`W, b`
+    are learnable weights between the output and the input in the formula. For instance,
+    :math:`W_{ix}, b_{ix}` are the weight and bias used to transform from input :math:`x` to :math:`i`.
+    Details can be found in paper `LONG SHORT-TERM MEMORY
+    <https://www.bioinf.jku.at/publications/older/2604.pdf>`_ and
+    `Long Short-Term Memory Recurrent Neural Network Architectures for Large Scale Acoustic Modeling
+    <https://static.googleusercontent.com/media/research.google.com/zh-CN//pubs/archive/43905.pdf>`_.
+
+    Args:
+        input_size (int): Number of features of input.
+        hidden_size (int):  Number of features of hidden layer.
+        layer_index (int): index of current layer of stacked LSTM . Default: 0.
+        has_bias (bool): Specifies whether has bias `b_ih` and `b_hh`. Default: True.
+        batch_first (bool): Specifies whether the first dimension of input is batch_size. Default: False.
+        dropout (float, int): If not 0, append `Dropout` layer on the outputs of each
+            LSTM layer except the last layer. Default 0. The range of dropout is [0.0, 1.0].
+        bidirectional (bool): Specifies whether this is a bidirectional LSTM. If set True,
+            number of directions will be 2 otherwise number of directions is 1. Default: False.
+
+    Inputs:
+        - **input** (Tensor) - Tensor of shape (seq_len, batch_size, `input_size`).
+        - **h** - data type mindspore.float32 or
+          mindspore.float16 and shape (num_directions * `num_layers`, batch_size, `hidden_size`).
+        - **c** - data type mindspore.float32 or
+          mindspore.float16 and shape (num_directions * `num_layers`, batch_size, `hidden_size`).
+          Data type of `h' and 'c' should be the same of `input`.
+
+    Outputs:
+        `output`, `h_n`, `c_n`, 'reserve', 'state'.
+
+        - **output** (Tensor) - Tensor of shape (seq_len, batch_size, num_directions * `hidden_size`).
+        - **h** - A Tensor with shape (num_directions * `num_layers`, batch_size, `hidden_size`).
+        - **c** - A Tensor with shape (num_directions * `num_layers`, batch_size, `hidden_size`).
+        - **reserve** - reserved
+        - **state** - reserved
+
+    Examples:
+        >>> class LstmNet(nn.Cell):
+        >>>     def __init__(self, input_size, hidden_size, layer_index, has_bias, batch_first, bidirectional):
+        >>>         super(LstmNet, self).__init__()
+        >>>         self.lstm = nn.LSTMCell(input_size=input_size,
+        >>>                             hidden_size=hidden_size,
+        >>>                             layer_index=layer_index,
+        >>>                             has_bias=has_bias,
+        >>>                             batch_first=batch_first,
+        >>>                             bidirectional=bidirectional,
+        >>>                             dropout=0.0)
+        >>>
+        >>>     def construct(self, inp, h0, c0):
+        >>>         return self.lstm(inp, (h0, c0))
+        >>>
+        >>> net = LstmNet(10, 12, 2, has_bias=True, batch_first=True, bidirectional=False)
+        >>> input = Tensor(np.ones([3, 5, 10]).astype(np.float32))
+        >>> h0 = Tensor(np.ones([1 * 2, 3, 12]).astype(np.float32))
+        >>> c0 = Tensor(np.ones([1 * 2, 3, 12]).astype(np.float32))
+        >>> output, hn, cn, _, _ = net(input, h0, c0)
+    """
+
+    def __init__(self,
+                 input_size,
+                 hidden_size,
+                 layer_index=0,
+                 has_bias=True,
+                 batch_first=False,
+                 dropout=0,
+                 bidirectional=False):
+        super(LSTMCell, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = 1
+        self.layer_index = layer_index
+        self.has_bias = has_bias
+        self.batch_first = validator.check_value_type("batch_first", batch_first, [bool], self.cls_name)
+        self.dropout = float(dropout)
+        self.bidirectional = bidirectional
+        self.num_directions = 1
+        if self.bidirectional:
+            self.num_directions = 2
         if self.batch_first:
             self.transpose1 = P.Transpose()
             self.transpose2 = P.Transpose()
+        w_np = np.ones([(self.input_size + self.hidden_size) * self.num_directions * self.hidden_size * 4, 1]).astype(
+            np.float32) * 0.01
+        if has_bias:
+            b_np = np.ones([self.num_directions * self.hidden_size * 4, 1]).astype(
+                np.float32) * 0.01
+        else:
+            b_np = np.zeros([self.num_directions * self.hidden_size * 4, 1]).astype(
+                np.float32) * 0.01
+        wb_np = np.concatenate((w_np, b_np), axis=0).reshape([-1, 1, 1])
+        self.w = Parameter(initializer(Tensor(wb_np), wb_np.shape), name='w' + str(self.layer_index))
         self.lstm = P.LSTM(input_size=self.input_size,
                            hidden_size=self.hidden_size,
-                           num_layers=self.num_layers,
+                           num_layers=1,
                            has_bias=self.has_bias,
                            bidirectional=self.bidirectional,
                            dropout=self.dropout)
 
-        num_directions = 2 if self.bidirectional else 1
-
-        weight_size = 0
-        gate_size = 4 * self.hidden_size
-        for layer in range(self.num_layers):
-            input_layer_size = self.input_size if layer == 0 else self.hidden_size * num_directions
-            increment_size = gate_size * input_layer_size
-            increment_size += gate_size * self.hidden_size
-            if self.has_bias:
-                increment_size += 2 * gate_size
-            weight_size += increment_size * num_directions
-
-        self.weight = Parameter(initializer(0.0, [weight_size, 1, 1]), name='weight')
-
-        self.fill = P.Fill()
-        self.shape = P.Shape()
-
-    def construct(self, x, hx):
+    def construct(self, x, h, c):
         if self.batch_first:
             x = self.transpose1(x, (1, 0, 2))
-        h0, c0 = hx
-        output, hn, cn, _, _ = self.lstm(x, h0, c0, self.weight)
+        output, hn, cn, _, _ = self.lstm(x, h, c, self.w)
         if self.batch_first:
             output = self.transpose2(output, (1, 0, 2))
-        return (output, (hn, cn))
+        return output, hn, cn, _, _

mindspore/ops/_grad/grad_array_ops.py

@@ -49,6 +49,7 @@ def get_bprop_dtype(self):
 
     def bprop(x, out, dout):
         return (zeros_like(x),)
+
     return bprop
 
 
@@ -61,6 +62,7 @@ def get_bprop_cast(self):
     def bprop(x, t, out, dout):
         dx = cast(dout, get_dtype(x))
         return dx, zeros_like(t)
+
     return bprop
 
 
@@ -70,6 +72,7 @@ def get_bprop_shape(self):
 
     def bprop(x, out, dout):
         return (zeros_like(x),)
+
     return bprop
 
 
@@ -82,6 +85,7 @@ def get_bprop_split(self):
         concat_op = P.Concat(axis)
         dx = concat_op(dout)
         return (dx,)
+
     return bprop
 
 
@@ -91,6 +95,7 @@ def get_bprop_rank(self):
 
     def bprop(x, out, dout):
         return (zeros_like(x),)
+
     return bprop
 
 
@@ -101,6 +106,7 @@ def get_bprop_reshape(self):
     def bprop(x, shp, out, dout):
         shapex = shape_op(x)
         return reshape(dout, shapex), zeros_like(shp)
+
     return bprop
 
 
@@ -111,6 +117,7 @@ def get_bprop_expand_dims(self):
     def bprop(x, axis, out, dout):
         shapex = shape_op(x)
         return reshape(dout, shapex), zeros_like(axis)
+
     return bprop
 
 
@@ -121,6 +128,7 @@ def get_bprop_squeeze(self):
     def bprop(x, out, dout):
         shapex = shape_op(x)
         return (reshape(dout, shapex),)
+
     return bprop
 
 
@@ -131,6 +139,7 @@ def get_bprop_flatten(self):
     def bprop(x, out, dout):
         dx = reshape(dout, shape_op(x))
         return (dx,)
+
     return bprop
 
 
@@ -165,6 +174,7 @@ def _tile_shape(multiples, shapex):
 @bprop_getters.register(P.Tile)
 def get_bprop_tile(self):
     """Generate bprop for Tile"""
+
     def bprop(x, multiples, out, dout):
         shapex = shape_op(x)
         r_shape = _tile_shape(multiples, shapex)
@@ -173,6 +183,7 @@ def get_bprop_tile(self):
         dx = reduce_sum(reshape(dout, r_shape), axis)
         dx = reshape(dx, shapex)
         return dx, zeros_like(multiples)
+
     return bprop
 
 
@@ -182,6 +193,7 @@ def get_bprop_transpose(self):
 
     def bprop(x, perm, out, dout):
         return transpose(dout, invert_permutation(perm)), zeros_like(perm)
+
     return bprop
 
 
@@ -197,6 +209,7 @@ def get_bprop_concat(self):
             slice_out = P.Slice()(dout, out_offset[i], shape_op(x[i]))
             dx = dx + (slice_out,)
         return (dx,)
+
     return bprop
 
 
@@ -214,12 +227,12 @@ def get_bprop_slice(self):
         dx = P.Pad(_slice_grad_pad(begin, size, shape_op(x)))(dout)
         return (dx, zeros_like(begin), zeros_like(size))
 
-    def bprop_gpu(x, begin, size, out, dout):
+    def bprop_grad(x, begin, size, out, dout):
         dx = dx = G.SliceGrad()(dout, x, begin, size)
         return (dx, zeros_like(begin), zeros_like(size))
 
-    if context.get_context('device_target') == "GPU":
-        return bprop_gpu
+    if context.get_context('device_target') == "GPU" or context.get_context('device_target') == "CPU":
+        return bprop_grad
     return bprop
 
 
@@ -248,6 +261,7 @@ def _generate_inverse_index(x_shape, axis):
 @bprop_getters.register(P.GatherV2)
 def get_bprop_gather_v2(self):
     """Generate bprop for GatherV2"""
+
     def bprop(x, indices, axis, out, dout):
         if F.rank(dout) == 0:
             dout = P.ExpandDims()(dout, -1)
@@ -264,6 +278,7 @@ def get_bprop_gather_v2(self):
         perm_2 = _generate_inverse_index(x_shp, axis)
         params_grad = transpose(params_grad, perm_2)
         return params_grad, zeros_like(indices), zeros_like(axis)
+
     return bprop
 
 
@@ -276,6 +291,7 @@ def get_bprop_pack(self):
         pack_grad = P.Unpack(axis)
         out = pack_grad(dout)
         return (out,)
+
     return bprop
 
 
@@ -288,6 +304,7 @@ def get_bprop_unpack(self):
         unpack_grad = P.Pack(axis)
         out = unpack_grad(dout)
         return (out,)
+
     return bprop
 
 
@@ -303,6 +320,7 @@ def get_bprop_strided_slice(self):
     def bprop(x, begin, end, strides, out, dout):
         dx = input_grad(dout, shape_op(x), begin, end, strides)
         return dx, zeros_like(begin), zeros_like(end), zeros_like(strides)
+
     return bprop
 
 
@@ -312,6 +330,7 @@ def get_bprop_eye(self):
 
     def bprop(n, m, t, out, dout):
         return zeros_like(n), zeros_like(m), zeros_like(t)
+
     return bprop
 
 
@@ -322,6 +341,7 @@ def get_bprop_select(self):
 
     def bprop(cond, x, y, out, dout):
         return zeros_like(cond), select(cond, dout, zeros_like(x)), select(cond, zeros_like(y), dout)
+
     return bprop
 
 
@@ -512,9 +532,11 @@ def get_bprop_unsorted_segment_min(self):
 def get_bprop_space_to_batch(self):
     """Generate bprop for SpaceToBatch"""
     space_to_batch_grad = P.BatchToSpace(self.block_size, self.paddings)
+
     def bprop(x, out, dout):
         dx = space_to_batch_grad(dout)
         return (dx,)
+
     return bprop
 
 
@@ -522,7 +544,9 @@ def get_bprop_space_to_batch(self):
 def get_bprop_batch_to_space(self):
     """Generate bprop for BatchToSpace"""
     batch_to_space_grad = P.SpaceToBatch(self.block_size, self.crops)
+
     def bprop(x, out, dout):
         dx = batch_to_space_grad(dout)
         return (dx,)
+
     return bprop

mindspore/ops/_grad/grad_math_ops.py

@@ -15,7 +15,6 @@
 
 """Define the grad rules of math related operations."""
 
-
 from functools import reduce
 import numpy as np
 from .. import functional as F
@@ -26,7 +25,6 @@ from ..functional import broadcast_gradient_args, reduced_shape, tuple_div
 from .grad_base import bprop_getters
 from ..primitive import constexpr
 
-
 shape_op = P.Shape()
 reduce_sum = P.ReduceSum()
 reshape = P.Reshape()
@@ -129,6 +127,7 @@ def bprop_matmul(self):
         else:
             dw = mul2(x, dout)
         return dx, dw
+
     return bprop
 
 
@@ -152,6 +151,7 @@ def bprop_batchmatmul(self):
         else:
             dw = mul2(x, dout)
         return dx, dw
+
     return bprop
 
 
@@ -161,6 +161,7 @@ def get_bprop_tensor_add(self):
 
     def bprop(x, y, out, dout):
         return binop_grad_common(x, y, dout, dout)
+
     return bprop
 
 
@@ -172,6 +173,7 @@ def get_bprop_neg(self):
     def bprop(x, out, dout):
         dx = neg_grad(dout)
         return (dx,)
+
     return bprop
 
 
@@ -182,6 +184,7 @@ def get_bprop_sub(self):
 
     def bprop(x, y, out, dout):
         return binop_grad_common(x, y, dout, neg_func(dout))
+
     return bprop
 
 
@@ -194,6 +197,7 @@ def get_bprop_mul(self):
         bc_dx = mul_func(y, dout)
         bc_dy = mul_func(x, dout)
         return binop_grad_common(x, y, bc_dx, bc_dy)
+
     return bprop
 
 
@@ -208,6 +212,7 @@ def get_bprop_real_div(self):
         bc_x = div_op(dout, y)
         bc_y = neg(mul_op(bc_x, out))
         return binop_grad_common(x, y, bc_x, bc_y)
+
     return bprop
 
 
@@ -222,6 +227,7 @@ def get_bprop_div(self):
         bc_x = div_op(dout, y)
         bc_y = neg(mul_op(bc_x, out))
         return binop_grad_common(x, y, bc_x, bc_y)
+
     return bprop
 
 
@@ -235,6 +241,7 @@ def get_bprop_floor(self):
     def bprop(x, out, dout):
         bc_x = fill_(dtype_(x), shape_(x), 0.)
         return (bc_x,)
+
     return bprop
 
 
@@ -249,6 +256,7 @@ def get_bprop_floordiv(self):
         bc_x = div_op(dout, y)
         bc_y = neg(mul_op(bc_x, out))
         return binop_grad_common(x, y, bc_x, bc_y)
+
     return bprop
 
 
@@ -260,6 +268,7 @@ def get_bprop_floormod(self):
         bc_x = dout
         bc_y = -dout * (x // y)
         return binop_grad_common(x, y, bc_x, bc_y)
+
     return bprop
 
 
@@ -274,6 +283,7 @@ def get_bprop_square(self):
         temp = mul_func(dout, x)
         dx = mul_func(fill_func(dtype(temp), shape_op(x), 2.0), temp)
         return (dx,)
+
     return bprop
 
 
@@ -290,6 +300,7 @@ def get_bprop_sqrt(self):
         temp = div_op(fill_func(dtype(x), shape_op(x), 0.5), sqrt(x))
         dx = mul_func(dout, temp)
         return (dx,)
+
     return bprop
 
 
@@ -298,9 +309,10 @@ def get_bprop_rsqrt(self):
     """Grad definition for `Rsqrt` operation."""
 
     def bprop(x, out, dout):
-        grad = F.fill(F.dtype(x), F.shape(x), -0.5) / (F.sqrt(x)*x)
+        grad = F.fill(F.dtype(x), F.shape(x), -0.5) / (F.sqrt(x) * x)
         dx = dout * grad
         return (dx,)
+
     return bprop
 
 
@@ -316,6 +328,7 @@ def get_bprop_reciprocal(self):
         g = neg(reciprocal(square(x)))
         dx = mul(dout, g)
         return (dx,)
+
     return bprop
 
 
@@ -328,6 +341,7 @@ def get_bprop_log(self):
         g = reciprocal(x)
         dx = g * dout
         return dx, 0
+
     return bprop
 
 
@@ -341,6 +355,7 @@ def get_bprop_log1p(self):
         g = reciprocal(x_1p)
         dx = g * dout
         return dx, 0
+
     return bprop
 
 
@@ -358,6 +373,7 @@ def get_bprop_erf(self):
         x_square = square(x)
         dx = dout * half_root_pi * exp(-x_square)
         return (dx,)
+
     return bprop
 
 
@@ -388,6 +404,7 @@ def get_bprop_pow(self):
         bc_dx = power * pow_op(x, power - 1.0) * dout
         bc_dpower = out * ln(x) * dout
         return binop_grad_common(x, power, bc_dx, bc_dpower)
+
     return bprop
 
 
@@ -400,6 +417,7 @@ def get_bprop_exp(self):
         g = exp_(x)
         dx = g * dout
         return (dx,)
+
     return bprop
 
 
@@ -411,6 +429,7 @@ def get_bprop_minimum(self):
     def bprop(x, y, out, dout):
         dx, dy = input_grad(x, y, dout)
         return dx, dy
+
     return bprop
 
 
@@ -422,6 +441,7 @@ def get_bprop_maximum(self):
     def bprop(x, y, out, dout):
         dx, dy = input_grad(x, y, dout)
         return dx, dy
+
     return bprop
 
 
@@ -432,6 +452,7 @@ def get_bprop_reducesum(self):
     def bprop(x, axis, out, dout):
         dx = _sum_grad(x, axis, dout)
         return dx, zeros_like(axis)
+
     return bprop
 
 
@@ -442,6 +463,7 @@ def get_bprop_cumsum(self):
 
     def bprop(x, axis, out, dout):
         return cumsum(dout, axis), zeros_like(axis)
+
     return bprop
 
 
@@ -500,6 +522,7 @@ def get_bprop_reduceprod(self):
         out = transpose(y, _invert_permutation(perm)) * grad
         dx = reshape(out, input_shape)
         return dx, zeros_like(axis)
+
     return bprop
 
 
@@ -515,6 +538,7 @@ def get_bprop_cumprod(self):
         prod = cumprod(x, axis)
         out = cumsum(prod * dout, axis)
         return out / x, zeros_like(axis)
+
     return bprop
 
 
@@ -524,6 +548,7 @@ def get_bprop_reduceall(self):
 
     def bprop(x, axis, out, dout):
         return zeros_like(x), zeros_like(axis)
+
     return bprop
 
 
@@ -534,6 +559,7 @@ def get_bprop_reducemax(self):
     def bprop(x, axis, out, dout):
         dx = _min_or_max_grad(x, axis, out, dout)
         return (dx, zeros_like(axis))
+
     return bprop
 
 
@@ -547,6 +573,7 @@ def get_bprop_argmaxwithvalue(self):
     def bprop(x, out, dout):
         dx = _argmin_or_argmax_grad(x, axis, keep_dims, op, out, dout)
         return (dx,)
+
     return bprop
 
 
@@ -557,6 +584,7 @@ def get_bprop_reducemin(self):
     def bprop(x, axis, out, dout):
         dx = _min_or_max_grad(x, axis, out, dout)
         return (dx, zeros_like(axis))
+
     return bprop
 
 
@@ -570,6 +598,7 @@ def get_bprop_argminwithvalue(self):
     def bprop(x, out, dout):
         dx = _argmin_or_argmax_grad(x, axis, keep_dims, op, out, dout)
         return (dx,)
+
     return bprop
 
 
@@ -585,6 +614,7 @@ def get_bprop_reduce_mean(self):
         div_shape = F.shape_mul(shape_op(x)) / F.shape_mul(shape_op(out))
         dx = div_op(grad, cast(F.scalar_to_array(div_shape), dtype(grad)))
         return dx, zeros_like(axis)
+
     return bprop
 
 
@@ -604,6 +634,7 @@ def get_bprop_not_equal(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -613,6 +644,7 @@ def get_bprop_greater(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -622,6 +654,7 @@ def get_bprop_greater_equal(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -631,6 +664,7 @@ def get_bprop_less(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -640,6 +674,7 @@ def get_bprop_less_equal(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -649,6 +684,7 @@ def get_bprop_logical_not(self):
 
     def bprop(x, out, dout):
         return (zeros_like(x),)
+
     return bprop
 
 
@@ -658,6 +694,7 @@ def get_bprop_logical_and(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -667,6 +704,7 @@ def get_bprop_logical_or(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -676,6 +714,7 @@ def get_bprop_npu_alloc_float_status(self):
 
     def bprop(out, dout):
         return ()
+
     return bprop
 
 
@@ -685,6 +724,7 @@ def get_bprop_npu_get_float_status(self):
 
     def bprop(x, out, dout):
         return (zeros_like(x),)
+
     return bprop
 
 
@@ -694,6 +734,7 @@ def get_bprop_npu_clear_float_status(self):
 
     def bprop(x, out, dout):
         return (zeros_like(x),)
+
     return bprop
 
 
@@ -703,6 +744,7 @@ def get_bprop_assign_add(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -712,6 +754,7 @@ def get_bprop_assign_sub(self):
 
     def bprop(x, y, out, dout):
         return zeros_like(x), zeros_like(y)
+
     return bprop
 
 
@@ -721,8 +764,9 @@ def get_bprop_sin(self):
     cos = P.Cos()
 
     def bprop(x, out, dout):
-        dx = dout*cos(x)
+        dx = dout * cos(x)
         return (dx,)
+
     return bprop
 
 
@@ -733,8 +777,9 @@ def get_bprop_cos(self):
     neg = P.Neg()
 
     def bprop(x, out, dout):
-        dx = dout*neg(sin(x))
+        dx = dout * neg(sin(x))
         return (dx,)
+
     return bprop
 
 
@@ -746,6 +791,7 @@ def get_bprop_acos(self):
     def bprop(x, out, dout):
         dx = input_grad(x, dout)
         return (dx,)
+
     return bprop
 
 
@@ -757,6 +803,7 @@ def get_bprop_acosh(self):
     def bprop(x, out, dout):
         dx = input_grad(out, dout)
         return (dx,)
+
     return bprop
 
 
@@ -768,6 +815,7 @@ def get_bprop_abs(self):
     def bprop(x, out, dout):
         dx = abs_grad(x, dout)
         return (dx,)
+
     return bprop
 
 
@@ -777,6 +825,7 @@ def get_bprop_scalar_cast(self):
 
     def bprop(x, t, out, dout):
         return F.scalar_cast(dout, F.typeof(x)), zeros_like(t)
+
     return bprop
 
 
@@ -789,6 +838,7 @@ def get_bprop_scalar_addn(self):
         for _ in range(len(x)):
             dx = dx + (dout,)
         return dx
+
     return bprop
 
 
@@ -798,6 +848,7 @@ def get_bprop_sign(self):
 
     def bprop(x, out, dout):
         return (zeros_like(x),)
+
     return bprop
 
 
@@ -807,6 +858,7 @@ def get_bprop_round(self):
 
     def bprop(x, out, dout):
         return (zeros_like(x),)
+
     return bprop
 
 
@@ -821,4 +873,5 @@ def get_bprop_atan2(self):
         bc_dx = tmp * y
         bc_dy = tmp * (-x)
         return binop_grad_common(x, y, bc_dx, bc_dy)
+
     return bprop

mindspore/ops/_grad/grad_nn_ops.py

@@ -21,6 +21,7 @@ from ..operations import _grad_ops as G
 from ..operations import _inner_ops as inner
 from ..composite.multitype_ops.zeros_like_impl import zeros_like
 from .grad_base import bprop_getters
+from ... import context
 
 
 @bprop_getters.register(P.BiasAdd)
@@ -550,6 +551,14 @@ def get_bprop_lstm(self):
         bidirectional=self.bidirectional,
         dropout=self.dropout
     )
+    lstm_grad = G.LSTMGrad(
+        input_size=self.input_size,
+        hidden_size=self.hidden_size,
+        num_layers=self.num_layers,
+        has_bias=self.has_bias,
+        bidirectional=self.bidirectional,
+        dropout=self.dropout
+    )
 
     def bprop(x, hx, cx, w, out, dout):
         y, _, _, reserve, state = out
@@ -558,6 +567,16 @@ def get_bprop_lstm(self):
         dw = lstm_grad_weight(F.depend(x, dx), hx, y, reserve, state)
         return dx, dhx, dcx, dw
 
+    #
+    def bprop_cpu(x, hx, cx, w, out, dout):
+        y, hy, cy, reserve, _ = out
+        dy, dhy, dcy, _, _ = dout
+        dx, dhx, dcx, dw = lstm_grad(x, hx, cx, w, y, hy, cy, dy, dhy, dcy, reserve)
+        return dx, dhx, dcx, dw
+
+    if context.get_context('device_target') == "CPU":
+        return bprop_cpu
+
     return bprop
 
 

mindspore/ops/operations/_grad_ops.py

@@ -107,6 +107,7 @@ class BiasAddGrad(Primitive):
 
 class BinaryCrossEntropyGrad(PrimitiveWithInfer):
     """Computes gradients for `BinaryCrossEntropy` operation."""
+
     @prim_attr_register
     def __init__(self, reduction='mean'):
         self.reduction = validator.check_string('reduction', reduction, ['none', 'mean', 'sum'], self.name)
@@ -665,6 +666,62 @@ class LSTMGradWeight(PrimitiveWithInfer):
         return hx_dtype
 
 
+class LSTMGrad(PrimitiveWithInfer):
+    """Computes the data and weight gradients of LSTM."""
+
+    @prim_attr_register
+    def __init__(self, input_size, hidden_size, num_layers, has_bias, bidirectional, dropout):
+        self.input_size = validator.check_integer('input_size', input_size, 0, Rel.GT, self.name)
+        self.hidden_size = validator.check_integer('hidden_size', hidden_size, 0, Rel.GT, self.name)
+        self.num_layers = validator.check_integer('num_layers', num_layers, 0, Rel.GT, self.name)
+        self.has_bias = validator.check_value_type('has_bias', has_bias, (bool,), self.name)
+        self.bidirectional = validator.check_value_type('bidirectional', bidirectional, (bool,), self.name)
+        self.dropout = validator.check_value_type("dropout", dropout, [float], self.name)
+        self.dropout = validator.check_number_range('dropout', dropout, 0, 1, Rel.INC_BOTH, self.name)
+
+        if bidirectional:
+            self.num_directions = 2
+        else:
+            self.num_directions = 1
+
+    def infer_shape(self, x_shape, hx_shape, cx_shape, w_shape, y_shape, hy_shape, cy_shape, dy_shape, dhy_shape,
+                    dcy_shape, reserve_shape):
+        # dhy and dcy should be same shape
+        validator.check_integer("h_shape", len(dhy_shape), 3, Rel.EQ, self.name)
+        validator.check_integer("h_shape", len(dhy_shape), len(dcy_shape), Rel.EQ, self.name)
+        validator.check_integer("h_shape[0]", dhy_shape[0], dcy_shape[0], Rel.EQ, self.name)
+        validator.check_integer("h_shape[1]", dhy_shape[1], dcy_shape[1], Rel.EQ, self.name)
+        validator.check_integer("h_shape[2]", dhy_shape[2], dcy_shape[2], Rel.EQ, self.name)
+
+        validator.check_integer("h_shape[0]", dhy_shape[0], self.num_layers * self.num_directions, Rel.EQ, self.name)
+        validator.check_integer("h_shape[2]", dhy_shape[2], self.hidden_size, Rel.EQ, self.name)
+
+        # dy: (seq_len, batch_size, hidden_size * num_directions)
+        validator.check_integer("dy_shape", len(dy_shape), 3, Rel.EQ, self.name)
+        validator.check_integer("dy[1]", dy_shape[1], dhy_shape[1], Rel.EQ, self.name)
+        validator.check_integer("dy[2]", dy_shape[2], self.hidden_size * self.num_directions, Rel.EQ, self.name)
+
+        # (seq_len, batch_size, input_size)
+        dx_shape = (y_shape[0], y_shape[1], self.input_size)
+        dhx_shape = dhy_shape
+        dcx_shape = dcy_shape
+        weight_size = 0
+        gate_size = 4 * self.hidden_size
+        for layer in range(self.num_layers):
+            for _ in range(self.num_directions):
+                input_layer_size = self.input_size if layer == 0 else self.hidden_size * self.num_directions
+                weight_size += gate_size * input_layer_size
+                weight_size += gate_size * self.hidden_size
+                if self.has_bias:
+                    weight_size += gate_size
+
+        return (dx_shape, dhx_shape, dcx_shape, (weight_size, 1, 1))
+
+    def infer_dtype(self, x_dtype, hx_dtype, cx_dtype, w_dtype, y_dtype, hy_dtype, cy_dtype, dy_dtype, dhy_dtype,
+                    dcy_dtype, reserve_dtype):
+        return (dy_dtype, dy_dtype, dy_dtype, hx_dtype)
+
+
 class PReLUGrad(PrimitiveWithInfer):
     r"""
     Gradients of PReLU operation.
@@ -1051,6 +1108,7 @@ class RefToEmbed(Primitive):
     __mindspore_signature__ = (
         ('variable', sig_rw.RW_REF, sig_kind.KIND_POSITIONAL_KEYWORD),
     )
+
     @prim_attr_register
     def __init__(self):
         pass

mindspore/ops/operations/nn_ops.py

@@ -35,9 +35,11 @@ def _check_positive_int_or_tuple(arg_name, arg_value, prim_name, allow_four=Fals
     """
     Checks whether an argument is a positive int or tuple with 2 or 4(when allow_four is True) positive int elements.
     """
+
     def _raise_message():
         raise ValueError(f"For '{prim_name}' attr '{arg_name}' should be an positive int number or a tuple of two "
                          f"{'or four ' if allow_four else ''}positive int numbers, but got {arg_value}")
+
     def _get_return_value():
         if isinstance(arg_value, int):
             ret = (1, 1, arg_value, arg_value) if ret_four else (arg_value, arg_value)
@@ -50,6 +52,7 @@ def _check_positive_int_or_tuple(arg_name, arg_value, prim_name, allow_four=Fals
         else:
             _raise_message()
         return ret
+
     validator.check_value_type(arg_name, arg_value, (int, tuple), prim_name)
     ret_value = _get_return_value()
     for item in ret_value:
@@ -58,6 +61,7 @@ def _check_positive_int_or_tuple(arg_name, arg_value, prim_name, allow_four=Fals
         _raise_message()
     return ret_value
 
+
 class Flatten(PrimitiveWithInfer):
     r"""
     Flattens a tensor without changing its batch size on the 0-th axis.
@@ -205,6 +209,7 @@ class Softplus(PrimitiveWithInfer):
         >>> softplus(input_x)
         [1.3132615, 2.126928, 3.0485873, 4.01815, 5.0067153]
     """
+
     @prim_attr_register
     def __init__(self):
         """init Softplus"""
@@ -301,6 +306,7 @@ class ReLUV2(PrimitiveWithInfer):
         ([[[[1., 0.], [0., 4.]], [[0., 6.], [7., 0.]]]],
          [[[[1, 0], [2, 0]], [[2, 0], [1, 0]]]])
     """
+
     @prim_attr_register
     def __init__(self):
         """init ReLUV2"""
@@ -398,6 +404,7 @@ class HSwish(PrimitiveWithInfer):
         >>> input_x = Tensor(np.array([-1, -2, 0, 2, 1]), mindspore.float16)
         >>> result = hswish(input_x)
     """
+
     @prim_attr_register
     def __init__(self):
         self.init_prim_io_names(inputs=['x'], outputs=['output'])
@@ -1077,6 +1084,7 @@ class MaxPoolWithArgmax(_Pool):
         >>> maxpool_arg_op = P.MaxPoolWithArgmax(padding="VALID", ksize=2, strides=1)
         >>> output_tensor, argmax = maxpool_arg_op(input_tensor)
     """
+
     def __init__(self, ksize=1, strides=1, padding="valid"):
         super(MaxPoolWithArgmax, self).__init__(ksize, strides, padding)
         self.is_tbe = context.get_context("device_target") == "Ascend"
@@ -1495,6 +1503,7 @@ class ApplyMomentum(PrimitiveWithInfer):
         ('gradient', sig_rw.RW_READ, sig_kind.KIND_POSITIONAL_KEYWORD),
         ('momentum', sig_rw.RW_READ, sig_kind.KIND_POSITIONAL_KEYWORD)
     )
+
     @prim_attr_register
     def __init__(self, use_nesterov=False, use_locking=False, gradient_scale=1.0):
         self.init_prim_io_names(inputs=['variable', 'accumulation', 'learning_rate', 'gradient', 'momentum'],
@@ -1584,6 +1593,7 @@ class L2Loss(PrimitiveWithInfer):
         >>> l2_loss(input_x)
         7.0
     """
+
     @prim_attr_register
     def __init__(self):
         """init L2Loss"""
@@ -2326,7 +2336,29 @@ class LSTM(PrimitiveWithInfer):
         y_shape = (x_shape[0], x_shape[1], self.hidden_size * self.num_directions)
 
         # set arbitrary shape for reserved space
-        reserved_shape = (1, 1)
+        type_size = 4
+        gates_ws_ld = self.get_good_ld(self.hidden_size * 4, type_size)
+        states_ws_ld = self.get_good_ld(max(self.hidden_size, self.input_size), type_size)
+        self.ws_gates_size = self.num_layers * self.num_directions * x_shape[0] * x_shape[1] * gates_ws_ld * type_size
+        self.ws_states_size = (self.num_layers + 1) * self.num_directions * (x_shape[0] + 1) * x_shape[
+            1] * states_ws_ld * type_size
+        self.ws_c_states_size = (self.num_layers + 1) * self.num_directions * (x_shape[0] + 1) * x_shape[
+            1] * states_ws_ld * type_size
+        self.ws_diff_states_size = (self.num_layers + 1) * self.num_directions * (x_shape[0] + 1) * (2 + 1) * x_shape[
+            1] * states_ws_ld * type_size
+        self.ws_grid_comp_size = 0
+        self.page_size = 4096
+        current_offset = 0
+        current_offset += self.ws_gates_size
+        current_offset = self.rnd_up(current_offset, self.page_size)
+        current_offset += self.ws_states_size
+        current_offset = self.rnd_up(current_offset, self.page_size)
+        current_offset += self.ws_c_states_size
+        current_offset = self.rnd_up(current_offset, self.page_size)
+        current_offset += self.ws_diff_states_size
+        current_offset = self.rnd_up(current_offset, self.page_size)
+        current_offset += self.ws_grid_comp_size
+        reserved_shape = (current_offset, 1)
         state_shape = (1, 1)
         return (y_shape, h_shape, c_shape, reserved_shape, state_shape)
 
@@ -2335,6 +2367,15 @@ class LSTM(PrimitiveWithInfer):
         validator.check_tensor_type_same(args, (mstype.float32, mstype.float16), self.name)
         return (x_dtype, x_dtype, x_dtype, x_dtype, x_dtype)
 
+    def rnd_up(self, current_offset, page_size):
+        return ((current_offset + page_size - 1) // page_size) * page_size
+
+    def get_good_ld(self, dim, type_size):
+        ld = self.rnd_up(dim, 64 // type_size)
+        if ld * 256 == 0:
+            return ld + 64 // type_size
+        return ld
+
 
 class SigmoidCrossEntropyWithLogits(PrimitiveWithInfer):
     r"""
@@ -3001,6 +3042,7 @@ class Dropout(PrimitiveWithInfer):
         >>> in = Tensor((20, 16, 50, 50))
         >>> out = dropout(in)
     """
+
     @prim_attr_register
     def __init__(self, drop_prob=0):
         self.drop_prob = validator.check_number_range("drop_prob", drop_prob, 0, 1, Rel.INC_BOTH, self.name)
@@ -3035,6 +3077,7 @@ class DropoutGrad(PrimitiveWithInfer):
         >>> in = Tensor((20, 16, 50, 50))
         >>> out = dropout_grad(in)
     """
+
     @prim_attr_register
     def __init__(self, drop_prob=0):
         self.drop_prob = validator.check_number_range("drop_prob", drop_prob, 0, 1, Rel.INC_BOTH, self.name)
@@ -3085,6 +3128,7 @@ class CTCLoss(PrimitiveWithInfer):
         >>> ctc_loss = P.CTCloss()
         >>> output = ctc_loss(inputs, labels_indices, labels_values, sequence_length)
     """
+
     @prim_attr_register
     def __init__(self, preprocess_collapse_repeated=False, ctc_merge_repeated=False,
                  ignore_longer_outputs_than_inputs=False):

tests/st/ops/cpu/test_lstm_op.py

@@ -0,0 +1,285 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed 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.
+# ============================================================================
+
+import pytest
+import mindspore.nn as nn
+from mindspore.common.api import ms_function
+import numpy as np
+import mindspore.context as context
+from mindspore.common.initializer import initializer
+from mindspore.ops import composite as C
+from mindspore.ops import operations as P
+from mindspore.common.tensor import Tensor
+from mindspore.common.parameter import ParameterTuple, Parameter
+
+context.set_context(device_target='CPU')
+
+
+class LstmNet(nn.Cell):
+    def __init__(self, seq_len, batch_size, input_size, hidden_size, num_layers, has_bias, bidirectional, dropout):
+        super(LstmNet, self).__init__()
+
+        num_directions = 1
+        if bidirectional:
+            num_directions = 2
+
+        self.lstm = P.LSTM(input_size, hidden_size, num_layers, has_bias, bidirectional, dropout)
+        input_np = np.array([[[0.6755, -1.6607, 0.1367], [0.4276, -0.7850, -0.3758]],
+                             [[-0.6424, -0.6095, 0.6639], [0.7918, 0.4147, -0.5089]],
+                             [[-1.5612, 0.0120, -0.7289], [-0.6656, -0.6626, -0.5883]],
+                             [[-0.9667, -0.6296, -0.7310], [0.1026, -0.6821, -0.4387]],
+                             [[-0.4710, 0.6558, -0.3144], [-0.8449, -0.2184, -0.1806]]
+                             ]).astype(np.float32)
+        self.x = Parameter(initializer(Tensor(input_np), [seq_len, batch_size, input_size]), name='x')
+
+        self.h = Parameter(initializer(
+            Tensor(
+                np.array([0.1, 0.1, 0.1, 0.1]).reshape((num_directions, batch_size, hidden_size)).astype(
+                    np.float32)),
+            [num_directions, batch_size, hidden_size]), name='h')
+
+        self.c = Parameter(initializer(
+            Tensor(
+                np.array([0.2, 0.2, 0.2, 0.2]).reshape((num_directions, batch_size, hidden_size)).astype(
+                    np.float32)),
+            [num_directions, batch_size, hidden_size]), name='c')
+
+        wih = np.array([[3.4021e-01, -4.6622e-01, 4.5117e-01],
+                        [-6.4257e-02, -2.4807e-01, 1.3550e-02],  # i
+                        [-3.2140e-01, 5.5578e-01, 6.3589e-01],
+                        [1.6547e-01, -7.9030e-02, -2.0045e-01],
+                        [-6.9863e-01, 5.9773e-01, -3.9062e-01],
+                        [-3.0253e-01, -1.9464e-01, 7.0591e-01],
+                        [-4.0835e-01, 3.6751e-01, 4.7989e-01],
+                        [-5.6894e-01, -5.0359e-01, 4.7491e-01]]).astype(np.float32)  # .reshape([1,-1])
+        whh = np.array([[-0.4820, -0.2350],
+                        [-0.1195, 0.0519],
+                        [0.2162, -0.1178],
+                        [0.6237, 0.0711],
+                        [0.4511, -0.3961],
+                        [-0.5962, 0.0906],
+                        [0.1867, -0.1225],
+                        [0.1831, 0.0850]]).astype(np.float32)  # .reshape([1,-1])
+        wih = wih.transpose((1, 0))
+        whh = whh.transpose((1, 0))
+        bih = np.zeros((1, 8)).astype(np.float32)
+        w_np = np.concatenate((wih, whh, bih), axis=0).reshape([-1, 1, 1])
+        self.w = Parameter(initializer(Tensor(w_np), w_np.shape), name='w')
+
+    @ms_function
+    def construct(self):
+        return self.lstm(self.x, self.h, self.c, self.w)
+
+
+@pytest.mark.level0
+@pytest.mark.platform_x86_cpu
+@pytest.mark.env_onecard
+def test_lstm():
+    seq_len = 5
+    batch_size = 2
+    input_size = 3
+    hidden_size = 2
+    num_layers = 1
+    has_bias = True
+    bidirectional = False
+    dropout = 0.0
+    num_directions = 1
+    if bidirectional:
+        num_directions = 2
+    net = LstmNet(seq_len, batch_size, input_size, hidden_size, num_layers, has_bias, bidirectional, dropout)
+    y, h, c, _, _ = net()
+    print(y)
+    print(c)
+    print(h)
+
+class MultiLayerBiLstmNet(nn.Cell):
+    def __init__(self, seq_len, batch_size, input_size, hidden_size, num_layers, has_bias, bidirectional, dropout):
+        super(MultiLayerBiLstmNet, self).__init__()
+
+        num_directions = 1
+        if bidirectional:
+            num_directions = 2
+
+        self.lstm = nn.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers, has_bias=has_bias,
+                            bidirectional=bidirectional, dropout=dropout)
+
+        input_np = np.array([[[-0.1887, -0.4144, -0.0235, 0.7489, 0.7522, 0.5969, 0.3342, 1.2198, 0.6786, -0.9404],
+                              [-0.8643, -1.6835, -2.4965, 2.8093, 0.1741, 0.2707, 0.7387, -0.0939, -1.7990, 0.4765]],
+
+                             [[-0.5963, -1.2598, -0.7226, 1.1365, -1.7320, -0.7302, 0.1221, -0.2111, -1.6173, -0.0706],
+                              [0.8964, 0.1737, -1.0077, -0.1389, 0.4889, 0.4391, 0.7911, 0.3614, -1.9533, -0.9936]],
+
+                             [[0.3260, -1.3312, 0.0601, 1.0726, -1.6010, -1.8733, -1.5775, 1.1579, -0.8801, -0.5742],
+                              [-2.2998, -0.6344, -0.5409, -0.9221, -0.6500, 0.1206, 1.5215, 0.7517, 1.3691, 2.0021]],
+
+                             [[-0.1245, -0.3690, 2.1193, 1.3852, -0.1841, -0.8899, -0.3646, -0.8575, -0.3131, 0.2026],
+                              [1.0218, -1.4331, 0.1744, 0.5442, -0.7808, 0.2527, 0.1566, 1.1484, -0.7766, -0.6747]],
+
+                             [[-0.6752, 0.9906, -0.4973, 0.3471, -0.1202, -0.4213, 2.0213, 0.0441, 0.9016, 1.0365],
+                              [1.2223, -1.3248, 0.1207, -0.8256, 0.1816, 0.7057, -0.3105, 0.5713, 0.2804,
+                               -1.0685]]]).astype(np.float32)
+
+        self.x = Parameter(initializer(Tensor(input_np), [seq_len, batch_size, input_size]), name='x')
+
+        self.h0 = Parameter(initializer(
+            Tensor(np.ones((num_directions, batch_size, hidden_size)).astype(np.float32)),
+            [num_directions, batch_size, hidden_size]), name='h0')
+        self.c0 = Parameter(initializer(
+            Tensor(np.ones((num_directions, batch_size, hidden_size)).astype(np.float32)),
+            [num_directions, batch_size, hidden_size]), name='c0')
+        self.h1 = Parameter(initializer(
+            Tensor(np.ones((num_directions, batch_size, hidden_size)).astype(np.float32)),
+            [num_directions, batch_size, hidden_size]), name='h1')
+        self.c1 = Parameter(initializer(
+            Tensor(np.ones((num_directions, batch_size, hidden_size)).astype(np.float32)),
+            [num_directions, batch_size, hidden_size]), name='c1')
+        self.h = ParameterTuple((self.h0, self.h1))
+        self.c = ParameterTuple((self.c0, self.c1))
+
+    @ms_function
+    def construct(self):
+        return self.lstm(self.x, (self.h, self.c))
+
+
+@pytest.mark.level0
+@pytest.mark.platform_x86_cpu
+@pytest.mark.env_onecard
+def test_multi_layer_bilstm():
+    seq_len = 5
+    batch_size = 2
+    input_size = 10
+    hidden_size = 2
+    num_layers = 2
+    has_bias = True
+    bidirectional = True
+    dropout = 0.0
+
+    num_directions = 1
+    if bidirectional:
+        num_directions = 2
+
+    net = MultiLayerBiLstmNet(seq_len, batch_size, input_size, hidden_size, num_layers, has_bias, bidirectional,
+                              dropout)
+    y, h, c, _, _ = net()
+    print(y)
+    print(h)
+    print(c)
+
+class Grad(nn.Cell):
+    def __init__(self, network):
+        super(Grad, self).__init__()
+        self.network = network
+        self.weights = ParameterTuple(network.trainable_params())
+        self.grad = C.GradOperation('grad',
+                                    get_by_list=True,
+                                    sens_param=True)
+
+    @ms_function
+    def construct(self, output_grad):
+        weights = self.weights
+        grads = self.grad(self.network, weights)(output_grad)
+        return grads
+
+class Net(nn.Cell):
+    def __init__(self, seq_len, batch_size, input_size, hidden_size, num_layers, has_bias, bidirectional, dropout):
+        super(Net, self).__init__()
+
+        num_directions = 1
+        if bidirectional:
+            num_directions = 2
+        input_np = np.array([[[0.6755, -1.6607, 0.1367], [0.4276, -0.7850, -0.3758]],
+                             [[-0.6424, -0.6095, 0.6639], [0.7918, 0.4147, -0.5089]],
+                             [[-1.5612, 0.0120, -0.7289], [-0.6656, -0.6626, -0.5883]],
+                             [[-0.9667, -0.6296, -0.7310], [0.1026, -0.6821, -0.4387]],
+                             [[-0.4710, 0.6558, -0.3144], [-0.8449, -0.2184, -0.1806]]
+                             ]).astype(np.float32)
+        self.x = Parameter(initializer(Tensor(input_np), [seq_len, batch_size, input_size]), name='x')
+        self.h = Parameter(initializer(
+            Tensor(
+                np.array([0.1, 0.1, 0.1, 0.1]).reshape((num_directions, batch_size, hidden_size)).astype(
+                    np.float32)),
+            [num_directions, batch_size, hidden_size]), name='h')
+        self.c = Parameter(initializer(
+            Tensor(
+                np.array([0.2, 0.2, 0.2, 0.2]).reshape((num_directions, batch_size, hidden_size)).astype(
+                    np.float32)),
+            [num_directions, batch_size, hidden_size]), name='c')
+        wih = np.array([[3.4021e-01, -4.6622e-01, 4.5117e-01],
+                        [-6.4257e-02, -2.4807e-01, 1.3550e-02],  # i
+                        [-3.2140e-01, 5.5578e-01, 6.3589e-01],
+                        [1.6547e-01, -7.9030e-02, -2.0045e-01],
+                        [-6.9863e-01, 5.9773e-01, -3.9062e-01],
+                        [-3.0253e-01, -1.9464e-01, 7.0591e-01],
+                        [-4.0835e-01, 3.6751e-01, 4.7989e-01],
+                        [-5.6894e-01, -5.0359e-01, 4.7491e-01]]).astype(np.float32).reshape([1,-1])
+        whh = np.array([[-0.4820, -0.2350],
+                        [-0.1195, 0.0519],
+                        [0.2162, -0.1178],
+                        [0.6237, 0.0711],
+                        [0.4511, -0.3961],
+                        [-0.5962, 0.0906],
+                        [0.1867, -0.1225],
+                        [0.1831, 0.0850]]).astype(np.float32).reshape([1,-1])
+        bih = np.zeros((1, 8)).astype(np.float32)
+        w_np = np.concatenate((wih, whh, bih), axis=1).reshape([-1, 1, 1])
+        self.w = Parameter(initializer(Tensor(w_np), w_np.shape), name='w')
+
+        self.lstm = P.LSTM(input_size=input_size, hidden_size=hidden_size, num_layers=num_layers,
+                           has_bias=has_bias, bidirectional=bidirectional, dropout=dropout)
+
+    @ms_function
+    def construct(self):
+        return self.lstm(self.x, self.h, self.c, self.w)[0]
+
+@pytest.mark.level0
+@pytest.mark.platform_x86_cpu
+@pytest.mark.env_onecard
+def test_grad():
+    seq_len = 5
+    batch_size = 2
+    input_size = 3
+    hidden_size = 2
+    num_layers = 1
+    has_bias = True
+    bidirectional = False
+    dropout = 0.0
+    num_directions = 1
+    if bidirectional:
+        num_directions = 2
+    net = Grad(Net(seq_len, batch_size, input_size, hidden_size, num_layers, has_bias, bidirectional, dropout))
+    dy = np.array([[[-3.5471e-01, 7.0540e-01],
+                    [2.7161e-01, 1.0865e+00]],
+
+                   [[-4.2431e-01, 1.4955e+00],
+                    [-4.0418e-01, -2.3282e-01]],
+
+                   [[-1.3654e+00, 1.9251e+00],
+                    [-4.6481e-01, 1.3138e+00]],
+
+                   [[1.2914e+00, -2.3753e-01],
+                    [5.3589e-01, -1.0981e-01]],
+
+                   [[-1.6032e+00, -1.8818e-01],
+                    [1.0065e-01, 9.2045e-01]]]).astype(np.float32)
+    dx, dhx, dcx, dw = net(Tensor(dy))
+    print(dx)
+    print(dhx)
+    print(dcx)
+    print(dw)
+
+# test_multi_layer_bilstm()
+# test_lstm()
+test_grad()
+# tf_lstm_test()

tests/st/ops/cpu/test_transpose_op.py

@@ -0,0 +1,145 @@
+# Copyright 2020 Huawei Technologies Co., Ltd
+#
+# Licensed 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.
+# ============================================================================
+
+import pytest
+import numpy as np
+from mindspore import Tensor
+from mindspore.ops import operations as P
+from mindspore.common.api import ms_function
+from mindspore.common.initializer import initializer
+from mindspore.common.parameter import Parameter
+import mindspore.nn as nn
+import mindspore.context as context
+
+context.set_context(device_target='CPU')
+
+
+class Transpose(nn.Cell):
+    def __init__(self):
+        super(Transpose, self).__init__()
+        self.transpose = P.Transpose()
+
+        self.x_2D = Parameter(initializer(Tensor(np.arange(5 * 6).reshape(5, 6).astype(np.float32)), [5, 6]),
+                              name='x_2D')
+        self.perm_2D = (1, 0)
+
+        self.x_3D = Parameter(initializer(Tensor(np.arange(2 * 2 * 4).reshape(2, 2, 4).astype(np.float32)), [2, 2, 4]),
+                              name='x_3D')
+        self.perm_3D = (1, 0, 2)
+
+        self.x_4D = Parameter(
+            initializer(Tensor(np.arange(2 * 3 * 4 * 5).reshape(2, 3, 4, 5).astype(np.float32)), [2, 3, 4, 5]),
+            name='x_4D')
+        self.perm_4D = (0, 1, 2, 3)
+
+        self.x_5D = Parameter(
+            initializer(Tensor(np.arange(1 * 2 * 3 * 4 * 5).reshape(1, 2, 3, 4, 5).astype(np.float32)),
+                        [1, 2, 3, 4, 5]), name='x_5D')
+        self.perm_5D = (1, 0, 3, 4, 2)
+
+    @ms_function
+    def construct(self):
+        return (self.transpose(self.x_2D, self.perm_2D), self.transpose(self.x_3D, self.perm_3D),
+                self.transpose(self.x_4D, self.perm_4D), self.transpose(self.x_5D, self.perm_5D))
+
+
+@pytest.mark.level0
+@pytest.mark.platform_x86_cpu
+@pytest.mark.env_onecard
+def test_transpose():
+    transpose = Transpose()
+    output = transpose()
+
+    expect0 = np.array([[[0, 6, 12, 18, 24],
+                         [1, 7, 13, 19, 25],
+                         [2, 8, 14, 20, 26],
+                         [3, 9, 15, 21, 27],
+                         [4, 10, 16, 22, 28],
+                         [5, 11, 17, 23, 29]]]).astype(np.float32)
+    expect1 = np.array([[[[0, 1, 2, 3],
+                          [8, 9, 10, 11]],
+                         [[4, 5, 6, 7],
+                          [12, 13, 14, 15]]]]).astype(np.float32)
+    expect2 = np.array([[[[[0, 1, 2, 3, 4],
+                           [5, 6, 7, 8, 9],
+                           [10, 11, 12, 13, 14],
+                           [15, 16, 17, 18, 19]],
+                          [[20, 21, 22, 23, 24],
+                           [25, 26, 27, 28, 29],
+                           [30, 31, 32, 33, 34],
+                           [35, 36, 37, 38, 39]],
+                          [[40, 41, 42, 43, 44],
+                           [45, 46, 47, 48, 49],
+                           [50, 51, 52, 53, 54],
+                           [55, 56, 57, 58, 59]]],
+
+                         [[[60, 61, 62, 63, 64],
+                           [65, 66, 67, 68, 69],
+                           [70, 71, 72, 73, 74],
+                           [75, 76, 77, 78, 79]],
+                          [[80, 81, 82, 83, 84],
+                           [85, 86, 87, 88, 89],
+                           [90, 91, 92, 93, 94],
+                           [95, 96, 97, 98, 99]],
+                          [[100, 101, 102, 103, 104],
+                           [105, 106, 107, 108, 109],
+                           [110, 111, 112, 113, 114],
+                           [115, 116, 117, 118, 119]]]]]).astype(np.float32)
+    expect3 = np.array([[[[[[0, 20, 40],
+                            [1, 21, 41],
+                            [2, 22, 42],
+                            [3, 23, 43],
+                            [4, 24, 44]],
+                           [[5, 25, 45],
+                            [6, 26, 46],
+                            [7, 27, 47],
+                            [8, 28, 48],
+                            [9, 29, 49]],
+                           [[10, 30, 50],
+                            [11, 31, 51],
+                            [12, 32, 52],
+                            [13, 33, 53],
+                            [14, 34, 54]],
+                           [[15, 35, 55],
+                            [16, 36, 56],
+                            [17, 37, 57],
+                            [18, 38, 58],
+                            [19, 39, 59]]]],
+
+                         [[[[60, 80, 100],
+                            [61, 81, 101],
+                            [62, 82, 102],
+                            [63, 83, 103],
+                            [64, 84, 104]],
+                           [[65, 85, 105],
+                            [66, 86, 106],
+                            [67, 87, 107],
+                            [68, 88, 108],
+                            [69, 89, 109]],
+                           [[70, 90, 110],
+                            [71, 91, 111],
+                            [72, 92, 112],
+                            [73, 93, 113],
+                            [74, 94, 114]],
+                           [[75, 95, 115],
+                            [76, 96, 116],
+                            [77, 97, 117],
+                            [78, 98, 118],
+                            [79, 99, 119]]]]]]).astype(np.float32)
+    assert (output[0].asnumpy() == expect0).all()
+    assert (output[1].asnumpy() == expect1).all()
+    assert (output[2].asnumpy() == expect2).all()
+    assert (output[3].asnumpy() == expect3).all()
+# test_transpose()