mv to_torch_tensor to common

Author: Han Wang

deepmd/pt_expt/common.py

@@ -0,0 +1,35 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import importlib
+from typing import (
+    Any,
+    overload,
+)
+
+import numpy as np
+
+from deepmd.pt_expt.utils import (
+    env,
+)
+
+torch = importlib.import_module("torch")
+
+
+@overload
+def to_torch_array(array: np.ndarray) -> torch.Tensor: ...
+
+
+@overload
+def to_torch_array(array: None) -> None: ...
+
+
+@overload
+def to_torch_array(array: torch.Tensor) -> torch.Tensor: ...
+
+
+def to_torch_array(array: Any) -> torch.Tensor | None:
+    """Convert input to a torch tensor on the pt-expt device."""
+    if array is None:
+        return None
+    if torch.is_tensor(array):
+        return array
+    return torch.as_tensor(array, device=env.DEVICE)

deepmd/pt_expt/utils/network.py

@@ -19,21 +19,13 @@
     make_fitting_network,
     make_multilayer_network,
 )
-from deepmd.pt_expt.utils import (
-    env,
+from deepmd.pt_expt.common import (
+    to_torch_array,
 )
 
 torch = importlib.import_module("torch")
 
 
-def _to_torch_array(value: Any) -> torch.Tensor | None:
-    if value is None:
-        return None
-    if torch.is_tensor(value):
-        return value
-    return torch.as_tensor(value, device=env.DEVICE)
-
-
 class TorchArrayParam(torch.nn.Parameter):
     def __new__(cls, data: Any = None, requires_grad: bool = True) -> Self:
         return torch.nn.Parameter.__new__(cls, data, requires_grad)
@@ -52,7 +44,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
         for name in ("w", "b", "idt"):
             if name in self._parameters or name in self._buffers:
                 continue
-            val = _to_torch_array(getattr(self, name))
+            val = to_torch_array(getattr(self, name))
             if val is None:
                 continue
             if self.trainable:
@@ -66,7 +58,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
     def __setattr__(self, name: str, value: Any) -> None:
         if name in {"w", "b", "idt"} and "_parameters" in self.__dict__:
-            val = _to_torch_array(value)
+            val = to_torch_array(value)
             if val is None:
                 return super().__setattr__(name, None)
             if getattr(self, "trainable", False):