[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

deepmd/descriptor/se_a.py

@@ -41,9 +41,9 @@
     GraphWithoutTensorError,
 )
 from deepmd.utils.graph import (
+    get_extra_embedding_net_variables_from_graph_def,
     get_pattern_nodes_from_graph_def,
     get_tensor_by_name_from_graph,
-    get_extra_embedding_net_variables_from_graph_def,
 )
 from deepmd.utils.network import (
     embedding_net,
@@ -1333,7 +1333,11 @@ def init_variables(
                 extra_suffix = "_one_side_ebd"
             else:
                 extra_suffix = "_two_side_ebd"
-            self.extra_embedding_net_variables = get_extra_embedding_net_variables_from_graph_def(graph_def, suffix, extra_suffix, self.layer_size)
+            self.extra_embedding_net_variables = (
+                get_extra_embedding_net_variables_from_graph_def(
+                    graph_def, suffix, extra_suffix, self.layer_size
+                )
+            )
 
     @property
     def explicit_ntypes(self) -> bool:

deepmd/descriptor/se_atten.py

@@ -42,9 +42,9 @@
 )
 from deepmd.utils.graph import (
     get_attention_layer_variables_from_graph_def,
+    get_extra_embedding_net_variables_from_graph_def,
     get_pattern_nodes_from_graph_def,
     get_tensor_by_name_from_graph,
-    get_extra_embedding_net_variables_from_graph_def,
 )
 from deepmd.utils.network import (
     embedding_net,
@@ -1311,7 +1311,11 @@ def init_variables(
                 ]
 
         if self.stripped_type_embedding:
-            self.two_side_embeeding_net_variables = get_extra_embedding_net_variables_from_graph_def(graph_def, suffix, "_two_side_ebd", self.layer_size)
+            self.two_side_embeeding_net_variables = (
+                get_extra_embedding_net_variables_from_graph_def(
+                    graph_def, suffix, "_two_side_ebd", self.layer_size
+                )
+            )
 
     def build_type_exclude_mask(
         self,

deepmd/utils/graph.py

@@ -236,14 +236,15 @@ def get_embedding_net_variables_from_graph_def(
         embedding_net_variables[item] = np.reshape(tensor_value, tensor_shape)
     return embedding_net_variables
 
+
 def get_variables_from_graph_def_as_numpy_array(graph_def: tf.GraphDef, pattern: str):
     """Get variables from the given tf.GraphDef object, with numpy array returns.
 
     Parameters
     ----------
     graph_def
         The input tf.GraphDef object
-    suffix : str
+    pattern : str
         The name of variable
 
     Returns
@@ -263,7 +264,10 @@ def get_variables_from_graph_def_as_numpy_array(graph_def: tf.GraphDef, pattern:
         tensor_value = get_tensor_by_type(node, dtype)
     return np.reshape(tensor_value, tensor_shape)
 
-def get_extra_embedding_net_variables_from_graph_def(graph_def: tf.GraphDef, suffix: str, extra_suffix: str, layer_size: int):
+
+def get_extra_embedding_net_variables_from_graph_def(
+    graph_def: tf.GraphDef, suffix: str, extra_suffix: str, layer_size: int
+):
     """Get extra embedding net variables from the given tf.GraphDef object.
     The "extra embedding net" means the embedding net with only type embeddings input,
     which occurs in "se_atten_v2" and "se_a_ebd_v2" descriptor.
@@ -274,10 +278,10 @@ def get_extra_embedding_net_variables_from_graph_def(graph_def: tf.GraphDef, suf
         The input tf.GraphDef object
     suffix : str
         The "common" suffix in the descriptor
-    extra_suffix: str
+    extra_suffix : str
         This value depends on the value of "type_one_side".
         It should always be "_one_side_ebd" or "_two_side_ebd"
-    layer_size: int
+    layer_size : int
         The layer size of the embedding net
 
     Returns
@@ -297,6 +301,7 @@ def get_extra_embedding_net_variables_from_graph_def(graph_def: tf.GraphDef, suf
         ] = get_variables_from_graph_def_as_numpy_array(graph_def, bias_pattern)
     return extra_embedding_net_variables
 
+
 def get_embedding_net_variables(model_file: str, suffix: str = "") -> Dict:
     """Get the embedding net variables with the given frozen model(model_file).