已合并
fix: support packed LSTM output shape transition #5761
ymdsall创建于 10 天前
fix: support packed LSTM output shape transition #5761
已合并
共 2 个文件变更+18-3
| @@ -789,6 +789,8 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> lstm( | |||
| 789 | double dropout, | 789 | double dropout, |
| 790 | bool train, | 790 | bool train, |
| 791 | bool bidirectional) { | 791 | bool bidirectional) { |
| 792 | + // The legacy LSTM.forward patch passes NPU batch_sizes and reshapes the output in Python. | ||
| 793 | + const bool should_reshape_output = batch_sizes.device().is_cpu(); | ||
| 792 | at::Tensor batch_sizes_cpu = batch_sizes.to("cpu"); | 794 | at::Tensor batch_sizes_cpu = batch_sizes.to("cpu"); |
| 793 | at::Tensor y; | 795 | at::Tensor y; |
| 794 | at::Tensor h; | 796 | at::Tensor h; |
| @@ -813,6 +815,9 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> lstm( | |||
| 813 | data, batch_sizes_cpu, hx, params, has_biases, num_layers, dropout, train, bidirectional); | 815 | data, batch_sizes_cpu, hx, params, has_biases, num_layers, dropout, train, bidirectional); |
| 814 | } | 816 | } |
| 815 | } | 817 | } |
| 818 | + if (should_reshape_output) { | ||
| 819 | + y = y.reshape({-1, y.size(-1)}); | ||
| 820 | + } | ||
| 816 | return std::tie(y, h, c); | 821 | return std::tie(y, h, c); |
| 817 | } | 822 | } |
| 818 | 823 | ||
| @@ -250,10 +250,20 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> lstm( | |||
| 250 | } | 250 | } |
| 251 | 251 | ||
| 252 | DO_COMPATIBILITY( | 252 | DO_COMPATIBILITY( |
| 253 | - aclnnLSTM, acl_op::lstm(data, batch_sizes, hx, params, has_biases, num_layers, dropout, train, bidirectional)); | 253 | + aclnnLSTM, |
| 254 | + acl_op::lstm(data, batch_sizes, hx, params, has_biases, num_layers, dropout, train, bidirectional)); | ||
| 255 | + // The legacy LSTM.forward patch passes NPU batch_sizes and reshapes the output in Python. | ||
| 256 | + const bool should_reshape_output = batch_sizes.device().is_cpu(); | ||
| 257 | + auto batch_sizes_npu = batch_sizes; | ||
| 258 | + if (batch_sizes_npu.device() != data.device()) { | ||
| 259 | + batch_sizes_npu = batch_sizes_npu.to(data.device()); | ||
| 260 | + } | ||
| 254 | auto output = at_npu::native::custom_ops::_lstm_npu( | 261 | auto output = at_npu::native::custom_ops::_lstm_npu( |
| 255 | - data, hx, params, has_biases, num_layers, dropout, train, bidirectional, false, batch_sizes); | 262 | + data, hx, params, has_biases, num_layers, dropout, train, bidirectional, false, batch_sizes_npu); |
| 263 | + const auto& output_y = std::get<0>(output); | ||
| 256 | return std::make_tuple( | 264 | return std::make_tuple( |
| 257 | - std::get<0>(output), std::get<1>(output), std::get<2>(output)); // 0 for output_y, 1 for output_h, 2 for output_c | 265 | + should_reshape_output ? output_y.reshape({-1, output_y.size(-1)}) : output_y, |
| 266 | + std::get<1>(output), | ||
| 267 | + std::get<2>(output)); // 0 for output_y, 1 for output_h, 2 for output_c | ||
| 258 | } | 268 | } |
| 259 | } // namespace op_api | 269 | } // namespace op_api |