| 8 | using namespace arm_common; |
| 9 | |
| 10 | LstmCellWeight::LstmCellWeight( |
| 11 | RefPtr weight_ptr, size_t hidden_size, size_t input_size, bool has_bias, |
| 12 | DType dtype) { |
| 13 | // weight_ih: [gate_hidden_size, input_size] |
| 14 | // weight_hh: [gate_hidden_size, hidden_size] |
| 15 | // bias_ih: [gate_hidden_size] |
| 16 | // bias_hh: [gate_hidden_size] |
| 17 | |
| 18 | size_t gate_hidden_size = 4 * hidden_size; |
| 19 | TensorLayout weight_ih_layout{{gate_hidden_size, input_size}, dtype}; |
| 20 | TensorLayout weight_hh_layout{{gate_hidden_size, hidden_size}, dtype}; |
| 21 | TensorLayout bias_layout{{gate_hidden_size}, dtype}; |
| 22 | m_weight_size = 0; |
| 23 | m_weight_ih = TensorND(weight_ih_layout, weight_ptr); |
| 24 | m_weight_size += weight_ih_layout.span().dist_byte(); |
| 25 | weight_ptr += weight_ih_layout.span().dist_byte(); |
| 26 | m_weight_hh = TensorND(weight_hh_layout, weight_ptr); |
| 27 | m_weight_size += weight_hh_layout.span().dist_byte(); |
| 28 | weight_ptr += weight_hh_layout.span().dist_byte(); |
| 29 | if (has_bias) { |
| 30 | m_bias_ih = TensorND(bias_layout, weight_ptr); |
| 31 | m_weight_size += bias_layout.span().dist_byte(); |
| 32 | weight_ptr += bias_layout.span().dist_byte(); |
| 33 | m_bias_hh = TensorND(bias_layout, weight_ptr); |
| 34 | m_weight_size += bias_layout.span().dist_byte(); |
| 35 | } |
| 36 | } |
| 37 | |
| 38 | LstmStates::LstmStates( |
| 39 | const SmallVector<RefPtr> ptr, size_t hidden_size, size_t batch_size, |