| 36 | } |
| 37 | |
| 38 | LstmStates::LstmStates( |
| 39 | const SmallVector<RefPtr> ptr, size_t hidden_size, size_t batch_size, |
| 40 | DType dtype) { |
| 41 | auto& h_ptr = ptr[0]; |
| 42 | auto& c_ptr = ptr[1]; |
| 43 | TensorLayout layout{{batch_size, hidden_size}, dtype}; |
| 44 | m_h = TensorND(layout, h_ptr); |
| 45 | m_c = TensorND(layout, c_ptr); |
| 46 | m_memory_size = layout.span().dist_byte(); |
| 47 | } |
| 48 | |
| 49 | TensorNDArray megdnn::arm_common::split_tensor( |
| 50 | _megdnn_tensor_in tensor, size_t nr_tensor, const TensorLayout& layout) { |