deepmd.pt_expt.utils.network ============================ .. py:module:: deepmd.pt_expt.utils.network Classes ------- .. autoapisummary:: deepmd.pt_expt.utils.network.TorchArrayParam deepmd.pt_expt.utils.network.NativeLayer deepmd.pt_expt.utils.network.NativeNet deepmd.pt_expt.utils.network.EmbeddingNet deepmd.pt_expt.utils.network.FittingNet deepmd.pt_expt.utils.network.NetworkCollection deepmd.pt_expt.utils.network.LayerNorm deepmd.pt_expt.utils.network.Identity Functions --------- .. autoapisummary:: deepmd.pt_expt.utils.network._torch_activation Module Contents --------------- .. py:class:: TorchArrayParam Bases: :py:obj:`torch.nn.Parameter` Parameter subclass that supports ``np.array(param)`` conversion. Note: this class is intentionally NOT used for model parameters. ``make_fx`` (``torch.fx.experimental.proxy_tensor``) uses ``ProxyTorchDispatchMode`` to intercept tensor operations. When an operand is a *subclass* of ``torch.Tensor`` (including subclasses of ``torch.nn.Parameter``), PyTorch invokes the ``__torch_function__`` protocol which the proxy dispatch mode does not handle, causing ``aten.mm`` and other ops to fail with "Multiple dispatch failed … returned NotImplemented". Using plain ``torch.nn.Parameter`` avoids this because the proxy mode is designed to work with the base ``Parameter`` type. ``TorchArrayParam`` is kept only for backward compatibility and should not be used for new code. .. !! processed by numpydoc !! .. py:method:: __array__(dtype: Any | None = None) -> numpy.ndarray .. py:class:: NativeLayer(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.utils.network.NativeLayer`, :py:obj:`torch.nn.Module` PyTorch layer wrapping dpmodel's ``NativeLayer``. Two aspects of the inherited dpmodel ``call()`` are incompatible with ``make_fx`` tracing (used to export ``forward_lower`` with ``autograd.grad``-based force/virial computation): 1. **Ellipsis indexing** (``self.w[...]``): On a ``torch.Tensor`` this triggers ``aten.alias``, an op that ``ProxyTorchDispatchMode`` does not support, resulting in "Multiple dispatch failed for ``aten.alias.default``". 2. **``array_api_compat`` wrappers** (``xp = array_api_compat .array_namespace(x); xp.matmul(…)``): The wrappers re-enter ``torch.matmul`` through Python, which goes through the ``__torch_function__`` protocol. Under the proxy dispatch mode this path also fails with "Multiple dispatch failed". This class therefore overrides ``call()`` with an implementation that uses plain ``torch`` ops exclusively (``torch.matmul``, ``torch.tanh``, etc.), avoiding both issues. Trainable weights are stored as plain ``torch.nn.Parameter`` (not ``TorchArrayParam``) for the same ``make_fx`` compatibility reason — see the ``TorchArrayParam`` docstring. .. !! processed by numpydoc !! .. py:method:: __call__(*args: Any, **kwargs: Any) -> Any Forward pass in NumPy implementation. .. !! processed by numpydoc !! .. py:method:: __setattr__(name: str, value: Any) -> None .. py:method:: call(x: torch.Tensor) -> torch.Tensor Forward pass using pure torch ops. Overrides dpmodel's ``call()`` to ensure compatibility with ``make_fx`` (``torch.fx.experimental.proxy_tensor``). The dpmodel implementation uses ``self.w[...]`` and ``array_api_compat.array_namespace(x).matmul(…)`` for backend-agnostic array operations. Both patterns break under ``make_fx``'s ``ProxyTorchDispatchMode``: - ``self.w[...]`` emits ``aten.alias`` which the proxy mode cannot dispatch. - ``array_api_compat`` re-enters ``torch.matmul`` via Python, hitting ``__torch_function__`` which the proxy mode returns ``NotImplemented`` for. This override uses ``torch.matmul``, ``torch.cat``, and ``_torch_activation`` directly, sidestepping both issues. .. !! processed by numpydoc !! .. py:method:: forward(x: torch.Tensor) -> torch.Tensor .. py:function:: _torch_activation(x: torch.Tensor, name: str) -> torch.Tensor Apply activation function using native torch ops. The dpmodel ``get_activation_fn`` returns closures that call ``array_api_compat.array_namespace(x).tanh(x)`` etc. Under ``make_fx`` proxy tracing, the ``array_api_compat`` indirection triggers ``__torch_function__`` dispatch failures. This function calls ``torch.tanh`` and friends directly to avoid the issue. .. !! processed by numpydoc !! .. py:class:: NativeNet(layers: list[dict] | None = None) Bases: :py:obj:`make_multilayer_network`\ (\ :py:obj:`NativeLayer`\ , :py:obj:`NativeOP`\ ) Native representation of a neural network. :Parameters: **layers** : :class:`python:list`\[:obj:`NativeLayer`], :obj:`optional` The layers of the network. .. !! processed by numpydoc !! .. py:attribute:: layers .. py:class:: EmbeddingNet(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.utils.network.EmbeddingNet`, :py:obj:`torch.nn.Module` The embedding network. :Parameters: **in_dim** Input dimension. **neuron** The number of neurons in each layer. The output dimension is the same as the dimension of the last layer. **activation_function** The activation function. **resnet_dt** Use time step at the resnet architecture. **precision** Floating point precision for the model parameters. **seed** : :class:`python:int`, :obj:`optional` Random seed. **bias** : :ref:`bool `, :obj:`Optional ` Whether to use bias in the embedding layer. **trainable** : :ref:`bool ` or :class:`python:list`\[:ref:`bool `], :obj:`Optional ` Whether the weights are trainable. If a list, each element corresponds to a layer. .. !! processed by numpydoc !! .. py:attribute:: layers .. py:method:: __call__(*args: Any, **kwargs: Any) -> Any .. py:method:: forward(x: torch.Tensor) -> torch.Tensor .. py:class:: FittingNet(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.utils.network.FittingNet`, :py:obj:`torch.nn.Module` The fitting network. It may be implemented as an embedding net connected with a linear output layer. :Parameters: **in_dim** Input dimension. **out_dim** Output dimension **neuron** The number of neurons in each hidden layer. **activation_function** The activation function. **resnet_dt** Use time step at the resnet architecture. **precision** Floating point precision for the model parameters. **bias_out** The last linear layer has bias. **seed** : :class:`python:int`, :obj:`optional` Random seed. **trainable** : :ref:`bool ` or :class:`python:list`\[:ref:`bool `], :obj:`optional` Whether the network is trainable. .. !! processed by numpydoc !! .. py:attribute:: layers .. py:method:: __call__(*args: Any, **kwargs: Any) -> Any .. py:method:: forward(x: torch.Tensor) -> torch.Tensor .. py:class:: NetworkCollection(*args: Any, **kwargs: Any) Bases: :py:obj:`deepmd.dpmodel.utils.network.NetworkCollection` A collection of networks for multiple elements. The number of dimensions for types might be 0, 1, or 2. - 0: embedding or fitting with type embedding, in () - 1: embedding with type_one_side, or fitting, in (type_i) - 2: embedding without type_one_side, in (type_i, type_j) :Parameters: **ndim** : :class:`python:int` The number of dimensions. **network_type** : :class:`python:str`, :obj:`optional` The type of the network. **networks** : :class:`python:dict`, :obj:`optional` The networks to initialize with. .. !! processed by numpydoc !! .. py:attribute:: NETWORK_TYPE_MAP :type: ClassVar[dict[str, type]] .. py:attribute:: _module_networks .. py:method:: __setitem__(key: int | tuple, value: Any) -> None .. py:class:: LayerNorm(num_in: int, eps: float = 1e-05, uni_init: bool = True, trainable: bool = True, precision: str = DEFAULT_PRECISION, seed: int | list[int] | None = None) Bases: :py:obj:`deepmd.dpmodel.utils.network.LayerNorm`, :py:obj:`NativeLayer` Implementation of Layer Normalization layer. :Parameters: **num_in** : :class:`python:int` The input dimension of the layer. **eps** : :class:`python:float`, :obj:`optional` A small value added to prevent division by zero in calculations. **uni_init** : :ref:`bool `, :obj:`optional` If initialize the weights to be zeros and ones. **trainable** : :ref:`bool `, :obj:`optional` If the weights are trainable. **precision** : :class:`python:str`, :obj:`optional` The precision of the layer. **seed** : :class:`python:int`, :obj:`optional` Random seed. .. !! processed by numpydoc !! .. py:class:: Identity Bases: :py:obj:`deepmd.dpmodel.utils.network.Identity` The unit operation of a native model. .. !! processed by numpydoc !! .. py:method:: forward(x: torch.Tensor) -> torch.Tensor