From 482956e7cdf147303c821600231cd359a2126c59 Mon Sep 17 00:00:00 2001 From: caic99 Date: Mon, 31 Mar 2025 07:35:42 +0000 Subject: [PATCH 01/11] perf: use `torch.split` in replace of slicing ops in repflow --- deepmd/pt/model/descriptor/repflow_layer.py | 65 ++++++++------------- 1 file changed, 23 insertions(+), 42 deletions(-) diff --git a/deepmd/pt/model/descriptor/repflow_layer.py b/deepmd/pt/model/descriptor/repflow_layer.py index 01f8477b7b..039d64a03f 100644 --- a/deepmd/pt/model/descriptor/repflow_layer.py +++ b/deepmd/pt/model/descriptor/repflow_layer.py @@ -397,42 +397,31 @@ def optim_angle_update( edge_ebd: torch.Tensor, feat: str = "edge", ) -> torch.Tensor: - angle_dim = angle_ebd.shape[-1] - node_dim = node_ebd.shape[-1] - edge_dim = edge_ebd.shape[-1] - sub_angle_idx = (0, angle_dim) - sub_node_idx = (angle_dim, angle_dim + node_dim) - sub_edge_idx_ij = (angle_dim + node_dim, angle_dim + node_dim + edge_dim) - sub_edge_idx_ik = ( - angle_dim + node_dim + edge_dim, - angle_dim + node_dim + 2 * edge_dim, - ) - if feat == "edge": + assert self.edge_angle_linear1 is not None matrix, bias = self.edge_angle_linear1.matrix, self.edge_angle_linear1.bias elif feat == "angle": + assert self.angle_self_linear is not None matrix, bias = self.angle_self_linear.matrix, self.angle_self_linear.bias else: raise NotImplementedError - assert angle_dim + node_dim + 2 * edge_dim == matrix.size()[0] + assert bias is not None - # nf * nloc * a_sel * a_sel * angle_dim - sub_angle_update = torch.matmul( - angle_ebd, matrix[sub_angle_idx[0] : sub_angle_idx[1]] + angle_dim = angle_ebd.shape[-1] + node_dim = node_ebd.shape[-1] + edge_dim = edge_ebd.shape[-1] + # angle_dim, node_dim, edge_dim, edge_dim + sub_angle, sub_node, sub_edge_ij, sub_edge_ik = torch.split( + matrix, [angle_dim, node_dim, edge_dim, edge_dim] ) + # nf * nloc * a_sel * a_sel * angle_dim + sub_angle_update = torch.matmul(angle_ebd, sub_angle) # nf * nloc * angle_dim - sub_node_update = torch.matmul( - node_ebd, matrix[sub_node_idx[0] : sub_node_idx[1]] - ) - + sub_node_update = torch.matmul(node_ebd, sub_node) # nf * nloc * a_nnei * angle_dim - sub_edge_update_ij = torch.matmul( - edge_ebd, matrix[sub_edge_idx_ij[0] : sub_edge_idx_ij[1]] - ) - sub_edge_update_ik = torch.matmul( - edge_ebd, matrix[sub_edge_idx_ik[0] : sub_edge_idx_ik[1]] - ) + sub_edge_update_ij = torch.matmul(edge_ebd, sub_edge_ij) + sub_edge_update_ik = torch.matmul(edge_ebd, sub_edge_ik) result_update = ( sub_angle_update @@ -450,36 +439,28 @@ def optim_edge_update( nlist: torch.Tensor, feat: str = "node", ) -> torch.Tensor: - node_dim = node_ebd.shape[-1] - edge_dim = edge_ebd.shape[-1] - sub_node_idx = (0, node_dim) - sub_node_ext_idx = (node_dim, 2 * node_dim) - sub_edge_idx = (2 * node_dim, 2 * node_dim + edge_dim) - if feat == "node": matrix, bias = self.node_edge_linear.matrix, self.node_edge_linear.bias elif feat == "edge": matrix, bias = self.edge_self_linear.matrix, self.edge_self_linear.bias else: raise NotImplementedError + assert bias is not None + + node_dim = node_ebd.shape[-1] + edge_dim = edge_ebd.shape[-1] assert 2 * node_dim + edge_dim == matrix.size()[0] + # node_dim, node_dim, edge_dim + node, node_ext, edge = torch.split(matrix, node_dim) # nf * nloc * node/edge_dim - sub_node_update = torch.matmul( - node_ebd, matrix[sub_node_idx[0] : sub_node_idx[1]] - ) - + sub_node_update = torch.matmul(node_ebd, node) # nf * nall * node/edge_dim - sub_node_ext_update = torch.matmul( - node_ebd_ext, matrix[sub_node_ext_idx[0] : sub_node_ext_idx[1]] - ) + sub_node_ext_update = torch.matmul(node_ebd_ext, node_ext) # nf * nloc * nnei * node/edge_dim sub_node_ext_update = _make_nei_g1(sub_node_ext_update, nlist) - # nf * nloc * nnei * node/edge_dim - sub_edge_update = torch.matmul( - edge_ebd, matrix[sub_edge_idx[0] : sub_edge_idx[1]] - ) + sub_edge_update = torch.matmul(edge_ebd, edge) result_update = ( sub_edge_update + sub_node_ext_update + sub_node_update[:, :, None, :] From 7a3226536e4956f7add0eb061b7102b2fb532a48 Mon Sep 17 00:00:00 2001 From: caic99 Date: Mon, 31 Mar 2025 07:43:00 +0000 Subject: [PATCH 02/11] update split usage for better readability --- deepmd/pt/model/descriptor/repflow_layer.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/deepmd/pt/model/descriptor/repflow_layer.py b/deepmd/pt/model/descriptor/repflow_layer.py index 039d64a03f..66cf30bf00 100644 --- a/deepmd/pt/model/descriptor/repflow_layer.py +++ b/deepmd/pt/model/descriptor/repflow_layer.py @@ -449,9 +449,8 @@ def optim_edge_update( node_dim = node_ebd.shape[-1] edge_dim = edge_ebd.shape[-1] - assert 2 * node_dim + edge_dim == matrix.size()[0] # node_dim, node_dim, edge_dim - node, node_ext, edge = torch.split(matrix, node_dim) + node, node_ext, edge = torch.split(matrix, [node_dim, node_dim, edge_dim]) # nf * nloc * node/edge_dim sub_node_update = torch.matmul(node_ebd, node) From 2cd31cc51ec9e2463117fc4c3bf0af85816df29f Mon Sep 17 00:00:00 2001 From: caic99 Date: Mon, 31 Mar 2025 09:07:06 +0000 Subject: [PATCH 03/11] update dp backend --- deepmd/dpmodel/descriptor/repflows.py | 63 ++++++++++----------------- 1 file changed, 23 insertions(+), 40 deletions(-) diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 56692589d7..36c9a553d8 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -796,42 +796,35 @@ def optim_angle_update( feat: str = "edge", ) -> np.ndarray: xp = array_api_compat.array_namespace(angle_ebd, node_ebd, edge_ebd) - angle_dim = angle_ebd.shape[-1] - node_dim = node_ebd.shape[-1] - edge_dim = edge_ebd.shape[-1] - sub_angle_idx = (0, angle_dim) - sub_node_idx = (angle_dim, angle_dim + node_dim) - sub_edge_idx_ij = (angle_dim + node_dim, angle_dim + node_dim + edge_dim) - sub_edge_idx_ik = ( - angle_dim + node_dim + edge_dim, - angle_dim + node_dim + 2 * edge_dim, - ) if feat == "edge": + assert self.edge_angle_linear1 is not None matrix, bias = self.edge_angle_linear1.w, self.edge_angle_linear1.b elif feat == "angle": + assert self.angle_self_linear is not None matrix, bias = self.angle_self_linear.w, self.angle_self_linear.b else: raise NotImplementedError - assert angle_dim + node_dim + 2 * edge_dim == matrix.shape[0] + assert bias is not None - # nf * nloc * a_sel * a_sel * angle_dim - sub_angle_update = xp.matmul( - angle_ebd, matrix[sub_angle_idx[0] : sub_angle_idx[1], :] + angle_dim = angle_ebd.shape[-1] + node_dim = node_ebd.shape[-1] + edge_dim = edge_ebd.shape[-1] + angle_dim = angle_ebd.shape[-1] + node_dim = node_ebd.shape[-1] + edge_dim = edge_ebd.shape[-1] + # angle_dim, node_dim, edge_dim, edge_dim + sub_angle, sub_node, sub_edge_ij, sub_edge_ik = xp.split( + matrix, [angle_dim, node_dim, edge_dim, edge_dim] ) + # nf * nloc * a_sel * a_sel * angle_dim + sub_angle_update = xp.matmul(angle_ebd, sub_angle) # nf * nloc * angle_dim - sub_node_update = xp.matmul( - node_ebd, matrix[sub_node_idx[0] : sub_node_idx[1], :] - ) - + sub_node_update = xp.matmul(node_ebd, sub_node) # nf * nloc * a_nnei * angle_dim - sub_edge_update_ij = xp.matmul( - edge_ebd, matrix[sub_edge_idx_ij[0] : sub_edge_idx_ij[1], :] - ) - sub_edge_update_ik = xp.matmul( - edge_ebd, matrix[sub_edge_idx_ik[0] : sub_edge_idx_ik[1], :] - ) + sub_edge_update_ij = xp.matmul(edge_ebd, sub_edge_ij) + sub_edge_update_ik = xp.matmul(edge_ebd, sub_edge_ik) result_update = ( sub_angle_update @@ -850,11 +843,6 @@ def optim_edge_update( feat: str = "node", ) -> np.ndarray: xp = array_api_compat.array_namespace(node_ebd, node_ebd_ext, edge_ebd, nlist) - node_dim = node_ebd.shape[-1] - edge_dim = edge_ebd.shape[-1] - sub_node_idx = (0, node_dim) - sub_node_ext_idx = (node_dim, 2 * node_dim) - sub_edge_idx = (2 * node_dim, 2 * node_dim + edge_dim) if feat == "node": matrix, bias = self.node_edge_linear.w, self.node_edge_linear.b @@ -862,24 +850,19 @@ def optim_edge_update( matrix, bias = self.edge_self_linear.w, self.edge_self_linear.b else: raise NotImplementedError - assert 2 * node_dim + edge_dim == matrix.shape[0] - + node_dim = node_ebd.shape[-1] + edge_dim = edge_ebd.shape[-1] + node, node_ext, edge = xp.split(matrix, [node_dim, node_dim, edge_dim]) # nf * nloc * node/edge_dim - sub_node_update = xp.matmul( - node_ebd, matrix[sub_node_idx[0] : sub_node_idx[1], :] - ) + sub_node_update = xp.matmul(node_ebd, node) # nf * nall * node/edge_dim - sub_node_ext_update = xp.matmul( - node_ebd_ext, matrix[sub_node_ext_idx[0] : sub_node_ext_idx[1], :] - ) + sub_node_ext_update = xp.matmul(node_ebd_ext, node_ext) # nf * nloc * nnei * node/edge_dim sub_node_ext_update = _make_nei_g1(sub_node_ext_update, nlist) # nf * nloc * nnei * node/edge_dim - sub_edge_update = xp.matmul( - edge_ebd, matrix[sub_edge_idx[0] : sub_edge_idx[1], :] - ) + sub_edge_update = xp.matmul(edge_ebd, edge) result_update = ( sub_edge_update + sub_node_ext_update + sub_node_update[:, :, xp.newaxis, :] From 9cc0452c32a8aed3a1447182360cc8430444b682 Mon Sep 17 00:00:00 2001 From: Chun Cai Date: Tue, 1 Apr 2025 10:54:30 +0800 Subject: [PATCH 04/11] remove duplicated statements Signed-off-by: Chun Cai --- deepmd/dpmodel/descriptor/repflows.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 36c9a553d8..bcea7c3c2a 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -807,9 +807,6 @@ def optim_angle_update( raise NotImplementedError assert bias is not None - angle_dim = angle_ebd.shape[-1] - node_dim = node_ebd.shape[-1] - edge_dim = edge_ebd.shape[-1] angle_dim = angle_ebd.shape[-1] node_dim = node_ebd.shape[-1] edge_dim = edge_ebd.shape[-1] From e437fceaa620b245629ef11f6224898098df2050 Mon Sep 17 00:00:00 2001 From: caic99 Date: Tue, 1 Apr 2025 04:42:01 +0000 Subject: [PATCH 05/11] fix ut --- deepmd/dpmodel/descriptor/repflows.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index bcea7c3c2a..3e8d33e91a 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -849,7 +849,8 @@ def optim_edge_update( raise NotImplementedError node_dim = node_ebd.shape[-1] edge_dim = edge_ebd.shape[-1] - node, node_ext, edge = xp.split(matrix, [node_dim, node_dim, edge_dim]) + assert node_dim * 2 + edge_dim == matrix.shape[0] + node, node_ext, edge, _ = xp.split(matrix, [node_dim, node_dim, edge_dim]) # nf * nloc * node/edge_dim sub_node_update = xp.matmul(node_ebd, node) From 229d1cfb81a438fa3a95bc1daaccbbc61e6b1baa Mon Sep 17 00:00:00 2001 From: caic99 Date: Tue, 1 Apr 2025 06:00:49 +0000 Subject: [PATCH 06/11] use unsqueeze and ellipsis to reduce slice op --- deepmd/pt/model/descriptor/repflow_layer.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/deepmd/pt/model/descriptor/repflow_layer.py b/deepmd/pt/model/descriptor/repflow_layer.py index 66cf30bf00..8ffeb39d8c 100644 --- a/deepmd/pt/model/descriptor/repflow_layer.py +++ b/deepmd/pt/model/descriptor/repflow_layer.py @@ -425,9 +425,9 @@ def optim_angle_update( result_update = ( sub_angle_update - + sub_node_update[:, :, None, None, :] - + sub_edge_update_ij[:, :, None, :, :] - + sub_edge_update_ik[:, :, :, None, :] + + sub_node_update.unsqueeze(2).unsqueeze(3) + + sub_edge_update_ij.unsqueeze(2) + + sub_edge_update_ik.unsqueeze(3) ) + bias return result_update @@ -462,7 +462,7 @@ def optim_edge_update( sub_edge_update = torch.matmul(edge_ebd, edge) result_update = ( - sub_edge_update + sub_node_ext_update + sub_node_update[:, :, None, :] + sub_edge_update + sub_node_ext_update + sub_node_update.unsqueeze(2) ) + bias return result_update @@ -590,7 +590,7 @@ def forward( nb, nloc, self.n_multi_edge_message, self.n_dim ) for head_index in range(self.n_multi_edge_message): - n_update_list.append(node_edge_update_mul_head[:, :, head_index, :]) + n_update_list.append(node_edge_update_mul_head[..., head_index, :]) else: n_update_list.append(node_edge_update) # update node_ebd @@ -625,14 +625,14 @@ def forward( edge_ebd_for_angle = self.a_compress_e_linear(edge_ebd) else: # use the first a_compress_dim dim for node and edge - node_ebd_for_angle = node_ebd[:, :, : self.n_a_compress_dim] - edge_ebd_for_angle = edge_ebd[:, :, :, : self.e_a_compress_dim] + node_ebd_for_angle = node_ebd[..., : self.n_a_compress_dim] + edge_ebd_for_angle = edge_ebd[..., : self.e_a_compress_dim] else: node_ebd_for_angle = node_ebd edge_ebd_for_angle = edge_ebd # nb x nloc x a_nnei x e_dim - edge_for_angle = edge_ebd_for_angle[:, :, : self.a_sel, :] + edge_for_angle = edge_ebd_for_angle[..., :self.a_sel, :] # nb x nloc x a_nnei x e_dim edge_for_angle = torch.where( a_nlist_mask.unsqueeze(-1), edge_for_angle, 0.0 @@ -680,8 +680,8 @@ def forward( # nb x nloc x a_nnei x a_nnei x e_dim weighted_edge_angle_update = ( - a_sw[:, :, :, None, None] - * a_sw[:, :, None, :, None] + a_sw[..., None, None] + * a_sw[..., None, :, None] * edge_angle_update ) # nb x nloc x a_nnei x e_dim From 7985409e36f14574115c1ccfb6d99e9e55d54264 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 1 Apr 2025 06:02:29 +0000 Subject: [PATCH 07/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/descriptor/repflow_layer.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/deepmd/pt/model/descriptor/repflow_layer.py b/deepmd/pt/model/descriptor/repflow_layer.py index 8ffeb39d8c..3245548ffc 100644 --- a/deepmd/pt/model/descriptor/repflow_layer.py +++ b/deepmd/pt/model/descriptor/repflow_layer.py @@ -632,7 +632,7 @@ def forward( edge_ebd_for_angle = edge_ebd # nb x nloc x a_nnei x e_dim - edge_for_angle = edge_ebd_for_angle[..., :self.a_sel, :] + edge_for_angle = edge_ebd_for_angle[..., : self.a_sel, :] # nb x nloc x a_nnei x e_dim edge_for_angle = torch.where( a_nlist_mask.unsqueeze(-1), edge_for_angle, 0.0 @@ -680,9 +680,7 @@ def forward( # nb x nloc x a_nnei x a_nnei x e_dim weighted_edge_angle_update = ( - a_sw[..., None, None] - * a_sw[..., None, :, None] - * edge_angle_update + a_sw[..., None, None] * a_sw[..., None, :, None] * edge_angle_update ) # nb x nloc x a_nnei x e_dim reduced_edge_angle_update = torch.sum( From 553d234d4e7200007460a5741a7ef44e2777c264 Mon Sep 17 00:00:00 2001 From: caic99 Date: Tue, 1 Apr 2025 06:10:39 +0000 Subject: [PATCH 08/11] fix: array api does not support .split --- deepmd/dpmodel/descriptor/repflows.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 3e8d33e91a..d1505c87ed 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -810,10 +810,12 @@ def optim_angle_update( angle_dim = angle_ebd.shape[-1] node_dim = node_ebd.shape[-1] edge_dim = edge_ebd.shape[-1] - # angle_dim, node_dim, edge_dim, edge_dim - sub_angle, sub_node, sub_edge_ij, sub_edge_ik = xp.split( - matrix, [angle_dim, node_dim, edge_dim, edge_dim] - ) + + # Array API does not provide a way to split the array + sub_angle = matrix[:angle_dim, ...] # angle_dim + sub_node = matrix[angle_dim:angle_dim+node_dim, ...] # node_dim + sub_edge_ij = matrix[angle_dim+node_dim:angle_dim+node_dim+edge_dim, ...] # edge_dim + sub_edge_ik = matrix[angle_dim+node_dim+edge_dim:, ...] # edge_dim # nf * nloc * a_sel * a_sel * angle_dim sub_angle_update = xp.matmul(angle_ebd, sub_angle) @@ -850,7 +852,11 @@ def optim_edge_update( node_dim = node_ebd.shape[-1] edge_dim = edge_ebd.shape[-1] assert node_dim * 2 + edge_dim == matrix.shape[0] - node, node_ext, edge, _ = xp.split(matrix, [node_dim, node_dim, edge_dim]) + # Array API does not provide a way to split the array + node = matrix[:node_dim, ...] # node_dim + node_ext = matrix[node_dim:2*node_dim, ...] # node_dim + edge = matrix[2*node_dim:2*node_dim + edge_dim, ...] # edge_dim + # nf * nloc * node/edge_dim sub_node_update = xp.matmul(node_ebd, node) From f8e9d8e7109f913f45cac818310ae623fea06157 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 1 Apr 2025 06:12:51 +0000 Subject: [PATCH 09/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/dpmodel/descriptor/repflows.py | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index d1505c87ed..5f2f589921 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -812,10 +812,12 @@ def optim_angle_update( edge_dim = edge_ebd.shape[-1] # Array API does not provide a way to split the array - sub_angle = matrix[:angle_dim, ...] # angle_dim - sub_node = matrix[angle_dim:angle_dim+node_dim, ...] # node_dim - sub_edge_ij = matrix[angle_dim+node_dim:angle_dim+node_dim+edge_dim, ...] # edge_dim - sub_edge_ik = matrix[angle_dim+node_dim+edge_dim:, ...] # edge_dim + sub_angle = matrix[:angle_dim, ...] # angle_dim + sub_node = matrix[angle_dim : angle_dim + node_dim, ...] # node_dim + sub_edge_ij = matrix[ + angle_dim + node_dim : angle_dim + node_dim + edge_dim, ... + ] # edge_dim + sub_edge_ik = matrix[angle_dim + node_dim + edge_dim :, ...] # edge_dim # nf * nloc * a_sel * a_sel * angle_dim sub_angle_update = xp.matmul(angle_ebd, sub_angle) @@ -853,9 +855,9 @@ def optim_edge_update( edge_dim = edge_ebd.shape[-1] assert node_dim * 2 + edge_dim == matrix.shape[0] # Array API does not provide a way to split the array - node = matrix[:node_dim, ...] # node_dim - node_ext = matrix[node_dim:2*node_dim, ...] # node_dim - edge = matrix[2*node_dim:2*node_dim + edge_dim, ...] # edge_dim + node = matrix[:node_dim, ...] # node_dim + node_ext = matrix[node_dim : 2 * node_dim, ...] # node_dim + edge = matrix[2 * node_dim : 2 * node_dim + edge_dim, ...] # edge_dim # nf * nloc * node/edge_dim sub_node_update = xp.matmul(node_ebd, node) From 60147423336d7a667a5f4bb496a43b5283c75124 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 1 Apr 2025 10:22:09 +0000 Subject: [PATCH 10/11] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/pt/model/descriptor/repflow_layer.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/deepmd/pt/model/descriptor/repflow_layer.py b/deepmd/pt/model/descriptor/repflow_layer.py index a5a66f14a3..f109109cfd 100644 --- a/deepmd/pt/model/descriptor/repflow_layer.py +++ b/deepmd/pt/model/descriptor/repflow_layer.py @@ -463,10 +463,7 @@ def optim_edge_update( sub_edge_update = torch.matmul(edge_ebd, edge) result_update = ( - bias - + sub_node_update.unsqueeze(2) - + sub_edge_update - + sub_node_ext_update + bias + sub_node_update.unsqueeze(2) + sub_edge_update + sub_node_ext_update ) return result_update From 8a9fd721df69e5ae6a70e5fd14f5d1a2430016f3 Mon Sep 17 00:00:00 2001 From: Chun Cai Date: Wed, 2 Apr 2025 14:17:05 +0800 Subject: [PATCH 11/11] Add assertion for matrix shape validation in dp backend --- deepmd/dpmodel/descriptor/repflows.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index 02728cf9bd..a49edb534c 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -810,7 +810,7 @@ def optim_angle_update( angle_dim = angle_ebd.shape[-1] node_dim = node_ebd.shape[-1] edge_dim = edge_ebd.shape[-1] - + assert angle_dim + node_dim + 2 * edge_dim == matrix.shape[0] # Array API does not provide a way to split the array sub_angle = matrix[:angle_dim, ...] # angle_dim sub_node = matrix[angle_dim : angle_dim + node_dim, ...] # node_dim