Star189
863
代码
代码
Issues196
Pull Requests327
项目讨论
Wiki
分析
Star189
863
已合并
add _lstm_npu and _lstm_npu_backward docs. #4382
shi-jiaxin9创建于 3月5日
add _lstm_npu and _lstm_npu_backward docs. #4382
已合并
shi-jiaxin9创建于 3月5日
从shi-jiaxin9/op-plugin_0305docsmaster合入到Ascend/op-pluginmaster
共 6 个文件变更+602-15
Mcodegen/templates/_op_plugin_docs.py+574-0
@@ -14016,3 +14016,577 @@ softmax_max_index, softmax_sum_index = torch_npu.npu_dense_lightning_indexer_sof
14016torch_npu.npu_dense_lightning_indexer_grad_kl_loss(q, k, q_index, k_index, weights, softmax_max, softmax_sum, softmax_max_index, softmax_sum_index, scale, query_rope=q_rope, key_rope=k_rope, actual_seq_qlen=actual_seq_qlen, actual_seq_klen=actual_seq_klen, layout=input_layout, sparse_mode=sparse_mode)14016torch_npu.npu_dense_lightning_indexer_grad_kl_loss(q, k, q_index, k_index, weights, softmax_max, softmax_sum, softmax_max_index, softmax_sum_index, scale, query_rope=q_rope, key_rope=k_rope, actual_seq_qlen=actual_seq_qlen, actual_seq_klen=actual_seq_klen, layout=input_layout, sparse_mode=sparse_mode)
14017"""14017"""
14018)14018)
14019+ 
14020+ 
14021+_add_torch_npu_docstr(
14022+ "_lstm_npu",
14023+ """
14024+接口原型:
14025+_lstm_npu(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)
14026+ 
14027+ 
14028+功能描述:
14029+LSTM（Long Short-Term Memory，长短时记忆）网络是一种特殊的循环神经网络（RNN）模型。进行LSTM网络计算，接收输入序列和初始状态，返回输出序列和最终状态。
14030+ 
14031+计算公式：
14032+
14033+ $$
14034+ \begin{aligned}
14035+ (1)\qquad f_t &=\sigma(W_f[h_{t-1}, x_t] + b_f) \\
14036+ (2)\qquad i_t &=\sigma(W_i[h_{t-1}, x_t] + b_i) \\
14037+ (3)\qquad o_t &=\sigma(W_o[h_{t-1}, x_t] + b_o) \\
14038+ (4)\qquad \tilde{c}_t &=tanh(W_c[h_{t-1}, x_t] + b_c) \\
14039+ (5)\qquad c_t &=f_t ⊙ c_{t-1} + i_t ⊙ \tilde{c}_t \\
14040+ (6)\qquad c_{o}^{t} &=tanh(c_t) \\
14041+ (7)\qquad h_t &=o_t ⊙ c_{o}^{t} \\
14042+ \end{aligned}
14043+ $$
14044+ 
14045+ - $x_t ∈ R^{d}$：LSTM单元的输入向量。
14046+ - $f_t ∈ (0, 1)^{h}$：遗忘门激活向量。
14047+ - $i_t ∈ (0, 1)^{h}$：输入门、更新门激活向量。
14048+ - $o_t ∈ (0, 1)^{h}$：输出门激活向量。
14049+ - $h_i ∈ (-1, 1)^{h}$：隐藏状态向量，也称为LSTM单元的输出向量。
14050+ - $\tilde{c}_t ∈ (-1, 1)^{h}$：cell输入激活向量。
14051+ - $c_t ∈ R^{h}$：cell状态向量。
14052+ - $W ∈ R^{h×d}，(U ∈ R^{h×h})∩(b ∈ R^{h})$：训练中需要学习的权重矩阵和偏置向量参数。
14053+ 
14054+ 
14055+参数说明:
14056+input (Tensor类型)：必选参数，LSTM单元的输入向量，数据格式支持ND，数据类型支持FLOAT16、FLOAT32，非连续Tensor。
14057+(1)若batch_sizes传入空指针：当batch_first=False时shape应为（time_step, batch_size, input_size）, 否则为（batch_size, time_step, input_size）。其中，batch_first表示batch是否在第一维；time_step表示时间维度；batch_size表示每个时刻需要处理的batch数量；input_size表示输入的特征数量。
14058+(2)若传入有效batch_sizes：shape应为(time_step * batch_size, input_size)，其内存排列与(time_step, batch_size, input_size)相同。
14059+hx (TensorList)：必选入参，表示LSTM运算中的初始hidden和cell状态列表，Device侧的TensorList,数据类型支持FLOAT16、FLOAT32，列表长度为2，列表中每个shape支持三维（D * num_layers, batch_size, hidden_size），若输入为空，则表示输入的初始hidden和cell状态为0。
14060+params (TensorList)：必选入参，表示LSTM运算中的权重和偏置张量列表,Device侧的TensorList，数据格式支持ND，数据类型支持FLOAT16、FLOAT32。
14061+ 列表长度为 2 * D * B * num_layers, 其中，num_layers对应参数numLayers，表示LSTM层数，bidirection为True时 D = 2， 否则 D = 1, has_biases为True时 B = 2, 否则 B = 1；
14062+ 其中bidirection为True， 且has_biases为True时，参数排布如下：[weight_ih_0, weight_hh_0, bias_ih_0, bias_hh_0, weight_ih_reverse_0, weight_hh_reverse_0, bias_ih_reverse_0, bias_hh_reverse_0],
14063+ 其中 weight_ih_0 表示第0层输入的权重参数，其shape为（4 * hidden_size, cur_input_size），其中cur_intput_size 表示LSTM每层计算时的输入的特征数量（首层为input_size, 后续层为hidden_size, 如果bidirection为True，则为2 * hidden_size）；
14064+ weight_hh_0 表示第0层隐藏层的权重参数，其shape为（4 * hidden_size, hidden_size）,bias_ih_0 表示第0层输入权重参数的偏置，其shape为（4 * hidden_size）,bias_hh_0 表示第0层隐藏层权重参数的偏置，其shape为（4 * hidden_size）。
14065+has_biases (bool)：必选入参，表示是否有biases。
14066+num_layers (int)：必选入参，表示LSTM层数。
14067+dropout (float)：表示随机掩码的概率。当前不支持该功能。
14068+train (bool)：必选入参，表示是否是训练模式。其中train = True时，在计算前向LSTM时会保存中间结果用于反向传播，train = False的时候，前向计算过程不保存中间结果。
14069+bidirectional(bool):必选入参，表示是否是双向。
14070+batch_first(bool):可选入参，表示输入数据格式是否是Batch在第一轴（B, T, H）。
14071+batch_sizes(Tensor类型):可选参数，表示每个时间步实际参与计算的有效Batch数。传入nullptr时，代表输入input为定长模式数据，否则为不定长模式。shape为(time_step,)，其中元素应按降序排列，元素值为正整数且最大不超过总Batch数量，且第一位元素值应与总Batch数量相等。
14072+ 
14073+输出说明:
14074+ 
14075+output (Tensor)：输出张量，表示LSTM运算中最后一层每个时间步的输出结果,数据类型支持FLOAT16、FLOAT32，非连续Tensor。当batch_first=False时shape支持三维（time_step, batch_size, D * hidden_size），否则支持三维（batch_size,time_step, D * hidden_size）。
14076+hy (Tensor)：输出张量，表示进行LSTM运算中每层最后一个时间步的隐藏层（公式（7）的输出）。shape支持三维（D * num_layers, batch_size, hidden_size）。数据格式支持ND，数据类型支持FLOAT16、FLOAT32。
14077+cy (Tensor)：输出张量，表示进行LSTM运算中每层最后一个时间步的Cell状态（公式（5）的输出）。shape支持三维（D * num_layers, batch_size, hidden_size）。数据格式支持ND，数据类型支持FLOAT16、FLOAT32。
14078+i_out(Tensor)：输出张量，表示LSTM运算中每层输入门的激活值（sigmoid输出，公式（2）的输出），数据类型支持FLOAT16、FLOAT32。原始列表长度为 D * num_layers，列表中每个元素为三维张量（time_step, batch_size, hidden_size）；
14079+最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, time_step, batch_size, hidden_size)；当train=False时，输出为空张量。
14080+j_out(Tensor)：输出张量，表示LSTM运算中每层的候选cell状态（tanh输出，公式（4）的输出），数据类型支持FLOAT16、FLOAT32。原始列表长度为 D * num_layers，列表中每个元素为三维张量（time_step, batch_size, hidden_size）；
14081+最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, time_step, batch_size, hidden_size)；当train=False时，输出为空张量。
14082+f_out(Tensor)：输出张量，表示LSTM运算中每层遗忘门的激活值（sigmoid输出），数据类型支持FLOAT16、FLOAT32。原始列表长度为 D * num_layers，列表中每个元素为三维张量（time_step, batch_size, hidden_size）；
14083+最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, time_step, batch_size, hidden_size)；当train=False时，输出为空张量。
14084+o_out(Tensor)：输出张量，表示LSTM运算中每层输出门的激活值（sigmoid输出，公式（3）的输出），数据类型支持FLOAT16、FLOAT32。原始列表长度为 D * num_layers，列表中每个元素为三维张量（time_step, batch_size, hidden_size）；
14085+最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, time_step, batch_size, hidden_size)；当train=False时，输出为空张量。
14086+h_out(Tensor)：输出张量，表示LSTM运算中每层的隐藏层（公式（7）的输出），数据类型支持FLOAT16、FLOAT32。train=True时：原始列表长度为 D * num_layers，列表中每个元素为三维张量（time_step, batch_size, hidden_size）；最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, time_step, batch_size, hidden_size)；
14087+train=False时：原始列表长度为 D * num_layers，列表中每个元素为二维张量（batch_size, hidden_size）；最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, batch_size, hidden_size)；
14088+c_out(Tensor)：输出张量，表示LSTM运算中每层的最终Cell状态（公式（5）的输出），数据类型支持FLOAT16、FLOAT32。train=True时：原始列表长度为 D * num_layers，列表中每个元素为三维张量（time_step, batch_size, hidden_size）；最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, time_step, batch_size, hidden_size)；
14089+train=False时：原始列表长度为 D * num_layers，列表中每个元素为二维张量（batch_size, hidden_size）；最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, batch_size, hidden_size)；
14090+tanh_c_out(Tensor)：输出张量，表示LSTM运算中每层最终cell状态经过tanh激活函数后的输出（公式（6）的输出），数据类型支持FLOAT16、FLOAT32。原始列表长度为 D * num_layers，列表中每个元素为三维张量（time_step, batch_size, hidden_size）；最终输出为将列表张量在维度0上堆叠后的张量，shape为 (D * num_layers, time_step, batch_size, hidden_size)；当train=False时，输出为空张量。
14091+ 
14092+ 
14093+约束说明
14094+确定性计算：_lstm_npu默认支持确定性实现。
14095+ 
14096+支持版本:
14097+PyTorch 2.6及更高版本
14098+ 
14099+支持的型号:
14100+Atlas A3 训练系列产品/Atlas A3 推理系列产品
14101+Atlas A2 训练系列产品/Atlas A2 推理系列产品
14102+Atlas 推理系列产品
14103+Atlas 训练系列产品
14104+ 
14105+调用示例:
14106+import torch
14107+import torch_npu
14108+ 
14109+torch.npu.set_device(0)
14110+device = torch.device("npu:0")
14111+dtype = torch.float32
14112+ 
14113+input_tensor = torch.randn((2, 3, 8), dtype=dtype, device=device)
14114+h0 = torch.randn((2, 2, 4), dtype=dtype, device=device)
14115+c0 = torch.randn((2, 2, 4), dtype=dtype, device=device)
14116+hx = [h0, c0]
14117+hidden_size = 4
14118+gate_size = 4 * hidden_size
14119+weight_ih = torch.randn((gate_size, 8), dtype=dtype, device=device)
14120+weight_hh = torch.randn((gate_size, 4), dtype=dtype, device=device)
14121+params = [weight_ih, weight_hh, weight_hh, weight_hh]
14122+c0 = torch.randn((1, 2, 16), dtype=dtype, device=device)
14123+batch_sizes = torch.randn((4*hidden_size), dtype=dtype, device=device)
14124+input_tensor = input_tensor.to(device, dtype=dtype)
14125+hx = [
14126+ t.to(device=device, dtype=dtype)
14127+ for t in hx
14128+]
14129+params = [
14130+ p.to(device=device, dtype=dtype)
14131+ for p in params
14132+]
14133+ 
14134+ 
14135+out, out_h, out_c, _, _, _, _, hn_list, cn_list, _ = torch_npu._lstm_npu(
14136+ input_tensor, # 1. input
14137+ hx, # 2. hx
14138+ params, # 3. params
14139+ False, # 4. has_biases
14140+ 2, # 5. num_layers
14141+ 0.0, # 6. dropout
14142+ False, # 7. train
14143+ False, # 8. bidirectional
14144+ batch_first=True, # 9. batch_first（位置参数，非关键字）
14145+ batch_sizes=None # 10. batch_sizes
14146+)
14147+ 
14148+"""
14149+)
14150+ 
14151+ 
14152+_add_torch_npu_docstr(
14153+ "_lstm_npu_backward",
14154+ """
14155+接口原型:
14156+_lstm_npu_backward(Tensor grad_y, Tensor grad_hy, Tensor grad_cy, Tensor input, Tensor[] hx, Tensor[] params, Tensor i, Tensor j, Tensor f, Tensor o, Tensor h, Tensor c, Tensor tanhc, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor[], Tensor[])
14157+ 
14158+ 
14159+功能描述:
14160+LSTM的反向传播，计算正向输入input、权重params、初始状态hx的梯度。
14161+ 
14162+计算公式：
14163+
14164+ <details>
14165+ <summary> 单层LSTM反向传播计算公式</summary>
14166+ 
14167+ | 组件 | 公式 |
14168+ |:---|:---|
14169+ | 输入拼接 | $\mathbf{z}_t = \begin{bmatrix} \mathbf{h}_{t-1} \\ \mathbf{x}_t \end{bmatrix}$ |
14170+ | 遗忘门 | $\mathbf{f}_t = \sigma(\mathbf{W}_f \mathbf{z}_t + \mathbf{b}_f)$ |
14171+ | 输入门 | $\mathbf{i}_t = \sigma(\mathbf{W}_i \mathbf{z}_t + \mathbf{b}_i)$ |
14172+ | 候选状态 | $\mathbf{g}_t = \tanh(\mathbf{W}_g \mathbf{z}_t + \mathbf{b}_c)$ |
14173+ | 输出门 | $\mathbf{o}_t = \sigma(\mathbf{W}_o \mathbf{z}_t + \mathbf{b}_o)$ |
14174+ | 细胞状态 | $\mathbf{c}_t = \mathbf{f}_t \odot \mathbf{c}_{t-1} + \mathbf{i}_t \odot \mathbf{g}_t$ |
14175+ | 隐藏状态 | $\mathbf{h}_t = \mathbf{o}_t \odot \tanh(\mathbf{c}_t)$ |
14176+ 
14177+ 其中：
14178+ 
14179+ - $\sigma$ 是 sigmoid 函数
14180+ - $\odot$ 表示逐元素乘法 (Hadamard product)
14181+ - $W_*$ 是可学习的权重矩阵
14182+ - $b_*$ 是可学习的偏置项
14183+ </details>
14184+ 
14185+ <details>
14186+ 
14187+ <summary> 反向传播变量定义</summary>
14188+ 
14189+ - 总损失：$L = \sum_{t=1}^{T} L_t$
14190+ - 隐藏状态梯度：$\delta\mathbf{h}_t = \frac{\partial L}{\partial \mathbf{h}_t}$
14191+ - 细胞状态梯度：$\delta\mathbf{c}_t = \frac{\partial L}{\partial \mathbf{c}_t}$
14192+ </details>
14193+ 
14194+ <details>
14195+ 
14196+ <summary> 反向传播算法（时间步 t -> t-1）</summary>
14197+ 
14198+ - **初始化**
14199+ 
14200+ $$
14201+ \delta\mathbf{h}_{T} = \mathbf{0}, \quad \delta\mathbf{c}_{T} = \mathbf{0}, \quad \mathbf{f}_{T} = \mathbf{0}
14202+ $$
14203+ 
14204+ - **循环 $t = T - 1$ 到 $0$**
14205+ 
14206+ 1.**当前隐藏状态梯度**
14207+ 
14208+ $$
14209+ \delta\mathbf{h}_t = \frac{\partial L_t}{\partial \mathbf{h}_t} + \delta\mathbf{h}_{\text{next}}
14210+ $$
14211+ 
14212+ 2.**当前细胞状态梯度**
14213+ 
14214+ $$
14215+ \delta\mathbf{c}_t = \delta\mathbf{h}_t \odot \mathbf{o}_t \odot (1 - \tanh^2(\mathbf{c}_t)) + \delta\mathbf{c}_{\text{next}} \odot \mathbf{f}_{\text{next}}
14216+ $$
14217+ 
14218+ 3.**门控梯度计算**
14219+ 
14220+ $$
14221+ \delta\mathbf{o}_t = \delta\mathbf{h}_t \odot \tanh(\mathbf{c}_t) \odot \mathbf{o}_t \odot (1 - \mathbf{o}_t)
14222+ $$
14223+ 
14224+ $$
14225+ \delta\mathbf{g}_t = \delta\mathbf{c}_t \odot \mathbf{i}_t \odot (1 - \mathbf{g}_t^2)
14226+ $$
14227+ 
14228+ $$
14229+ \delta\mathbf{i}_t = \delta\mathbf{c}_t \odot \mathbf{g}_t \odot \mathbf{i}_t \odot (1 - \mathbf{i}_t)
14230+ $$
14231+ 
14232+ $$
14233+ \delta\mathbf{f}_t = \delta\mathbf{c}_t \odot \mathbf{c}_{t-1} \odot \mathbf{f}_t \odot (1 - \mathbf{f}_t)
14234+ $$
14235+ 
14236+ 4.**参数梯度累加**
14237+ 
14238+ $$
14239+ \frac{\partial L}{\partial \mathbf{W}_f} \mathrel{+}= \delta\mathbf{f}_t \mathbf{z}_t^\top
14240+ $$
14241+ 
14242+ $$
14243+ \frac{\partial L}{\partial \mathbf{b}_f} \mathrel{+}= \delta\mathbf{f}_t
14244+ $$
14245+ 
14246+ $$
14247+ \frac{\partial L}{\partial \mathbf{W}_i} \mathrel{+}= \delta\mathbf{i}_t \mathbf{z}_t^\top
14248+ $$
14249+ 
14250+ $$
14251+ \frac{\partial L}{\partial \mathbf{b}_i} \mathrel{+}= \delta\mathbf{i}_t
14252+ $$
14253+ 
14254+ $$
14255+ \frac{\partial L}{\partial \mathbf{W}_g} \mathrel{+}= \delta\mathbf{g}_t \mathbf{z}_t^\top
14256+ $$
14257+ 
14258+ $$
14259+ \frac{\partial L}{\partial \mathbf{b}_g} \mathrel{+}= \delta\mathbf{g}_t
14260+ $$
14261+ 
14262+ $$
14263+ \frac{\partial L}{\partial \mathbf{W}_o} \mathrel{+}= \delta\mathbf{o}_t \mathbf{z}_t^\top
14264+ $$
14265+ 
14266+ $$
14267+ \frac{\partial L}{\partial \mathbf{b}_o} \mathrel{+}= \delta\mathbf{o}_t
14268+ $$
14269+ 
14270+ 5.**传播到前一时刻**
14271+ 
14272+ $$
14273+ \delta\mathbf{z}_t = \mathbf{W}_f^\top \delta\mathbf{f}_t + \mathbf{W}_i^\top \delta\mathbf{i}_t + \mathbf{W}_g^\top \delta\mathbf{g}_t + \mathbf{W}_o^\top \delta\mathbf{o}_t
14274+ $$
14275+ 
14276+ $$
14277+ \delta\mathbf{h}_{\text{prev}} = \delta\mathbf{z}_t[1:\dim(\mathbf{h}_{t-1})]
14278+ $$
14279+ 
14280+ $$
14281+ \delta\mathbf{c}_{\text{prev}} = \delta\mathbf{c}_t \odot \mathbf{f}_t
14282+ $$
14283+ 
14284+ 6.**更新传播变量**
14285+ 
14286+ $$
14287+ \delta\mathbf{h}_{\text{next}} \leftarrow \delta\mathbf{h}_{\text{prev}}
14288+ $$
14289+ 
14290+ $$
14291+ \delta\mathbf{c}_{\text{next}} \leftarrow \delta\mathbf{c}_{\text{prev}}
14292+ $$
14293+ 
14294+ $$
14295+ \mathbf{f}_{\text{next}} \leftarrow \mathbf{f}_t
14296+ $$
14297+ 
14298+ </details>
14299+ 
14300+ <details>
14301+ <summary> 梯度计算原理</summary>
14302+ 
14303+ - **细胞状态梯度推导**
14304+ 
14305+ $$
14306+ \delta\mathbf{c}_t = \frac{\partial L}{\partial \mathbf{h}_t} \frac{\partial \mathbf{h}_t}{\partial \mathbf{c}_t} + \frac{\partial L}{\partial \mathbf{c}_{t+1}} \frac{\partial \mathbf{c}_{t+1}}{\partial \mathbf{c}_t}
14307+ $$
14308+ 
14309+ 其中：
14310+ 
14311+ $$
14312+ \frac{\partial \mathbf{h}_t}{\partial \mathbf{c}_t} = \mathbf{o}_t \odot (1 - \tanh^2(\mathbf{c}_t))
14313+ $$
14314+ 
14315+ $$
14316+ \frac{\partial \mathbf{c}_{t+1}}{\partial \mathbf{c}_t} = \mathbf{f}_{t+1}
14317+ $$
14318+ 
14319+ - **遗忘门梯度推导**
14320+ 
14321+ $$
14322+ \delta\mathbf{f}_t = \frac{\partial L}{\partial \mathbf{a}_f^t} = \delta\mathbf{c}_t \odot \mathbf{c}_{t-1} \odot \mathbf{f}_t \odot (1 - \mathbf{f}_t)
14323+ $$
14324+ 
14325+ - **参数梯度推导**
14326+ 
14327+ $$
14328+ \frac{\partial L}{\partial \mathbf{W}_f} = \sum_{t=1}^{T} \delta\mathbf{f}_t \mathbf{z}_t^\top
14329+ $$
14330+ 
14331+ - **LSTM 梯度流动特性**
14332+ 
14333+ **长程依赖处理**
14334+ 
14335+ $$
14336+ \frac{\partial \mathbf{c}_T}{\partial \mathbf{c}_1} = \prod_{k=2}^{T} \mathbf{f}_k \quad \text{(对角矩阵)}
14337+ $$
14338+ 
14339+ </details>
14340+ 
14341+ <details>
14342+ <summary> 多层LSTMBackward反向传播</summary>
14343+ 在多层LSTM网络中，层与层之间的梯度传播仅关注隐藏状态的传递（忽略单层内部细节，如门控机制或单元状态）。设：
14344+ 
14345+ - $\mathbf{h}^{(l)}$：第 $l$ 层的隐藏状态（$l = 1, 2, \dots, L$，其中 $L$ 为总层数）
14346+ - $L$：损失函数
14347+ - $\frac{\partial L}{\partial \mathbf{h}^{(l)}}$：损失函数对第 $l$ 层隐藏状态的梯度
14348+ 
14349+ **核心传播公式**
14350+ 
14351+ 梯度从顶层（$l = L$）向底层（$l = 1$）传播，层间关系由链式法则给出：
14352+ 
14353+ $$
14354+ \frac{\partial L}{\partial \mathbf{h}^{(l-1)}} = \frac{\partial L}{\partial \mathbf{h}^{(l)}} \cdot \frac{\partial \mathbf{h}^{(l)}}{\partial \mathbf{h}^{(l-1)}}
14355+ $$
14356+ 
14357+ 其中：
14358+ 
14359+ - $\frac{\partial L}{\partial \mathbf{h}^{(l)}}$：当前层 $l$ 的梯度（已由上一层反向传播得到）
14360+ - $\frac{\partial \mathbf{h}^{(l)}}{\partial \mathbf{h}^{(l-1)}}$：第 $l$ 层隐藏状态对第 $l-1$ 层隐藏状态的雅可比矩阵
14361+ - $\cdot$：矩阵乘法（梯度传播本质为向量-矩阵乘法）
14362+ 
14363+ 即每层的输出的梯度dx为上一层输入的梯度dy。
14364+ </details>
14365+ 
14366+ 
14367+参数说明:
14368+grad_y (Tensor类型)：必选参数，LSTM正向最后一层输出hidden的梯度。对应公式中的∂L/∂h^(l)。双向时数据沿最后一维按前后向排布。数据类型与input一致。数据格式支持ND，数据类型支持FLOAT16、FLOAT32。非连续Tensor。
14369+若传入有效batchSizesOptional，shape为[time_step * batch_size, hidden_size * D];若传入空指针batchSizesOptional，shape为[time_step, batch_size, hidden_size * D] 或 [batch_size, time_step, hidden_size * D]。
14370+grad_hy (Tensor类型)：必选入参，LSTM正向每层输出hidden在T时刻从下一个时间步传来的梯度。对应δh_next。多层双向时数据沿第0维按先双向后逐层排布。数据类型与input一致。数据格式支持ND，数据类型支持FLOAT16、FLOAT32。非连续Tensor,shape为[numLayers * D, batch_size, hidden_size]。
14371+grad_cy (Tensor类型)：必选入参，LSTM每层输出cell在T时刻从下一个时间步传来的梯度。对应δc_next。多层双向时数据沿第0维按先双向后逐层排布。数据类型与input一致。数据格式支持ND，数据类型支持FLOAT16、FLOAT32。非连续Tensor。shape为[numLayers * D, batch_size, hidden_size]。
14372+input (Tensor类型)：必选参数，LSTM单元的输入向量，数据格式支持ND，数据类型支持FLOAT16、FLOAT32，非连续Tensor。若传入有效batchSizesOptional，shape为[time_step * batch_size, input_size];
14373+若传入空指针batchSizesOptional，shape为[time_step, batch_size, input_size] 或 [batch_size, time_step, input_size],batch_size表示序列组数；time_step表示时间维度；input_size表示输入的特征数量。
14374+hx (TensorList)：必选入参，LSTM每层的初始hidden和cell状态。对应0时刻的h(t-1)与c(t-1)。列表长度为2，包含h_0和c_0。多层双向时每个tensor数据沿第0维按先双向后逐层排布。数据类型与input一致。数据格式支持ND，数据类型支持FLOAT16、FLOAT32。非连续Tensor。列表内每个tensor shape为[D * num_layers, batch_size, hidden_size]。
14375+params (TensorList)：必选入参，LSTM每层的权重和偏置张量列表，对应公式中的w与b。bidirection为True时 `D = 2`，否则 `D = 1`，hasBiases为True时 `B = 2`，否则 `B = 1`。列表长度为 D * B * num_layers。当bidirection和hasBias均为True时排布为：[weight_ih_0, weight_hh_0, bias_ih_0, bias_hh_0, weight_ih_reverse_0, weight_hh_reverse_0, bias_ih_reverse_0, bias_hh_reverse_0]。hasBias为False时无bias项；bidirection为False时无reverse项。多层时逐层排布。数据类型与input一致。
14376+数据格式支持ND，数据类型支持FLOAT16、FLOAT32。非连续Tensor。weight_ih: [4*hidden_size, cur_input_size];weight_hh: [4*hidden_size, hidden_size];bias_ih: [4*hidden_size];bias_hh: [4*hidden_size]。
14377+i (Tensor类型)：必选入参，LSTM正向中每层输出的输入门的激活值。对应公式中的i。
14378+ 原始输入形态：为Tensor列表，列表长度为 D * num_layers；多层双向场景下，列表内Tensor按「先双向、后多层」的顺序排布；
14379+ 列表中单个Tensor特征：非连续Tensor，数据类型与input一致，支持FLOAT16、FLOAT32；shape为[time_step, batch_size, hidden_size]；
14380+ 代码处理后形态：将列表内Tensor在维度0上堆叠为单个大Tensor，最终shape为[D * num_layers, time_step, batch_size, hidden_size]；若列表为空则输出空Tensor。
14381+j (Tensor类型)：必选入参，LSTM正向中每层输出的候选cell状态的激活值。对应公式中的g。
14382+ 原始输入形态：为Tensor列表，列表长度为 D * num_layers；多层双向场景下，列表内Tensor按「先双向、后多层」的顺序排布；
14383+ 列表中单个Tensor特征：非连续Tensor，数据类型与input一致，支持FLOAT16、FLOAT32；shape为[time_step, batch_size, hidden_size]；
14384+ 代码处理后形态：将列表内Tensor在维度0上堆叠为单个大Tensor，最终shape为[D * num_layers, time_step, batch_size, hidden_size]；若列表为空则输出空Tensor。
14385+f (Tensor类型)：必选入参，LSTM正向中每层遗忘门的激活值。对应公式中的f。
14386+ 原始输入形态：为Tensor列表，列表长度为 D * num_layers；多层双向场景下，列表内Tensor按「先双向、后多层」的顺序排布；
14387+ 列表中单个Tensor特征：非连续Tensor，数据类型与input一致，支持FLOAT16、FLOAT32；shape为[time_step, batch_size, hidden_size]；
14388+ 代码处理后形态：将列表内Tensor在维度0上堆叠为单个大Tensor，最终shape为[D * num_layers, time_step, batch_size, hidden_size]；若列表为空则输出空Tensor。
14389+o (Tensor类型)：必选入参，LSTM正向中每层输出门的激活值。对应公式中的o。
14390+ 原始输入形态：为Tensor列表，列表长度为 D * num_layers；多层双向场景下，列表内Tensor按「先双向、后多层」的顺序排布；
14391+ 列表中单个Tensor特征：非连续Tensor，数据类型与input一致，支持FLOAT16、FLOAT32；shape为[time_step, batch_size, hidden_size]；
14392+ 代码处理后形态：将列表内Tensor在维度0上堆叠为单个大Tensor，最终shape为[D * num_layers, time_step, batch_size, hidden_size]；若列表为空则输出空Tensor。
14393+h (Tensor类型)：必选入参，LSTM正向中每层的隐藏hidden状态。对应公式中的h。
14394+ 原始输入形态：为Tensor列表，列表长度为 D * num_layers；多层双向场景下，列表内Tensor按「先双向、后多层」的顺序排布；
14395+ 列表中单个Tensor特征：非连续Tensor，数据类型与input一致，支持FLOAT16、FLOAT32；shape为[time_step, batch_size, hidden_size]；
14396+ 代码处理后形态：将列表内Tensor在维度0上堆叠为单个大Tensor，最终shape为[D * num_layers, time_step, batch_size, hidden_size]；若列表为空则输出空Tensor。
14397+c (Tensor类型)：必选入参，LSTM正向中每层的最终cell状态。对应公式中的c。
14398+ 原始输入形态：为Tensor列表，列表长度为 D * num_layers；多层双向场景下，列表内Tensor按「先双向、后多层」的顺序排布；
14399+ 列表中单个Tensor特征：非连续Tensor，数据类型与input一致，支持FLOAT16、FLOAT32；shape为[time_step, batch_size, hidden_size]；
14400+ 代码处理后形态：将列表内Tensor在维度0上堆叠为单个大Tensor，最终shape为[D * num_layers, time_step, batch_size, hidden_size]；若列表为空则输出空Tensor。
14401+tanhc (Tensor类型)：必选入参，LSTM正向中每层最终cell状态经过tanh激活函数后的输出。对应公式中的tanh(c)。
14402+ 原始输入形态：为Tensor列表，列表长度为 D * num_layers；多层双向场景下，列表内Tensor按「先双向、后多层」的顺序排布；
14403+ 列表中单个Tensor特征：非连续Tensor，数据类型与input一致，支持FLOAT16、FLOAT32；shape为[time_step, batch_size, hidden_size]；
14404+ 代码处理后形态：将列表内Tensor在维度0上堆叠为单个大Tensor，最终shape为[D * num_layers, time_step, batch_size, hidden_size]；若列表为空则输出空Tensor。
14405+has_biases(bool):必选入参，表示是否有偏置b。
14406+num_layers(int):必选入参，表示LSTM层数。值大于0。
14407+dropout (float)：表示随机掩码的概率。当前不支持该功能。
14408+train (bool)：必选入参，表示是否是训练模式。其中train = True时，在计算前向LSTM时会保存中间结果用于反向传播，train = False的时候，前向计算过程不保存中间结果。
14409+bidirectional(bool):必选入参，表示是否是双向。bidirectional为true时D为2，bidirectional为false时D为1。
14410+batch_first(bool):可选入参，表示输入数据格式是否是Batch在第一轴（B, T, H）。
14411+batch_sizes(Tensor类型):可选参数，变长LSTM输入序列各个时刻的有效序列batch数。变长序列时支持。shape为[time_step]。
14412+ 
14413+ 
14414+输出说明:
14415+dx_out (Tensor)：输出张量，输入input上的梯度，对应公式中的δx。shape与input一致。数据类型与input一致。shape为[time_step, batch_size, input_size] 或 [batch_size, time_step, input_size]，数据格式支持ND，数据类型支持FLOAT16、FLOAT32。非连续Tensor。
14416+out_hx_prev (TensorList)：输出张量列表，由张量dh_prev_out和dc_prev_out拼接结果，列表长度为2。dh_prev_out是LSTM每层初始hidden的梯度，对应t=0时的δh_prev。数据类型与input一致。shape为[D * num_layers, batch_size, hidden_size]，数据格式支持ND，数据类型支持FLOAT16、FLOAT32。非连续Tensor。dc_prev_out是多层双向时数据沿第0维按先双向后逐层排布。数据类型与input一致。shape为[D * num_layers, batch_size, hidden_size]，数据格式支持ND，数据类型支持FLOAT16、FLOAT32。非连续Tensor。
14417+dparams_out(TensorList)：输出张量列表，权重和偏置的梯度张量列表。对应公式中的δw和δb。列表长度为 D * B * num_layers。排布与输入params一致。数据类型与input一致。dweight_ih: [4*hidden_size, cur_input_size]，dweight_hh: [4*hidden_size, hidden_size]，dbias: [4*hidden_size]
14418+ 
14419+ 
14420+约束说明
14421+确定性计算：_lstm_npu_backward默认支持确定性实现。
14422+ 
14423+支持版本:
14424+PyTorch 2.6及更高版本
14425+ 
14426+支持的型号:
14427+Atlas A3 训练系列产品/Atlas A3 推理系列产品
14428+Atlas A2 训练系列产品/Atlas A2 推理系列产品
14429+Atlas 推理系列产品
14430+Atlas 训练系列产品
14431+ 
14432+调用示例:
14433+ 
14434+import torch
14435+import torch_npu
14436+ 
14437+def _lstm_backward_npu(
14438+ grad_y: torch.Tensor,
14439+ grad_hy: torch.Tensor,
14440+ grad_cy: torch.Tensor,
14441+ input: torch.Tensor,
14442+ hx: List[torch.Tensor],
14443+ params: List[torch.Tensor],
14444+ i: torch.Tensor,
14445+ j: torch.Tensor,
14446+ f: torch.Tensor,
14447+ o: torch.Tensor,
14448+ h: torch.Tensor,
14449+ c: torch.Tensor,
14450+ tanhc: torch.Tensor,
14451+ has_biases: bool = True,
14452+ num_layers: int = 1,
14453+ dropout: float = 0.0,
14454+ train: bool = True,
14455+ bidirectional: bool = False,
14456+ batch_first: Optional[bool] = None,
14457+ batch_sizes: Optional[torch.Tensor] = None
14458+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[torch.Tensor]]:
14459+ 
14460+ # 1. 强制所有张量转到NPU设备，且确保类型为float32（解决double dtype警告）
14461+ npu_device = torch.device('npu:0')
14462+ dtype = torch.float32 # NPU不支持double，统一用float32
14463+
14464+ # 类型+设备转换
14465+ def to_npu_tensor(t: torch.Tensor) -> torch.Tensor:
14466+ if t is None:
14467+ return t
14468+ return t.to(device=npu_device, dtype=dtype)
14469+
14470+ grad_y = to_npu_tensor(grad_y)
14471+ grad_hy = to_npu_tensor(grad_hy)
14472+ grad_cy = to_npu_tensor(grad_cy)
14473+ input = to_npu_tensor(input)
14474+ i = to_npu_tensor(i)
14475+ j = to_npu_tensor(j)
14476+ f = to_npu_tensor(f)
14477+ o = to_npu_tensor(o)
14478+ h = to_npu_tensor(h)
14479+ c = to_npu_tensor(c)
14480+ tanhc = to_npu_tensor(tanhc)
14481+
14482+ hx = [to_npu_tensor(t) for t in hx]
14483+ params = [to_npu_tensor(p) for p in params]
14484+ if batch_sizes is not None:
14485+ batch_sizes = to_npu_tensor(batch_sizes)
14486+
14487+ # 2. 处理batch_first默认值（接口声明中默认False）
14488+ if batch_first is None:
14489+ batch_first = False
14490+
14491+ # 3. 核心调用：严格匹配接口声明的入参类型
14492+ # 关键：i/j/f/o/h/c/tanhc直接传原始Tensor，不做chunk拆分！
14493+ output = torch_npu._lstm_backward_npu(
14494+ grad_y, # 0: Tensor grad_y
14495+ grad_hy, # 1: Tensor grad_hy
14496+ grad_cy, # 2: Tensor grad_cy
14497+ input, # 3: Tensor input
14498+ hx, # 4: Tensor[] hx (Python列表)
14499+ params, # 5: Tensor[] params (Python列表)
14500+ i, # 6: Tensor i （核心修复：传单个Tensor，非tuple）
14501+ j, # 7: Tensor j
14502+ f, # 8: Tensor f
14503+ o, # 9: Tensor o
14504+ h, # 10: Tensor h
14505+ c, # 11: Tensor c
14506+ tanhc, # 12: Tensor tanhc
14507+ has_biases, # 13: bool has_biases
14508+ num_layers, # 14: int num_layers
14509+ dropout, # 15: float dropout
14510+ train, # 16: bool train
14511+ bidirectional, # 17: bool bidirectional
14512+ batch_first=batch_first, # 18: bool? batch_first（关键字参数，接口声明要求）
14513+ batch_sizes=batch_sizes # 19: Tensor? batch_sizes（关键字参数）
14514+ )
14515+
14516+ # 4. 解析输出
14517+ input_grad = output[0] # 输入序列梯度
14518+ hx_prev_grad = output[1] # 初始hx梯度
14519+ param_grads_list = output[2] # 参数梯度列表
14520+ 
14521+ return input_grad, hx_prev_grad, param_grads_list
14522+ 
14523+if __name__ == "__main__":
14524+ # 1. 初始化测试参数（严格匹配接口要求）
14525+ seq_len = 8
14526+ batch_size = 4
14527+ input_size = 6
14528+ hidden_size = 4
14529+ num_layers = 1
14530+ bidirectional = False
14531+ num_dir = 2 if bidirectional else 1
14532+
14533+ # 2. 构造NPU张量（float32，避免double dtype警告）
14534+ device = torch.device('npu:0')
14535+ dtype = torch.float32
14536+
14537+ # 输入序列 [seq_len, batch, input_size]
14538+ input = torch.randn(seq_len, batch_size, input_size, device=device, dtype=dtype)
14539+ # 输出梯度 [seq_len, batch, hidden_size*num_dir]
14540+ grad_y = torch.randn(seq_len, batch_size, hidden_size * num_dir, device=device, dtype=dtype)
14541+ # 最后一层h/c梯度 [num_layers*num_dir, batch, hidden_size]
14542+ grad_hy = torch.randn(num_layers * num_dir, batch_size, hidden_size, device=device, dtype=dtype)
14543+ grad_cy = torch.randn(num_layers * num_dir, batch_size, hidden_size, device=device, dtype=dtype)
14544+
14545+ # 初始h0/c0 [num_layers*num_dir, batch, hidden_size]
14546+ h0 = torch.randn(num_layers * num_dir, batch_size, hidden_size, device=device, dtype=dtype)
14547+ c0 = torch.randn(num_layers * num_dir, batch_size, hidden_size, device=device, dtype=dtype)
14548+ hx = [h0, c0]
14549+
14550+ # LSTM参数（w_ih, w_hh, b_ih, b_hh）
14551+ w_ih = torch.randn(4 * hidden_size, input_size, device=device, dtype=dtype)
14552+ w_hh = torch.randn(4 * hidden_size, hidden_size, device=device, dtype=dtype)
14553+ b_ih = torch.randn(4 * hidden_size, device=device, dtype=dtype)
14554+ b_hh = torch.randn(4 * hidden_size, device=device, dtype=dtype)
14555+ params = [w_ih, w_hh, b_ih, b_hh]
14556+
14557+ # 前向中间结果（核心：不拆分，直接传原始Tensor）
14558+ # 形状：[seq_len, batch, num_layers*num_dir*hidden_size]
14559+ i = torch.randn(seq_len, batch_size, num_layers * num_dir * hidden_size, device=device, dtype=dtype)
14560+ j = torch.randn_like(i)
14561+ f = torch.randn_like(i)
14562+ o = torch.randn_like(i)
14563+ h = torch.randn_like(i)
14564+ c = torch.randn_like(i)
14565+ tanhc = torch.tanh(c).to(dtype=dtype) # 确保dtype一致
14566+
14567+ 
14568+ input_grad, h_prev_grad, c_prev_grad, param_grads = _lstm_backward_npu(
14569+ grad_y=grad_y,
14570+ grad_hy=grad_hy,
14571+ grad_cy=grad_cy,
14572+ input=input,
14573+ hx=hx,
14574+ params=params,
14575+ i=i, # 直接传单个Tensor，无chunk
14576+ j=j,
14577+ f=f,
14578+ o=o,
14579+ h=h,
14580+ c=c,
14581+ tanhc=tanhc,
14582+ has_biases=True,
14583+ num_layers=num_layers,
14584+ dropout=0.0,
14585+ train=True,
14586+ bidirectional=bidirectional,
14587+ batch_first=False,
14588+ batch_sizes=None
14589+ )
14590+ 
14591+"""
14592+)
Mop_plugin/config/derivatives.yaml+5-0
@@ -183,6 +183,11 @@ backward:
183 input, weight, bias, h, c: npu_lstm_data_backward(grads[0], grads[1], grads[2], input, batch_sizes, weight, bias, h, c, result0, result1, result2, result3, result4, result5, result6, result7, direction)183 input, weight, bias, h, c: npu_lstm_data_backward(grads[0], grads[1], grads[2], input, batch_sizes, weight, bias, h, c, result0, result1, result2, result3, result4, result5, result6, result7, direction)
184 version: [v2.1, newest]184 version: [v2.1, newest]
185 185 
186+- name: _lstm_npu(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)
187+ output_differentiability: [true, true, true, false, false, false, false, false, false, false]
188+ input, hx, params: _lstm_npu_backward(grads[0], grads[1], grads[2], input, hx, params, result3, result4, result5, result6, result7, result8, result9, has_biases, num_layers, dropout, train, bidirectional, batch_first, batch_sizes)
189+ version: all_version
190+ 
186- name: npu_max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)191- name: npu_max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices)
187 self: npu_max_backward(grad, dim, indices, self.sym_sizes(), keepdim)192 self: npu_max_backward(grad, dim, indices, self.sym_sizes(), keepdim)
188 version: [v2.1, newest]193 version: [v2.1, newest]
Mop_plugin/config/op_plugin_functions.yaml+1-1
@@ -6195,7 +6195,7 @@ custom:
6195 op_api: all_version6195 op_api: all_version
6196 tags: nondeterministic_seeded6196 tags: nondeterministic_seeded
6197 6197 
6198- - func: _lstm_npu_backward(Tensor grad_y, Tensor grad_hy, Tensor grad_cy, Tensor input, Tensor[] hx, Tensor[] params, Tensor i, Tensor j, Tensor f, Tensor o, Tensor h, Tensor c, Tensor tanhc, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor, Tensor, Tensor[])6198+ - func: _lstm_npu_backward(Tensor grad_y, Tensor grad_hy, Tensor grad_cy, Tensor input, Tensor[] hx, Tensor[] params, Tensor i, Tensor j, Tensor f, Tensor o, Tensor h, Tensor c, Tensor tanhc, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor[], Tensor[])
6199 op_api: all_version6199 op_api: all_version
6200 6200 
6201 - func: npu_masked_fill_range(Tensor self, Tensor start, Tensor end, Tensor value, int axis=-1) -> Tensor6201 - func: npu_masked_fill_range(Tensor self, Tensor start, Tensor end, Tensor value, int axis=-1) -> Tensor
Mop_plugin/ops/opapi/LSTMBackwardKernelNpuOpApi.cpp+7-4
@@ -31,7 +31,7 @@ std::vector<at::Tensor> squeeze_chunk_result(const at::TensorList& chunk_result)
31 return squeezed_result;31 return squeezed_result;
32}32}
33 33 
34-std::tuple<at::Tensor, at::Tensor, at::Tensor, std::vector<at::Tensor>> _lstm_npu_backward(34+std::tuple<at::Tensor, std::vector<at::Tensor>, std::vector<at::Tensor>> _lstm_npu_backward(
35 const at::Tensor &grad_y,35 const at::Tensor &grad_y,
36 const at::Tensor &grad_hy,36 const at::Tensor &grad_hy,
37 const at::Tensor &grad_cy,37 const at::Tensor &grad_cy,
@@ -71,7 +71,6 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, std::vector<at::Tensor>> _lstm_np
71 }71 }
72 72 
73 int64_t list_length = D * num_layers;73 int64_t list_length = D * num_layers;
74- 
75 const int64_t split_dim = 0;74 const int64_t split_dim = 0;
76 auto i_chunk_origin = at::chunk(i, list_length, split_dim);75 auto i_chunk_origin = at::chunk(i, list_length, split_dim);
77 auto j_chunk_origin = at::chunk(j, list_length, split_dim);76 auto j_chunk_origin = at::chunk(j, list_length, split_dim);
@@ -127,10 +126,14 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor, std::vector<at::Tensor>> _lstm_np
127 out_c_prev,126 out_c_prev,
128 param_list_);127 param_list_);
129 128 
129+ std::vector<at::Tensor> hx_prev_vec;
130+ hx_prev_vec.push_back(out_h_prev);
131+ hx_prev_vec.push_back(out_c_prev);
132+ at::TensorList out_hx_prev = at::TensorList(hx_prev_vec);
130 return std::make_tuple(133 return std::make_tuple(
131 out0,134 out0,
132- out_h_prev,135+ out_hx_prev.vec(),
133- out_c_prev,
134 param_list_.vec());136 param_list_.vec());
135}137}
138+ 
136}139}
Mop_plugin/utils/KernelNpuOutputSize.cpp+14-9
@@ -2275,13 +2275,14 @@ c10::SmallVector<int64_t, SIZE> lstm_backward_npu_output_size(const at::Tensor &
2275{2275{
2276 c10::SmallVector<int64_t, SIZE> output_shape;2276 c10::SmallVector<int64_t, SIZE> output_shape;
2277 if (batch_sizes.has_value()) {2277 if (batch_sizes.has_value()) {
2278- output_shape = {input.size(0), input.size(1)};2278+ const at::Tensor& bs_tensor = batch_sizes.value();
2279- } else {2279+ if (bs_tensor.numel() > 0) {
2280- if (batch_first) {2280+ output_shape = {input.size(0), input.size(1)};
2281- output_shape = {input.size(1), input.size(0), input.size(2)};
2282 } else {2281 } else {
2283 output_shape = {input.size(0), input.size(1), input.size(2)};2282 output_shape = {input.size(0), input.size(1), input.size(2)};
2284 }2283 }
2284+ } else {
2285+ output_shape = {input.size(0), input.size(1), input.size(2)};
2285 }2286 }
2286 return output_shape;2287 return output_shape;
2287}2288}
@@ -2295,11 +2296,15 @@ c10::SmallVector<int64_t, SIZE> lstm_backward_npu_hc_prev_output_size(const at::
2295 int64_t time_step;2296 int64_t time_step;
2296 if (batch_sizes.has_value()) {2297 if (batch_sizes.has_value()) {
2297 const at::Tensor& bs_tensor = batch_sizes.value();2298 const at::Tensor& bs_tensor = batch_sizes.value();
2298- time_step = bs_tensor.size(0);2299+ if (bs_tensor.numel() > 0) {
2299- TORCH_CHECK(time_step > 0 && input.size(0) % time_step == 0,2300+ time_step = bs_tensor.size(0);
2300- "The input batch_sizes is invalid, time_step must > 0 and input.size(0) must be divisible by time_step, but got ",2301+ TORCH_CHECK(time_step > 0 && input.size(0) % time_step == 0,
2301- "time_step=", time_step, ", input_size0=", input.size(0), ", batch_sizes_shape=", bs_tensor.sizes(), OPS_ERROR(ErrCode::PARAM));2302+ "The _lstm_npu_backward input batch_sizes is invalid, time_step must > 0 and input.size(0) must be divisible by time_step, but got ",
2302- batch_size = input.size(0) / time_step;2303+ "time_step=", time_step, ", input_size0=", input.size(0), ", batch_sizes_shape=", bs_tensor.sizes(), OPS_ERROR(ErrCode::PARAM));
2304+ batch_size = input.size(0) / time_step;
2305+ } else {
2306+ batch_size = batch_first ? input.size(0) : input.size(1);
2307+ }
2303 } else {2308 } else {
2304 batch_size = batch_first ? input.size(0) : input.size(1);2309 batch_size = batch_first ? input.size(0) : input.size(1);
2305 }2310 }
Mtest/core_tests/torch_npu_OpApi_schema_all.json+1-1
@@ -4771,7 +4771,7 @@
4771 "func: _lstm_npu(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)": {4771 "func: _lstm_npu(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor)": {
4772 "version": ["all_version"]4772 "version": ["all_version"]
4773 },4773 },
4774- "func: _lstm_npu_backward(Tensor grad_y, Tensor grad_hy, Tensor grad_cy, Tensor input, Tensor[] hx, Tensor[] params, Tensor i, Tensor j, Tensor f, Tensor o, Tensor h, Tensor c, Tensor tanhc, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor, Tensor, Tensor[])": {4774+ "func: _lstm_npu_backward(Tensor grad_y, Tensor grad_hy, Tensor grad_cy, Tensor input, Tensor[] hx, Tensor[] params, Tensor i, Tensor j, Tensor f, Tensor o, Tensor h, Tensor c, Tensor tanhc, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, bool? batch_first=False, Tensor? batch_sizes=None) -> (Tensor, Tensor[], Tensor[])": {
4775 "version": ["all_version"]4775 "version": ["all_version"]
4776 },4776 },
4777 "func: save_npugraph_tensor(Tensor self, *, str? save_path=None) -> ()": {4777 "func: save_npugraph_tensor(Tensor self, *, str? save_path=None) -> ()": {