diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml index 851a509269..2694175826 100644 --- a/.github/workflows/test_python.yml +++ b/.github/workflows/test_python.yml @@ -63,12 +63,36 @@ jobs: DP_CI_IMPORT_PADDLE_BEFORE_TF: 1 FLAGS_use_stride_compute_kernel: 0 - name: Test TF2 eager mode - run: pytest --cov=deepmd --cov-append source/tests/consistent/io/test_io.py source/jax2tf_tests + run: | + run_pytest_allow_no_tests() { + set +e + pytest "$@" + local status=$? + set -e + if [ "$status" -eq 5 ]; then + # pytest-split may leave an individual shard with no selected + # tests after path/-k filtering. Other shards still cover the + # selected tests, so do not fail the whole matrix for exit 5. + return 0 + fi + return "$status" + } + + run_pytest_allow_no_tests --cov=deepmd --cov-append \ + source/tests/consistent/io/test_io.py \ + source/jax2tf_tests \ + --splits 12 \ + --group ${{ matrix.group }} + run_pytest_allow_no_tests --cov=deepmd --cov-append \ + source/tests/consistent \ + -k tf2 \ + --splits 12 \ + --group ${{ matrix.group }} env: NUM_WORKERS: 0 DP_TEST_TF2_ONLY: 1 DP_DTYPE_PROMOTION_STRICT: 1 - if: matrix.group == 1 + DP_CI_IMPORT_PADDLE_BEFORE_TF: 1 - run: mv .test_durations .test_durations_${{ matrix.group }} - name: Upload partial durations uses: actions/upload-artifact@v7 diff --git a/deepmd/_vendors/__init__.py b/deepmd/_vendors/__init__.py new file mode 100644 index 0000000000..feef50104b --- /dev/null +++ b/deepmd/_vendors/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Vendored third-party modules used by DeePMD-kit.""" diff --git a/deepmd/_vendors/ndtensorflow/__init__.py b/deepmd/_vendors/ndtensorflow/__init__.py new file mode 100644 index 0000000000..a484073031 --- /dev/null +++ b/deepmd/_vendors/ndtensorflow/__init__.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +from typing import ( + Final, +) + +from . import ( + fft, + linalg, +) +from ._array import ( + Array, +) +from ._info import ( + __array_namespace_info__, +) +from ._namespace import * +from ._namespace import __all__ as _namespace_all + +__array_api_version__: Final = "2025.12" + +__all__ = sorted( + set(_namespace_all) + | { + "Array", + "__array_api_version__", + "__array_namespace_info__", + "fft", + "linalg", + } +) + + +def __dir__() -> list[str]: + return __all__ diff --git a/deepmd/_vendors/ndtensorflow/_array.py b/deepmd/_vendors/ndtensorflow/_array.py new file mode 100644 index 0000000000..3b734971b2 --- /dev/null +++ b/deepmd/_vendors/ndtensorflow/_array.py @@ -0,0 +1,553 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +import math +from collections.abc import ( + Callable, + Iterator, +) +from enum import ( + IntEnum, +) +from typing import ( + Any, +) + +import tensorflow as tf +from tensorflow.python.framework import ( + composite_tensor_gradient, +) + + +class DLDeviceType(IntEnum): + CPU = 1 + CUDA = 2 + + +class _ArrayGradient(composite_tensor_gradient.CompositeTensorGradient): + def get_gradient_components(self, value: Array) -> tf.Tensor: + return value._tensor + + def replace_gradient_components( + self, + value: Array, + component_grads: tf.Tensor | None, + ) -> Array | None: + del value + if component_grads is None: + return None + return Array._from_tensor(component_grads) + + +class Array(tf.experimental.ExtensionType): + """User-facing TensorFlow-backed array object. + + The object owns the Array API surface. TensorFlow tensors remain plain + TensorFlow tensors; no TensorFlow class is patched. + """ + + __array_priority__ = 1 + __composite_gradient__ = _ArrayGradient() + + _tensor: tf.Tensor + + def __init__(self, tensor: Any | None = None, /) -> None: + if tensor is None: + raise TypeError( + "'Array' cannot be instantiated without data. Use " + "'ndtensorflow.asarray' or another creation function instead." + ) + if not isinstance(tensor, tf.Tensor): + tensor = tf.convert_to_tensor(tensor) + self._tensor = tensor + + @classmethod + def _from_tensor(cls, tensor: Any, /) -> Array: + if isinstance(tensor, Array): + return tensor + return cls(tensor) + + def _replace_tensor(self, tensor: tf.Tensor) -> None: + self.__dict__["_tensor"] = tensor + + def unwrap(self) -> tf.Tensor: + """Return the wrapped TensorFlow tensor.""" + return self._tensor + + def __tf_tensor__( + self, + dtype: tf.DType | None = None, + name: str | None = None, + ) -> tf.Tensor: + del name + return tf.cast(self._tensor, dtype) if dtype is not None else self._tensor + + @property + def device(self) -> str: + device = self._tensor.device + marker = "/device:" + if marker in device: + return device.rsplit(marker, maxsplit=1)[-1] + return device + + @property + def dtype(self) -> tf.DType: + return self._tensor.dtype + + @property + def mT(self) -> Array: # noqa: N802 + from deepmd._vendors import ndtensorflow as xp + + return xp.matrix_transpose(self) + + @property + def ndim(self) -> int: + return len(self.shape) + + @property + def shape(self) -> tuple[int | tf.Tensor, ...]: + static_shape = self._tensor.shape.as_list() + if all(dim is not None for dim in static_shape): + return tuple(static_shape) + dynamic_shape = tf.shape(self._tensor) + return tuple( + dynamic_shape[ii] if dim is None else dim + for ii, dim in enumerate(static_shape) + ) + + @property + def size(self) -> int | tf.Tensor: + shape = self.shape + if all(isinstance(dim, int) for dim in shape): + return math.prod(shape) + return tf.size(self._tensor) + + @property + def T(self) -> Array: # noqa: N802 + from deepmd._vendors import ndtensorflow as xp + + return xp.permute_dims(self, tuple(range(self.ndim - 1, -1, -1))) + + def astype( + self, + dtype: tf.DType, + /, + *, + copy: bool = True, + device: str | None = None, + ) -> Array: + from deepmd._vendors import ndtensorflow as xp + + return xp.astype(self, dtype, copy=copy, device=device) + + def reshape(self, *shape: Any, copy: bool | None = None) -> Array: + from deepmd._vendors import ndtensorflow as xp + + if len(shape) == 1 and isinstance(shape[0], tuple | list): + shape = tuple(shape[0]) + return xp.reshape(self, tuple(shape), copy=copy) + + def ravel(self) -> Array: + from deepmd._vendors import ndtensorflow as xp + + return xp.reshape(self, (-1,)) + + def squeeze(self, axis: int | tuple[int, ...] | None = None) -> Array: + if axis is None: + return type(self)._from_tensor(tf.squeeze(self._tensor)) + from deepmd._vendors import ndtensorflow as xp + + return xp.squeeze(self, axis=axis) + + def to_device(self, device: str, /, *, stream: int | Any | None = None) -> Array: + del stream + from deepmd._vendors import ndtensorflow as xp + + return xp.asarray(self, device=device, copy=True) + + def __dlpack__( + self, + *, + stream: int | Any | None = None, + max_version: tuple[int, int] | None = None, + dl_device: tuple[int, int] | None = None, + copy: bool | None = None, + ) -> Any: + del stream, max_version, dl_device, copy + if hasattr(self._tensor, "__dlpack__"): + return self._tensor.__dlpack__() + raise BufferError( + "TensorFlow exposes DLPack conversion only outside the stable " + "Tensor API in this environment." + ) + + def __dlpack_device__(self) -> tuple[DLDeviceType, int]: + device = self.device.upper() + if "GPU" in device: + index = int(device.rsplit(":", maxsplit=1)[-1]) + return (DLDeviceType.CUDA, index) + return (DLDeviceType.CPU, 0) + + def __array_namespace__(self, /, *, api_version: str | None = None) -> Any: + del api_version + from deepmd._vendors import ndtensorflow as xp + + return xp + + def __array__(self, dtype: Any | None = None) -> Any: + if not tf.executing_eagerly(): + raise TypeError("cannot convert a TensorFlow graph tensor to a NumPy array") + array = self._tensor.numpy() + return array.astype(dtype) if dtype is not None else array + + def __len__(self) -> int: + if self.ndim == 0: + raise TypeError("len() of unsized array") + dim = self.shape[0] + if not isinstance(dim, int): + raise TypeError("len() requires a statically known leading dimension") + return dim + + def __iter__(self) -> Iterator[Array]: + if self.ndim == 0: + raise ValueError("iteration over a 0-d array") + return (self[i] for i in range(len(self))) + + def __getitem__(self, key: Any, /) -> Array: + key = _normalize_index_key(key) + if isinstance(key, tf.Tensor) and key.dtype == tf.bool: + rank = key.shape.rank + if rank is None: + raise IndexError("boolean index rank must be statically known") + if rank > self.ndim: + raise IndexError( + "boolean index shape is incompatible with indexed array" + ) + if rank == 0: + tensor = tf.expand_dims(self._tensor, 0) + mask = tf.reshape(key, (1,)) + return type(self)._from_tensor(tf.boolean_mask(tensor, mask)) + tensor_shape = self._tensor.shape.as_list() + key_shape = key.shape.as_list() + if any(key_dim == 0 for key_dim in key_shape): + shape = tf.concat( + [tf.constant([0], dtype=tf.int32), tf.shape(self._tensor)[rank:]], + axis=0, + ) + return type(self)._from_tensor(tf.zeros(shape, dtype=self.dtype)) + if any( + tensor_dim is not None and key_dim is not None and tensor_dim != key_dim + for tensor_dim, key_dim in zip( + tensor_shape[:rank], key_shape, strict=True + ) + ): + raise IndexError( + "boolean index shape is incompatible with indexed array" + ) + + shape_assert = tf.debugging.assert_equal( + tf.shape(key), + tf.shape(self._tensor)[:rank], + message="boolean index shape is incompatible with indexed array", + ) + with tf.control_dependencies([shape_assert]): + return type(self)._from_tensor(tf.boolean_mask(self._tensor, key)) + if ( + isinstance(key, tf.Tensor) + and _is_integer_index_tensor(key) + and key.shape.rank != 0 + ): + return type(self)._from_tensor( + tf.gather( + self._tensor, _normalize_integer_index(key, self.shape[0]), axis=0 + ) + ) + scalar = _scalar_integer_getitem(self._tensor, key) + if scalar is not None: + return type(self)._from_tensor(scalar) + advanced = _advanced_integer_getitem(self._tensor, key) + if advanced is not None: + return type(self)._from_tensor(advanced) + return type(self)._from_tensor(self._tensor[key]) + + def __setitem__(self, key: Any, value: Any, /) -> None: + key = _normalize_index_key(key) + value_tensor = _value_to_tensor(value, dtype=self.dtype) + + if isinstance(key, tf.Tensor) and key.dtype == tf.bool: + self._replace_tensor(_boolean_setitem(self._tensor, key, value_tensor)) + return + + variable = tf.Variable(self._tensor) + try: + variable[key].assign(value_tensor) + except Exception as exc: # pragma: no cover - TensorFlow controls details. + raise TypeError( + f"unsupported TensorFlow assignment index: {key!r}" + ) from exc + self._replace_tensor(tf.convert_to_tensor(variable)) + + def _scalar_value(self) -> Any: + if self.ndim != 0: + raise TypeError("only 0-d arrays can be converted to Python scalars") + if not tf.executing_eagerly(): + raise TypeError( + "cannot convert a TensorFlow graph tensor to a Python scalar" + ) + return self._tensor.numpy().item() + + def __bool__(self, /) -> bool: + return bool(self._scalar_value()) + + def __complex__(self, /) -> complex: + return complex(self._scalar_value()) + + def __float__(self, /) -> float: + return float(self._scalar_value()) + + def __index__(self, /) -> int: + return int(self._scalar_value()) + + def __int__(self, /) -> int: + return int(self._scalar_value()) + + def __repr__(self) -> str: + return f"ndtensorflow.asarray({self._tensor!r})" + + def __eq__(self, other: Any) -> Array: # type: ignore[override] + if not _is_supported_operand(other): + return NotImplemented + from deepmd._vendors import ndtensorflow as xp + + return xp.equal(self, other) + + def __ne__(self, other: Any) -> Array: # type: ignore[override] + if not _is_supported_operand(other): + return NotImplemented + from deepmd._vendors import ndtensorflow as xp + + return xp.not_equal(self, other) + + def __abs__(self) -> Array: + from deepmd._vendors import ndtensorflow as xp + + return xp.abs(self) + + def __invert__(self) -> Array: + from deepmd._vendors import ndtensorflow as xp + + return xp.bitwise_invert(self) + + def __neg__(self) -> Array: + from deepmd._vendors import ndtensorflow as xp + + return xp.negative(self) + + def __pos__(self) -> Array: + from deepmd._vendors import ndtensorflow as xp + + return xp.positive(self) + + +def _is_supported_operand(value: Any) -> bool: + return isinstance(value, Array | tf.Tensor | bool | int | float | complex) + + +def _normalize_index_key(key: Any) -> Any: + if isinstance(key, Array): + tensor = key.unwrap() + if tensor.dtype.is_integer and tensor.shape.rank == 0: + return tf.cast(tensor, tf.int64) + return tensor + if isinstance(key, tuple): + return tuple(_normalize_index_key(item) for item in key) + if isinstance(key, slice): + return slice( + _normalize_slice_bound(key.start), + _normalize_slice_bound(key.stop), + _normalize_slice_bound(key.step), + ) + return key + + +def _normalize_slice_bound(value: Any) -> Any: + value = _normalize_index_key(value) + if isinstance(value, tf.Tensor) and value.dtype.is_integer: + return tf.cast(value, tf.int64) + return value + + +def _is_integer_index_tensor(key: tf.Tensor) -> bool: + return key.dtype in (tf.int32, tf.int64) + + +def _normalize_integer_index(index: tf.Tensor, dim: int | None) -> tf.Tensor: + index = tf.cast(index, tf.int64) + if dim is None: + return index + return tf.where(index < 0, index + tf.cast(dim, tf.int64), index) + + +def _advanced_integer_getitem(tensor: tf.Tensor, key: Any) -> tf.Tensor | None: + if not isinstance(key, tuple): + return None + if not any( + isinstance(item, tf.Tensor) and _is_integer_index_tensor(item) for item in key + ): + return None + if all(not isinstance(item, tf.Tensor) or item.shape.rank == 0 for item in key): + return None + if any( + not ( + isinstance(item, int) + or isinstance(item, tf.Tensor) + and _is_integer_index_tensor(item) + ) + for item in key + ): + return None + if len(key) > tensor.shape.rank: + return None + + broadcast_shape = tf.TensorShape(()) + for item in key: + if isinstance(item, tf.Tensor): + broadcast_shape = tf.broadcast_static_shape(broadcast_shape, item.shape) + out_shape = tuple(broadcast_shape.as_list()) + + coords = [] + for axis, item in enumerate(key): + dim = tensor.shape[axis] + if isinstance(item, int): + item = item + dim if item < 0 and dim is not None else item + coord = tf.fill(out_shape, tf.cast(item, tf.int64)) + else: + coord = tf.broadcast_to(_normalize_integer_index(item, dim), out_shape) + coords.append(coord) + indices = tf.stack(coords, axis=-1) + return tf.gather_nd(tensor, indices) + + +def _scalar_integer_getitem(tensor: tf.Tensor, key: Any) -> tf.Tensor | None: + if key == (): + return tensor if tensor.shape.rank == 0 else None + if not isinstance(key, tuple): + return None + if len(key) != tensor.shape.rank: + return None + if tensor.shape.rank <= 7: + return None + if not all(isinstance(item, int) for item in key): + return None + coords = [] + for axis, item in enumerate(key): + dim = tensor.shape[axis] + item = item + dim if item < 0 and dim is not None else item + coords.append(item) + out = tensor + for coord in coords: + out = tf.gather(out, coord, axis=0) + return out + + +def _value_to_tensor(value: Any, dtype: tf.DType) -> tf.Tensor: + if isinstance(value, Array): + return tf.cast(value.unwrap(), dtype) + return tf.convert_to_tensor(value, dtype=dtype) + + +def _boolean_setitem( + tensor: tf.Tensor, + key: tf.Tensor, + value: tf.Tensor, +) -> tf.Tensor: + rank = key.shape.rank + if rank is None: + raise IndexError("boolean index rank must be statically known") + if rank == 0: + return tf.where(key, tf.broadcast_to(value, tf.shape(tensor)), tensor) + tensor_rank = tensor.shape.rank + if tensor_rank is None: + raise IndexError("indexed array rank must be statically known") + if rank > tensor_rank: + raise IndexError("boolean index shape is incompatible with indexed array") + tensor_shape = tensor.shape.as_list() + key_shape = key.shape.as_list() + if any( + tensor_dim is not None and key_dim is not None and tensor_dim != key_dim + for tensor_dim, key_dim in zip(tensor_shape[:rank], key_shape, strict=True) + ): + raise IndexError("boolean index shape is incompatible with indexed array") + + shape_assert = tf.debugging.assert_equal( + tf.shape(key), + tf.shape(tensor)[:rank], + message="boolean index shape is incompatible with indexed array", + ) + with tf.control_dependencies([shape_assert]): + indices = tf.where(key) + updates_shape = tf.concat( + [[tf.shape(indices)[0]], tf.shape(tensor)[rank:]], + axis=0, + ) + updates = tf.broadcast_to(value, updates_shape) + return tf.tensor_scatter_nd_update(tensor, indices, updates) + + +def _binary_forward(name: str) -> Callable[[Array, Any], Array]: + def method(self: Array, other: Any, /) -> Array: + if not _is_supported_operand(other): + return NotImplemented + from deepmd._vendors import ndtensorflow as xp + + return getattr(xp, name)(self, other) + + return method + + +def _binary_reflected(name: str) -> Callable[[Array, Any], Array]: + def method(self: Array, other: Any, /) -> Array: + if not _is_supported_operand(other): + return NotImplemented + from deepmd._vendors import ndtensorflow as xp + + return getattr(xp, name)(other, self) + + return method + + +Array.__add__ = _binary_forward("add") # type: ignore[attr-defined] +Array.__radd__ = _binary_reflected("add") # type: ignore[attr-defined] +Array.__and__ = _binary_forward("bitwise_and") # type: ignore[attr-defined] +Array.__rand__ = _binary_reflected("bitwise_and") # type: ignore[attr-defined] +Array.__floordiv__ = _binary_forward("floor_divide") # type: ignore[attr-defined] +Array.__rfloordiv__ = _binary_reflected("floor_divide") # type: ignore[attr-defined] +Array.__ge__ = _binary_forward("greater_equal") # type: ignore[attr-defined] +Array.__le__ = _binary_reflected("greater_equal") # type: ignore[attr-defined] +Array.__gt__ = _binary_forward("greater") # type: ignore[attr-defined] +Array.__lt__ = _binary_reflected("greater") # type: ignore[attr-defined] +Array.__lshift__ = _binary_forward("bitwise_left_shift") # type: ignore[attr-defined] +Array.__rlshift__ = _binary_reflected("bitwise_left_shift") # type: ignore[attr-defined] +Array.__matmul__ = _binary_forward("matmul") # type: ignore[attr-defined] +Array.__rmatmul__ = _binary_reflected("matmul") # type: ignore[attr-defined] +Array.__mod__ = _binary_forward("remainder") # type: ignore[attr-defined] +Array.__rmod__ = _binary_reflected("remainder") # type: ignore[attr-defined] +Array.__mul__ = _binary_forward("multiply") # type: ignore[attr-defined] +Array.__rmul__ = _binary_reflected("multiply") # type: ignore[attr-defined] +Array.__or__ = _binary_forward("bitwise_or") # type: ignore[attr-defined] +Array.__ror__ = _binary_reflected("bitwise_or") # type: ignore[attr-defined] +Array.__pow__ = _binary_forward("pow") # type: ignore[attr-defined] +Array.__rpow__ = _binary_reflected("pow") # type: ignore[attr-defined] +Array.__rshift__ = _binary_forward("bitwise_right_shift") # type: ignore[attr-defined] +Array.__rrshift__ = _binary_reflected("bitwise_right_shift") # type: ignore[attr-defined] +Array.__sub__ = _binary_forward("subtract") # type: ignore[attr-defined] +Array.__rsub__ = _binary_reflected("subtract") # type: ignore[attr-defined] +Array.__truediv__ = _binary_forward("divide") # type: ignore[attr-defined] +Array.__rtruediv__ = _binary_reflected("divide") # type: ignore[attr-defined] +Array.__xor__ = _binary_forward("bitwise_xor") # type: ignore[attr-defined] +Array.__rxor__ = _binary_reflected("bitwise_xor") # type: ignore[attr-defined] + + +__all__ = ["Array", "DLDeviceType"] diff --git a/deepmd/_vendors/ndtensorflow/_info.py b/deepmd/_vendors/ndtensorflow/_info.py new file mode 100644 index 0000000000..9fd4494397 --- /dev/null +++ b/deepmd/_vendors/ndtensorflow/_info.py @@ -0,0 +1,133 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +from builtins import bool as py_bool +from typing import ( + TypedDict, +) + +import tensorflow as tf + +from ._namespace import bool as bool_dtype +from ._namespace import ( + complex64, + complex128, + float32, + float64, + int8, + int16, + int32, + int64, + isdtype, + uint8, + uint16, + uint32, + uint64, +) + +DefaultDataTypes = TypedDict( + "DefaultDataTypes", + { + "real floating": tf.DType, + "complex floating": tf.DType, + "integral": tf.DType, + "indexing": tf.DType, + }, +) + + +class DataTypes(TypedDict, total=False): + bool: tf.DType + int8: tf.DType + int16: tf.DType + int32: tf.DType + int64: tf.DType + uint8: tf.DType + uint16: tf.DType + uint32: tf.DType + uint64: tf.DType + float32: tf.DType + float64: tf.DType + complex64: tf.DType + complex128: tf.DType + + +Capabilities = TypedDict( + "Capabilities", + { + "boolean indexing": py_bool, + "data-dependent shapes": py_bool, + "max dimensions": int | None, + }, +) + + +def _device_name(device: tf.config.LogicalDevice) -> str: + name = device.name + if name.startswith("/device:"): + return name.removeprefix("/device:") + return name + + +class Info: + """Namespace returned by ``__array_namespace_info__``.""" + + def capabilities(self) -> Capabilities: + return { + "boolean indexing": True, + "data-dependent shapes": True, + "max dimensions": None, + } + + def default_device(self) -> str: + devices = self.devices() + return devices[0] if devices else "CPU:0" + + def default_dtypes(self, *, device: str | None = None) -> DefaultDataTypes: + del device + return { + "real floating": float32, + "complex floating": complex64, + "integral": int32, + "indexing": int64, + } + + def devices(self) -> tuple[str, ...]: + return tuple( + _device_name(device) for device in tf.config.list_logical_devices() + ) + + def dtypes( + self, + *, + device: str | None = None, + kind: None | str | tuple[str, ...] = None, + ) -> DataTypes: + del device + dtypes = { + "bool": bool_dtype, + "int8": int8, + "int16": int16, + "int32": int32, + "int64": int64, + "uint8": uint8, + "uint16": uint16, + "uint32": uint32, + "uint64": uint64, + "float32": float32, + "float64": float64, + "complex64": complex64, + "complex128": complex128, + } + if kind is None: + return dtypes + return {name: dtype for name, dtype in dtypes.items() if isdtype(dtype, kind)} + + +def __array_namespace_info__() -> Info: + return Info() + + +__all__ = ["Info", "__array_namespace_info__"] diff --git a/deepmd/_vendors/ndtensorflow/_namespace.py b/deepmd/_vendors/ndtensorflow/_namespace.py new file mode 100644 index 0000000000..e5ee9bd6db --- /dev/null +++ b/deepmd/_vendors/ndtensorflow/_namespace.py @@ -0,0 +1,2281 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +import math +from builtins import abs as py_abs +from builtins import all as py_all +from builtins import any as py_any +from builtins import bool as py_bool +from builtins import max as py_max +from collections import ( + namedtuple, +) +from collections.abc import ( + Sequence, +) +from contextlib import ( + nullcontext, +) +from functools import reduce as _reduce +from functools import wraps as _wraps +from typing import ( + Any, + Literal, +) + +import tensorflow as tf + +from ._array import ( + Array, +) + +DType = tf.DType +Device = str + +bool = tf.bool +int8 = tf.int8 +int16 = tf.int16 +int32 = tf.int32 +int64 = tf.int64 +uint8 = tf.uint8 +uint16 = tf.uint16 +uint32 = tf.uint32 +uint64 = tf.uint64 +float16 = tf.float16 +bfloat16 = tf.bfloat16 +float32 = tf.float32 +float64 = tf.float64 +complex64 = tf.complex64 +complex128 = tf.complex128 + +newaxis = None +e = math.e +inf = math.inf +nan = float("nan") +pi = math.pi + +UniqueAllResult = namedtuple( + "UniqueAllResult", ["values", "indices", "inverse_indices", "counts"] +) +UniqueCountsResult = namedtuple("UniqueCountsResult", ["values", "counts"]) +UniqueInverseResult = namedtuple("UniqueInverseResult", ["values", "inverse_indices"]) + +_py_scalars = (py_bool, int, float, complex) +_bool_dtypes = {tf.bool} +_signed_dtypes = {tf.int8, tf.int16, tf.int32, tf.int64} +_unsigned_dtypes = {tf.uint8, tf.uint16, tf.uint32, tf.uint64} +_real_floating_dtypes = {tf.float16, tf.bfloat16, tf.float32, tf.float64} +_complex_floating_dtypes = {tf.complex64, tf.complex128} +_integral_dtypes = _signed_dtypes | _unsigned_dtypes +_numeric_dtypes = _integral_dtypes | _real_floating_dtypes | _complex_floating_dtypes +_all_dtypes = _bool_dtypes | _numeric_dtypes + +_dtype_bits = { + tf.bool: 1, + tf.int8: 8, + tf.int16: 16, + tf.int32: 32, + tf.int64: 64, + tf.uint8: 8, + tf.uint16: 16, + tf.uint32: 32, + tf.uint64: 64, + tf.float16: 16, + tf.bfloat16: 16, + tf.float32: 32, + tf.float64: 64, + tf.complex64: 64, + tf.complex128: 128, +} + +_float_for_bits = { + 16: tf.float16, + 32: tf.float32, + 64: tf.float64, +} + +_complex_for_bits = { + 32: tf.complex64, + 64: tf.complex128, +} + +_finfo = { + tf.float16: { + "bits": 16, + "eps": 0.0009765625, + "max": 65504.0, + "min": -65504.0, + "smallest_normal": 0.00006103515625, + }, + tf.bfloat16: { + "bits": 16, + "eps": 0.0078125, + "max": 3.3895313892515355e38, + "min": -3.3895313892515355e38, + "smallest_normal": 1.1754943508222875e-38, + }, + tf.float32: { + "bits": 32, + "eps": 1.1920928955078125e-07, + "max": 3.4028234663852886e38, + "min": -3.4028234663852886e38, + "smallest_normal": 1.1754943508222875e-38, + }, + tf.float64: { + "bits": 64, + "eps": 2.220446049250313e-16, + "max": 1.7976931348623157e308, + "min": -1.7976931348623157e308, + "smallest_normal": 2.2250738585072014e-308, + }, +} + +_iinfo = { + tf.int8: {"bits": 8, "min": -128, "max": 127}, + tf.int16: {"bits": 16, "min": -32768, "max": 32767}, + tf.int32: {"bits": 32, "min": -2147483648, "max": 2147483647}, + tf.int64: {"bits": 64, "min": -9223372036854775808, "max": 9223372036854775807}, + tf.uint8: {"bits": 8, "min": 0, "max": 255}, + tf.uint16: {"bits": 16, "min": 0, "max": 65535}, + tf.uint32: {"bits": 32, "min": 0, "max": 4294967295}, + tf.uint64: {"bits": 64, "min": 0, "max": 18446744073709551615}, +} + + +def _device_context(device: Device | None): + return tf.device(device) if device is not None else nullcontext() + + +def _dlpack_module() -> Any | None: + dlpack = getattr(tf, "dlpack", None) + if dlpack is not None: + return dlpack + experimental = getattr(tf, "experimental", None) + return getattr(experimental, "dlpack", None) + + +def _same_device(requested: Device | None, actual: Device) -> py_bool: + if requested is None: + return True + return requested == actual or str(requested).endswith(str(actual)) + + +def _unwrap(value: Any) -> Any: + if isinstance(value, Array): + return value.unwrap() + return value + + +def _unwrap_nested(value: Any) -> Any: + if isinstance(value, Array): + return value.unwrap() + if isinstance(value, tuple): + return tuple(_unwrap_nested(item) for item in value) + if isinstance(value, list): + return [_unwrap_nested(item) for item in value] + return value + + +def _wrap(value: Any) -> Any: + if isinstance(value, tf.Tensor): + return Array._from_tensor(value) + if isinstance(value, tuple) and hasattr(value, "_fields"): + return type(value)(*(_wrap(item) for item in value)) + if isinstance(value, tuple): + return tuple(_wrap(item) for item in value) + if isinstance(value, list): + return [_wrap(item) for item in value] + return value + + +def _python_scalar_dtype(x: complex) -> DType | None: + if isinstance(x, py_bool): + return tf.bool + if isinstance(x, int): + if _iinfo[tf.int32]["min"] <= x <= _iinfo[tf.int32]["max"]: + return tf.int32 + if _iinfo[tf.int64]["min"] <= x <= _iinfo[tf.int64]["max"]: + return tf.int64 + if 0 <= x <= _iinfo[tf.uint64]["max"]: + return tf.uint64 + return tf.int64 + if isinstance(x, float): + if math.isfinite(x) and py_abs(x) > _finfo[tf.float32]["max"]: + return tf.float64 + return tf.float32 + if isinstance(x, complex): + return tf.complex128 + return None + + +def _real_dtype_for(dtype: DType) -> DType: + dtype = tf.as_dtype(dtype) + if dtype == tf.complex64: + return tf.float32 + if dtype == tf.complex128: + return tf.float64 + return dtype + + +def _is_integer(dtype: DType) -> py_bool: + return dtype in _integral_dtypes + + +def _is_complex(dtype: DType) -> py_bool: + return dtype in _complex_floating_dtypes + + +def _accumulation_dtype(dtype: DType) -> DType: + if dtype in (tf.uint8, tf.uint16, tf.uint32): + return tf.uint32 + if dtype == tf.uint64: + return tf.uint64 + if dtype in (tf.int8, tf.int16, tf.int32): + return tf.int32 + if dtype == tf.int64: + return tf.int64 + return dtype + + +def _normalize_axis(axis: int, ndim: int) -> int: + if axis < 0: + axis += ndim + if axis < 0 or axis >= ndim: + raise IndexError(f"axis {axis} is out of bounds for array of dimension {ndim}") + return axis + + +def _normalize_axes(axis: int | Sequence[int] | None, ndim: int) -> tuple[int, ...]: + if axis is None: + return tuple(range(ndim)) + if isinstance(axis, int): + axis = (axis,) + axes = tuple(_normalize_axis(a, ndim) for a in axis) + if len(set(axes)) != len(axes): + raise ValueError("repeated axis") + return axes + + +def _shape_tuple(x: Array | tf.Tensor) -> tuple[int | tf.Tensor, ...]: + tensor = _unwrap(x) + static_shape = tensor.shape.as_list() + if py_all(dim is not None for dim in static_shape): + return tuple(static_shape) + dynamic_shape = tf.shape(tensor) + return tuple( + dynamic_shape[ii] if dim is None else dim for ii, dim in enumerate(static_shape) + ) + + +def _normalize_shape_arg(shape: int | Sequence[Any] | tf.Tensor) -> Any: + if isinstance(shape, Array): + return shape.unwrap() + if isinstance(shape, tf.Tensor): + return shape + if isinstance(shape, int): + return (shape,) + return tuple(_unwrap(dim) for dim in shape) + + +def _shape_arg_tensor( + shape: int | Sequence[Any] | tf.Tensor, + dtype: DType = tf.int32, +) -> tf.Tensor: + shape = _normalize_shape_arg(shape) + if isinstance(shape, tf.Tensor): + return tf.cast(shape, dtype) + return tf.stack( + [ + tf.cast(dim, dtype) + if isinstance(dim, tf.Tensor) + else tf.constant(dim, dtype) + for dim in shape + ] + ) + + +def _shape_arg_for_tf(shape: int | Sequence[Any] | tf.Tensor) -> Any: + shape = _normalize_shape_arg(shape) + if isinstance(shape, tuple) and py_any(isinstance(dim, tf.Tensor) for dim in shape): + return _shape_arg_tensor(shape) + return shape + + +def _shape_product(shape: Sequence[int | tf.Tensor]) -> int | tf.Tensor: + if not shape: + return 1 + if py_all(isinstance(dim, int) for dim in shape): + return math.prod(shape) + return tf.reduce_prod(_shape_arg_tensor(shape)) + + +def _dtype_of(x: Array | tf.Tensor | DType | complex) -> DType: + if isinstance(x, tf.DType): + return x + if isinstance(x, Array): + return x.dtype + if isinstance(x, tf.Tensor): + return x.dtype + dtype = _python_scalar_dtype(x) + if dtype is not None: + return dtype + return tf.convert_to_tensor(_unwrap_nested(x)).dtype + + +def _promote_signed_unsigned(signed: DType, unsigned: DType) -> DType: + signed_bits = _dtype_bits[signed] + unsigned_bits = _dtype_bits[unsigned] + for dtype in (tf.int16, tf.int32, tf.int64): + bits = _dtype_bits[dtype] + if bits >= signed_bits and bits > unsigned_bits: + return dtype + return tf.float64 + + +def _promote_dtypes(dtype1: DType, dtype2: DType) -> DType: + dtype1 = tf.as_dtype(dtype1) + dtype2 = tf.as_dtype(dtype2) + if dtype1 == dtype2: + return dtype1 + if dtype1 == tf.bool: + return dtype2 + if dtype2 == tf.bool: + return dtype1 + if dtype1 in _complex_floating_dtypes or dtype2 in _complex_floating_dtypes: + bits = py_max( + _dtype_bits[_real_dtype_for(dtype1)], + _dtype_bits[_real_dtype_for(dtype2)], + ) + return _complex_for_bits[bits] + if dtype1 in _real_floating_dtypes or dtype2 in _real_floating_dtypes: + bits = py_max( + _dtype_bits[dtype1] if dtype1 in _real_floating_dtypes else 0, + _dtype_bits[dtype2] if dtype2 in _real_floating_dtypes else 0, + ) + return _float_for_bits[py_max(bits, 32)] + if dtype1 in _signed_dtypes and dtype2 in _signed_dtypes: + return dtype1 if _dtype_bits[dtype1] >= _dtype_bits[dtype2] else dtype2 + if dtype1 in _unsigned_dtypes and dtype2 in _unsigned_dtypes: + return dtype1 if _dtype_bits[dtype1] >= _dtype_bits[dtype2] else dtype2 + if dtype1 in _signed_dtypes and dtype2 in _unsigned_dtypes: + return _promote_signed_unsigned(dtype1, dtype2) + if dtype1 in _unsigned_dtypes and dtype2 in _signed_dtypes: + return _promote_signed_unsigned(dtype2, dtype1) + raise TypeError(f"Cannot promote {dtype1!r} and {dtype2!r}") + + +def _promote_scalar(dtype: DType, scalar: complex) -> DType: + if isinstance(scalar, py_bool): + return dtype if dtype != tf.bool else tf.bool + if isinstance(scalar, int): + return dtype if dtype in _numeric_dtypes else _promote_dtypes(dtype, tf.int32) + if isinstance(scalar, float): + if dtype in _real_floating_dtypes | _complex_floating_dtypes: + return dtype + return _promote_dtypes(dtype, tf.float64) + if isinstance(scalar, complex): + if dtype in _complex_floating_dtypes: + return dtype + return _promote_dtypes(dtype, tf.complex128) + return _promote_dtypes(dtype, _dtype_of(scalar)) + + +def _iter_nested_scalars(obj: Any): + obj = _unwrap(obj) + if isinstance(obj, tf.Tensor): + return + if isinstance(obj, Sequence) and not isinstance( + obj, str | bytes | bytearray | memoryview + ): + for item in obj: + yield from _iter_nested_scalars(item) + else: + yield obj + + +def _infer_nested_dtype(obj: Any) -> DType | None: + dtypes = [] + for scalar in _iter_nested_scalars(obj): + if scalar is None: + return None + dtypes.append(_dtype_of(scalar)) + if not dtypes: + return None + return _reduce(_promote_dtypes, dtypes) + + +def _coerce_scalar_to_dtype(obj: Any, dtype: DType) -> Any: + if dtype == tf.bool: + return py_bool(obj) + if dtype in _integral_dtypes: + return int(obj) + if dtype in _real_floating_dtypes: + return float(obj) + if dtype in _complex_floating_dtypes: + return complex(obj) + return obj + + +def _coerce_nested_to_dtype(obj: Any, dtype: DType) -> Any: + obj = _unwrap(obj) + if isinstance(obj, tf.Tensor): + return obj + if hasattr(obj, "shape") and hasattr(obj, "dtype"): + try: + return tf.convert_to_tensor(obj, dtype=dtype) + except (TypeError, ValueError): + pass + if isinstance(obj, Sequence) and not isinstance( + obj, str | bytes | bytearray | memoryview + ): + return [_coerce_nested_to_dtype(item, dtype) for item in obj] + return _coerce_scalar_to_dtype(obj, dtype) + + +def _negative_zero(dtype: DType) -> tf.Tensor: + dtype = tf.as_dtype(dtype) + if dtype == tf.float16: + return tf.bitcast(tf.constant(0x8000, dtype=tf.uint16), tf.float16) + if dtype == tf.bfloat16: + return tf.bitcast(tf.constant(0x8000, dtype=tf.uint16), tf.bfloat16) + if dtype == tf.float32: + return tf.bitcast(tf.constant(0x80000000, dtype=tf.uint32), tf.float32) + if dtype == tf.float64: + return tf.bitcast(tf.constant(0x8000000000000000, dtype=tf.uint64), tf.float64) + raise TypeError(f"{dtype!r} is not a real floating dtype") + + +def _python_scalar_to_tensor(x: Any, dtype: DType | None) -> tf.Tensor | None: + if dtype is None: + dtype = _python_scalar_dtype(x) + if dtype is None: + return None + dtype = tf.as_dtype(dtype) + if isinstance(x, float) and dtype in _real_floating_dtypes and x == 0.0: + if math.copysign(1.0, x) < 0: + return _negative_zero(dtype) + return tf.zeros((), dtype=dtype) + if isinstance(x, complex) and dtype in _complex_floating_dtypes: + real_dtype = _real_dtype_for(dtype) + real_part = _python_scalar_to_tensor(x.real, real_dtype) + imag_part = _python_scalar_to_tensor(x.imag, real_dtype) + if real_part is None: + real_part = tf.convert_to_tensor(x.real, dtype=real_dtype) + if imag_part is None: + imag_part = tf.convert_to_tensor(x.imag, dtype=real_dtype) + return tf.complex(real_part, imag_part) + return None + + +def _astype_tensor(x: tf.Tensor, dtype: DType, copy: py_bool = False) -> tf.Tensor: + if x.dtype == dtype: + return tf.identity(x) if copy else x + return tf.cast(x, dtype) + + +def _to_tensor(x: Array | tf.Tensor | complex, dtype: DType | None = None) -> tf.Tensor: + x = _unwrap(x) + if isinstance(x, tf.Tensor): + return _astype_tensor(x, dtype) if dtype is not None else x + out = _python_scalar_to_tensor(x, dtype) + if out is not None: + return out + return tf.convert_to_tensor(_unwrap_nested(x), dtype=dtype) + + +def result_type(*arrays_and_dtypes: Array | tf.Tensor | DType | complex) -> DType: + if not arrays_and_dtypes: + raise ValueError("At least one array or dtype must be provided") + if py_all(isinstance(x, _py_scalars) for x in arrays_and_dtypes): + raise ValueError("At least one array or dtype must be provided") + scalars = [] + others = [] + for x in arrays_and_dtypes: + if isinstance(x, _py_scalars): + scalars.append(x) + else: + others.append(x) + dtype = _dtype_of(others[0]) + for other in others[1:]: + dtype = _promote_dtypes(dtype, _dtype_of(other)) + for scalar in scalars: + dtype = _promote_scalar(dtype, scalar) + return dtype + + +def _result_type_with_scalars(x1: Any, x2: Any) -> DType: + x1_is_scalar = isinstance(x1, _py_scalars) + x2_is_scalar = isinstance(x2, _py_scalars) + if x1_is_scalar and x2_is_scalar: + return _promote_dtypes(_dtype_of(x1), _dtype_of(x2)) + if x1_is_scalar: + return _promote_scalar(_dtype_of(x2), x1) + if x2_is_scalar: + return _promote_scalar(_dtype_of(x1), x2) + return result_type(x1, x2) + + +def _known_unequal(x1: int | None, x2: int | None) -> py_bool: + return x1 is not None and x2 is not None and x1 != x2 + + +def _promote_two( + x1: Array | tf.Tensor | complex, x2: Array | tf.Tensor | complex +) -> tuple[tf.Tensor, tf.Tensor]: + dtype = result_type(x1, x2) + return _to_tensor(x1, dtype), _to_tensor(x2, dtype) + + +def _two_arg(f): + @_wraps(f) + def _f(x1: Any, x2: Any, /, **kwargs: Any) -> Array: + x1, x2 = _promote_two(x1, x2) + return Array._from_tensor(f(x1, x2, **kwargs)) + + return _f + + +def _logical_two_arg(f): + @_wraps(f) + def _f(x1: Any, x2: Any, /, **kwargs: Any) -> Array: + return Array._from_tensor( + f( + tf.cast(_unwrap(x1), tf.bool), + tf.cast(_unwrap(x2), tf.bool), + **kwargs, + ) + ) + + return _f + + +def _unary(f): + @_wraps(f) + def _f(x: Array, /, **kwargs: Any) -> Array: + return Array._from_tensor(f(_unwrap(x), **kwargs)) + + return _f + + +def _signed_zero_like(x: tf.Tensor) -> tf.Tensor: + neg_zero = tf.broadcast_to(_negative_zero(x.dtype), tf.shape(x)) + return tf.where(signbit(Array._from_tensor(x)).unwrap(), neg_zero, tf.zeros_like(x)) + + +def _moveaxis_permutation( + ndim: int, + source: int | Sequence[int], + destination: int | Sequence[int], +) -> list[int]: + if isinstance(source, int): + source = (source,) + if isinstance(destination, int): + destination = (destination,) + if len(source) != len(destination): + raise ValueError("`source` and `destination` must have the same number of axes") + source_ = tuple(_normalize_axis(axis, ndim) for axis in source) + destination_ = tuple(_normalize_axis(axis, ndim) for axis in destination) + if len(set(source_)) != len(source_) or len(set(destination_)) != len(destination_): + raise ValueError("repeated axis") + order = [axis for axis in range(ndim) if axis not in source_] + for dest, src in sorted(zip(destination_, source_, strict=True)): + order.insert(dest, src) + return order + + +def _moveaxis( + x: Array | tf.Tensor, + source: int | Sequence[int], + destination: int | Sequence[int], +) -> tf.Tensor: + tensor = _unwrap(x) + return tf.transpose( + tensor, _moveaxis_permutation(tensor.shape.rank, source, destination) + ) + + +def asarray( + obj: Any, + /, + *, + dtype: DType | None = None, + device: Device | None = None, + copy: py_bool | None = None, + **kwargs: Any, +) -> Array: + if copy is False and not isinstance(obj, Array | tf.Tensor): + raise ValueError("Unable to avoid copy while creating a TensorFlow tensor") + with _device_context(device): + if isinstance(obj, Array): + same_dtype = dtype is None or obj.dtype == dtype + if copy is False: + if not same_dtype or not _same_device(device, obj.device): + raise ValueError("Unable to avoid copy while converting an Array") + return obj + tensor = obj.unwrap() if same_dtype else tf.cast(obj.unwrap(), dtype) + if device is not None or copy is True: + tensor = tf.identity(tensor) + return Array._from_tensor(tensor) + if isinstance(obj, tf.Tensor): + tensor = _unwrap(obj) + same_dtype = dtype is None or tensor.dtype == dtype + if copy is False and not same_dtype: + raise ValueError("Unable to avoid copy while converting dtype") + out = tensor if same_dtype else tf.cast(tensor, dtype) + if device is not None or copy is True: + out = tf.identity(out) + return Array._from_tensor(out) + try: + if dtype is None: + dtype = _infer_nested_dtype(obj) + out = _python_scalar_to_tensor(obj, dtype) + if out is None: + obj_ = ( + _coerce_nested_to_dtype(obj, dtype) + if dtype is not None + else _unwrap_nested(obj) + ) + out = tf.convert_to_tensor(obj_, dtype=dtype, **kwargs) + except (TypeError, ValueError): + obj_ = list(obj) + if dtype is not None: + obj_ = _coerce_nested_to_dtype(obj_, dtype) + out = tf.convert_to_tensor(obj_, dtype=dtype, **kwargs) + return Array._from_tensor(tf.identity(out) if copy is True else out) + + +def astype( + x: Array, + dtype: DType, + /, + *, + copy: py_bool = True, + device: Device | None = None, +) -> Array: + with _device_context(device): + return Array._from_tensor( + _astype_tensor(_unwrap(x), dtype, copy=copy or device is not None) + ) + + +def from_dlpack( + x: Any, + /, + *, + device: Device | None = None, + copy: py_bool | None = None, +) -> Array: + if isinstance(x, Array | tf.Tensor): + return asarray(x, device=device, copy=copy) + dlpack = _dlpack_module() + if dlpack is None: + raise BufferError("TensorFlow DLPack import is not available") + capsule = x.__dlpack__() if hasattr(x, "__dlpack__") else x + with _device_context(device): + out = dlpack.from_dlpack(capsule) + if device is not None or copy is True: + out = tf.identity(out) + return Array._from_tensor(out) + + +def arange( + start: float, + /, + stop: float | None = None, + step: float = 1, + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + if stop is None: + start, stop = 0, start + if ( + isinstance(_unwrap(start), tf.Tensor) + or isinstance(_unwrap(stop), tf.Tensor) + or isinstance(_unwrap(step), tf.Tensor) + ): + if dtype is None: + dtype = _dtype_of(stop) + if dtype not in _numeric_dtypes: + dtype = tf.int32 + start_ = tf.cast(_to_tensor(start), dtype) + stop_ = tf.cast(_to_tensor(stop), dtype) + step_ = tf.cast(_to_tensor(step), dtype) + with _device_context(device): + return Array._from_tensor(tf.range(start_, stop_, step_, dtype=dtype)) + with _device_context(device): + if step > 0 and stop <= start or step < 0 and stop >= start: + if dtype is None: + dtype = ( + tf.int32 + if py_all(isinstance(i, int) for i in (start, stop, step)) + else tf.float32 + ) + return Array._from_tensor(tf.zeros((0,), dtype=dtype)) + if dtype is None: + if py_all(isinstance(i, int) for i in (start, stop, step)): + return Array._from_tensor(tf.range(start, stop, step, dtype=tf.int32)) + return Array._from_tensor( + tf.cast(tf.range(start, stop, step, dtype=tf.float64), tf.float32) + ) + work_dtype = tf.int64 if dtype in _integral_dtypes else tf.float64 + return Array._from_tensor( + tf.cast(tf.range(start, stop, step, dtype=work_dtype), dtype) + ) + + +def empty( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + if isinstance(shape, int): + shape = (shape,) + with _device_context(device): + return Array._from_tensor( + tf.zeros(_shape_arg_for_tf(shape), dtype=dtype or tf.float32) + ) + + +def empty_like( + x: Array, + /, + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + with _device_context(device): + return Array._from_tensor(tf.zeros_like(_unwrap(x), dtype=dtype)) + + +def eye( + n_rows: int, + n_cols: int | None = None, + /, + *, + k: int = 0, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + if n_cols is None: + n_cols = n_rows + with _device_context(device): + if k >= n_cols or k <= -n_rows: + return Array._from_tensor( + tf.zeros((n_rows, n_cols), dtype=dtype or tf.float32) + ) + rows = tf.range(n_rows, dtype=tf.int32)[:, newaxis] + cols = tf.range(n_cols, dtype=tf.int32)[newaxis, :] + return Array._from_tensor(tf.cast(cols - rows == k, dtype or tf.float32)) + + +def full( + shape: int | tuple[int, ...], + fill_value: complex, + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + if isinstance(shape, int): + shape = (shape,) + with _device_context(device): + value = _to_tensor(fill_value, dtype=dtype) + return Array._from_tensor(tf.broadcast_to(value, _shape_arg_for_tf(shape))) + + +def full_like( + x: Array, + /, + fill_value: complex, + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + return full(_shape_tuple(x), fill_value, dtype=dtype or x.dtype, device=device) + + +def linspace( + start: float, + stop: float, + /, + num: int, + *, + dtype: DType | None = None, + device: Device | None = None, + endpoint: py_bool = True, + **kwargs: object, +) -> Array: + del kwargs + with _device_context(device): + if num == 0: + return Array._from_tensor(tf.zeros((0,), dtype=dtype or tf.float32)) + out_dtype = dtype or tf.float32 + work_dtype = ( + out_dtype + if out_dtype in _real_floating_dtypes | _complex_floating_dtypes + else tf.float32 + ) + start_ = tf.convert_to_tensor(start, dtype=work_dtype) + stop_ = tf.convert_to_tensor(stop, dtype=work_dtype) + # pylint: disable-next=no-explicit-dtype + out = tf.linspace(start_, stop_, num if endpoint else num + 1) + if not endpoint: + out = out[:-1] + return Array._from_tensor(tf.cast(out, out_dtype)) + + +def ones( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + with _device_context(device): + return Array._from_tensor( + tf.ones(_shape_arg_for_tf(shape), dtype=dtype or tf.float32) + ) + + +def ones_like( + x: Array, + /, + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + with _device_context(device): + return Array._from_tensor(tf.ones_like(_unwrap(x), dtype=dtype)) + + +def zeros( + shape: int | tuple[int, ...], + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + with _device_context(device): + return Array._from_tensor( + tf.zeros(_shape_arg_for_tf(shape), dtype=dtype or tf.float32) + ) + + +def zeros_like( + x: Array, + /, + *, + dtype: DType | None = None, + device: Device | None = None, + **kwargs: object, +) -> Array: + del kwargs + with _device_context(device): + return Array._from_tensor(tf.zeros_like(_unwrap(x), dtype=dtype)) + + +def tril(x: Array, /, *, k: int = 0) -> Array: + tensor = _unwrap(x) + rows = tf.range(tensor.shape[-2], dtype=tf.int32)[:, newaxis] + cols = tf.range(tensor.shape[-1], dtype=tf.int32)[newaxis, :] + return Array._from_tensor( + tf.where(cols - rows <= k, tensor, tf.zeros((), dtype=tensor.dtype)) + ) + + +def triu(x: Array, /, *, k: int = 0) -> Array: + tensor = _unwrap(x) + rows = tf.range(tensor.shape[-2], dtype=tf.int32)[:, newaxis] + cols = tf.range(tensor.shape[-1], dtype=tf.int32)[newaxis, :] + return Array._from_tensor( + tf.where(cols - rows >= k, tensor, tf.zeros((), dtype=tensor.dtype)) + ) + + +def can_cast(from_: DType | Array, to: DType, /) -> py_bool: + from_dtype = _dtype_of(from_) + to = tf.as_dtype(to) + if from_dtype == to: + return True + if from_dtype == tf.bool: + return to == tf.bool + if from_dtype in _signed_dtypes: + if to in _signed_dtypes: + return _dtype_bits[from_dtype] <= _dtype_bits[to] + if to in _real_floating_dtypes: + return _dtype_bits[to] > _dtype_bits[from_dtype] + if to in _complex_floating_dtypes: + return _dtype_bits[to] // 2 > _dtype_bits[from_dtype] + return False + if from_dtype in _unsigned_dtypes: + if to in _unsigned_dtypes: + return _dtype_bits[from_dtype] <= _dtype_bits[to] + if to in _signed_dtypes: + return _dtype_bits[from_dtype] < _dtype_bits[to] + if to in _real_floating_dtypes: + return _dtype_bits[to] > _dtype_bits[from_dtype] + if to in _complex_floating_dtypes: + return _dtype_bits[to] // 2 > _dtype_bits[from_dtype] + return False + if from_dtype in _real_floating_dtypes: + if to in _real_floating_dtypes: + return _dtype_bits[from_dtype] <= _dtype_bits[to] + if to in _complex_floating_dtypes: + return _dtype_bits[from_dtype] <= _dtype_bits[to] // 2 + return False + return ( + from_dtype in _complex_floating_dtypes + and to in _complex_floating_dtypes + and _dtype_bits[from_dtype] <= _dtype_bits[to] + ) + + +def isdtype( + dtype: DType, + kind: DType | str | tuple[DType | str, ...], + *, + _tuple: py_bool = True, +) -> py_bool: + dtype = tf.as_dtype(dtype) + if isinstance(kind, tuple) and _tuple: + return py_any(isdtype(dtype, k, _tuple=False) for k in kind) + if isinstance(kind, str): + if kind == "bool": + return dtype in _bool_dtypes + if kind == "signed integer": + return dtype in _signed_dtypes + if kind == "unsigned integer": + return dtype in _unsigned_dtypes + if kind == "integral": + return dtype in _integral_dtypes + if kind == "real floating": + return dtype in _real_floating_dtypes + if kind == "complex floating": + return dtype in _complex_floating_dtypes + if kind == "numeric": + return dtype in _numeric_dtypes + raise ValueError(f"Unrecognized data type kind: {kind!r}") + return dtype == tf.as_dtype(kind) + + +class _FInfo: + def __init__(self, dtype: DType): + real_dtype = _real_dtype_for(dtype) + info = _finfo[real_dtype] + self.bits = info["bits"] + self.eps = info["eps"] + self.max = info["max"] + self.min = info["min"] + self.smallest_normal = info["smallest_normal"] + self.dtype = real_dtype + + +class _IInfo: + def __init__(self, dtype: DType): + info = _iinfo[dtype] + self.bits = info["bits"] + self.max = info["max"] + self.min = info["min"] + self.dtype = dtype + + +def finfo(type_: DType | Array, /) -> _FInfo: + return _FInfo(_dtype_of(type_)) + + +def iinfo(type_: DType | Array, /) -> _IInfo: + return _IInfo(_dtype_of(type_)) + + +def abs(x: Array, /) -> Array: + tensor = _unwrap(x) + if tensor.dtype in _unsigned_dtypes or tensor.dtype == tf.bool: + return Array._from_tensor(tf.identity(tensor)) + return Array._from_tensor(tf.abs(tensor)) + + +acos = _unary(tf.acos) +acosh = _unary(tf.acosh) +asin = _unary(tf.asin) +asinh = _unary(tf.asinh) +atan = _unary(tf.atan) +atan2 = _two_arg(tf.atan2) +atanh = _unary(tf.atanh) +add = _two_arg(tf.add) +conj = _unary(tf.math.conj) +cos = _unary(tf.cos) +cosh = _unary(tf.cosh) +divide = _two_arg(tf.divide) +equal = _two_arg(tf.equal) +exp = _unary(tf.exp) +greater = _two_arg(tf.greater) +greater_equal = _two_arg(tf.greater_equal) +less = _two_arg(tf.less) +less_equal = _two_arg(tf.less_equal) +logical_and = _logical_two_arg(tf.logical_and) +logical_not = _unary(tf.logical_not) +logical_or = _logical_two_arg(tf.logical_or) +logical_xor = _logical_two_arg(tf.math.logical_xor) +maximum = _two_arg(tf.maximum) +minimum = _two_arg(tf.minimum) +multiply = _two_arg(tf.multiply) +not_equal = _two_arg(tf.not_equal) +positive = _unary(tf.identity) +sin = _unary(tf.sin) +sinh = _unary(tf.sinh) +square = _unary(tf.square) +sqrt = _unary(tf.sqrt) +subtract = _two_arg(tf.subtract) +tan = _unary(tf.tan) + + +def expm1(x: Array, /) -> Array: + tensor = _unwrap(x) + out = tf.math.expm1(tensor) + if _is_complex(tensor.dtype): + real_part = tf.math.real(tensor) + imag_part = tf.math.imag(tensor) + + plus_inf_real = tf.math.is_inf(real_part) & (real_part > 0) + exp_out = tf.exp(tensor) - tf.cast(1, tensor.dtype) + inf_real = tf.fill(tf.shape(real_part), tf.cast(math.inf, real_part.dtype)) + zero_imag_out = tf.complex(inf_real, imag_part) + out = tf.where( + plus_inf_real, tf.where(imag_part == 0, zero_imag_out, exp_out), out + ) + + minus_inf_real = tf.math.is_inf(real_part) & (real_part < 0) + minus_inf_imag = tf.where( + tf.math.is_nan(imag_part), + tf.zeros_like(imag_part), + _signed_zero_like(imag_part), + ) + out = tf.where( + minus_inf_real, tf.complex(-tf.ones_like(real_part), minus_inf_imag), out + ) + + nan_real_zero_imag = tf.math.is_nan(real_part) & (imag_part == 0) + out = tf.where(nan_real_zero_imag, tf.complex(real_part, imag_part), out) + + zero = tf.zeros((), dtype=real_part.dtype) + zero_out = tf.complex(tf.zeros_like(real_part), imag_part) + out = tf.where((real_part == zero) & (imag_part == zero), zero_out, out) + return Array._from_tensor(out) + + +def tanh(x: Array, /) -> Array: + tensor = _unwrap(x) + out = tf.math.tanh(tensor) + if tensor.dtype in _real_floating_dtypes | _complex_floating_dtypes: + out = tf.where(tensor == tf.zeros((), dtype=tensor.dtype), tensor, out) + if _is_complex(tensor.dtype): + real_part = tf.math.real(tensor) + imag_part = tf.math.imag(tensor) + inf_real = tf.math.is_inf(real_part) + real_out = tf.where( + real_part > 0, tf.ones_like(real_part), -tf.ones_like(real_part) + ) + imag_out = tf.where( + tf.math.is_nan(imag_part), + tf.zeros_like(imag_part), + _signed_zero_like(imag_part), + ) + out = tf.where(inf_real, tf.complex(real_out, imag_out), out) + return Array._from_tensor(out) + + +def remainder(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + out = tf.math.floormod(x1, x2) + if out.dtype in _real_floating_dtypes: + signed_zero = tf.broadcast_to(_signed_zero_like(x2), tf.shape(out)) + out = tf.where(out == 0, signed_zero, out) + return Array._from_tensor(out) + + +def bitwise_and(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + if x1.dtype == tf.bool: + return Array._from_tensor(tf.logical_and(x1, x2)) + return Array._from_tensor(tf.bitwise.bitwise_and(x1, x2)) + + +def bitwise_left_shift(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + out = tf.bitwise.left_shift(x1, x2) + return Array._from_tensor( + tf.where( + x2 >= tf.cast(_dtype_bits[x1.dtype], x2.dtype), + tf.zeros((), dtype=x1.dtype), + out, + ) + ) + + +def bitwise_invert(x: Array, /) -> Array: + tensor = _unwrap(x) + if tensor.dtype == tf.bool: + return Array._from_tensor(tf.logical_not(tensor)) + return Array._from_tensor(tf.bitwise.invert(tensor)) + + +def bitwise_or(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + if x1.dtype == tf.bool: + return Array._from_tensor(tf.logical_or(x1, x2)) + return Array._from_tensor(tf.bitwise.bitwise_or(x1, x2)) + + +def bitwise_right_shift(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + return Array._from_tensor(tf.bitwise.right_shift(x1, x2)) + + +def bitwise_xor(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + if x1.dtype == tf.bool: + return Array._from_tensor(tf.math.logical_xor(x1, x2)) + return Array._from_tensor(tf.bitwise.bitwise_xor(x1, x2)) + + +def ceil(x: Array, /) -> Array: + tensor = _unwrap(x) + return Array._from_tensor( + tf.identity(tensor) if _is_integer(tensor.dtype) else tf.math.ceil(tensor) + ) + + +def floor(x: Array, /) -> Array: + tensor = _unwrap(x) + return Array._from_tensor( + tf.identity(tensor) if _is_integer(tensor.dtype) else tf.math.floor(tensor) + ) + + +def trunc(x: Array, /) -> Array: + tensor = _unwrap(x) + if _is_integer(tensor.dtype): + return Array._from_tensor(tf.identity(tensor)) + return Array._from_tensor( + tf.where(tensor < 0, tf.math.ceil(tensor), tf.math.floor(tensor)) + ) + + +def copysign(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + return Array._from_tensor( + tf.where(signbit(Array._from_tensor(x2)).unwrap(), -tf.abs(x1), tf.abs(x1)) + ) + + +def hypot(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + return Array._from_tensor(tf.sqrt(tf.square(x1) + tf.square(x2))) + + +def imag(x: Array, /) -> Array: + tensor = _unwrap(x) + if _is_complex(tensor.dtype): + return Array._from_tensor(tf.math.imag(tensor)) + return Array._from_tensor(tf.zeros_like(tensor)) + + +def isfinite(x: Array, /) -> Array: + tensor = _unwrap(x) + if _is_integer(tensor.dtype) or tensor.dtype == tf.bool: + return Array._from_tensor(tf.ones(_shape_tuple(tensor), dtype=tf.bool)) + if _is_complex(tensor.dtype): + return Array._from_tensor( + tf.math.is_finite(tf.math.real(tensor)) + & tf.math.is_finite(tf.math.imag(tensor)) + ) + return Array._from_tensor(tf.math.is_finite(tensor)) + + +def isinf(x: Array, /) -> Array: + tensor = _unwrap(x) + if _is_integer(tensor.dtype) or tensor.dtype == tf.bool: + return Array._from_tensor(tf.zeros(_shape_tuple(tensor), dtype=tf.bool)) + if _is_complex(tensor.dtype): + return Array._from_tensor( + tf.math.is_inf(tf.math.real(tensor)) | tf.math.is_inf(tf.math.imag(tensor)) + ) + return Array._from_tensor(tf.math.is_inf(tensor)) + + +def isnan(x: Array, /) -> Array: + tensor = _unwrap(x) + if _is_integer(tensor.dtype) or tensor.dtype == tf.bool: + return Array._from_tensor(tf.zeros(_shape_tuple(tensor), dtype=tf.bool)) + if _is_complex(tensor.dtype): + return Array._from_tensor( + tf.math.is_nan(tf.math.real(tensor)) | tf.math.is_nan(tf.math.imag(tensor)) + ) + return Array._from_tensor(tf.math.is_nan(tensor)) + + +def _complex_log(x: tf.Tensor) -> tf.Tensor: + return tf.complex( + tf.math.log(tf.abs(x)), tf.atan2(tf.math.imag(x), tf.math.real(x)) + ) + + +def log(x: Array, /) -> Array: + tensor = _unwrap(x) + return Array._from_tensor( + _complex_log(tensor) if _is_complex(tensor.dtype) else tf.math.log(tensor) + ) + + +def floor_divide(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + return Array._from_tensor(tf.math.floordiv(x1, x2)) + + +def log1p(x: Array, /) -> Array: + tensor = _unwrap(x) + if _is_complex(tensor.dtype): + return log(Array._from_tensor(tf.cast(1, tensor.dtype) + tensor)) + return Array._from_tensor(tf.math.log1p(tensor)) + + +def log2(x: Array, /) -> Array: + tensor = _unwrap(x) + out = _complex_log(tensor) if _is_complex(tensor.dtype) else tf.math.log(tensor) + return Array._from_tensor(out / tf.cast(math.log(2.0), tensor.dtype)) + + +def log10(x: Array, /) -> Array: + tensor = _unwrap(x) + out = _complex_log(tensor) if _is_complex(tensor.dtype) else tf.math.log(tensor) + return Array._from_tensor(out / tf.cast(math.log(10.0), tensor.dtype)) + + +def negative(x: Array, /) -> Array: + tensor = _unwrap(x) + if tensor.dtype in _unsigned_dtypes: + return Array._from_tensor(tf.cast(-tf.cast(tensor, tf.int64), tensor.dtype)) + return Array._from_tensor(tf.negative(tensor)) + + +def pow(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + out_dtype = x1.dtype + work_dtype = tf.int64 if out_dtype in _integral_dtypes else out_dtype + return Array._from_tensor( + tf.cast(tf.pow(tf.cast(x1, work_dtype), tf.cast(x2, work_dtype)), out_dtype) + ) + + +def logaddexp(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + shape = tuple(tf.broadcast_static_shape(x1.shape, x2.shape).as_list()) + x1 = tf.broadcast_to(x1, shape) + x2 = tf.broadcast_to(x2, shape) + return Array._from_tensor(tf.reduce_logsumexp(tf.stack([x1, x2]), axis=0)) + + +def nextafter(x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + return Array._from_tensor(tf.math.nextafter(x1, x2)) + + +def real(x: Array, /) -> Array: + tensor = _unwrap(x) + if _is_complex(tensor.dtype): + return Array._from_tensor(tf.math.real(tensor)) + return Array._from_tensor(tf.identity(tensor)) + + +def reciprocal(x: Array, /) -> Array: + tensor = _unwrap(x) + return Array._from_tensor(tf.math.reciprocal(tensor)) + + +def round(x: Array, /, *, decimals: int = 0) -> Array: + tensor = _unwrap(x) + if tensor.dtype in _integral_dtypes: + return Array._from_tensor(tf.identity(tensor)) + if _is_complex(tensor.dtype): + real_part = round( + Array._from_tensor(tf.math.real(tensor)), decimals=decimals + ).unwrap() + imag_part = round( + Array._from_tensor(tf.math.imag(tensor)), decimals=decimals + ).unwrap() + return Array._from_tensor(tf.complex(real_part, imag_part)) + if decimals == 0: + return Array._from_tensor(tf.round(tensor)) + factor = tf.cast(10**decimals, tensor.dtype) + return Array._from_tensor(tf.round(tensor * factor) / factor) + + +def sign(x: Array, /) -> Array: + tensor = _unwrap(x) + if tensor.dtype in _unsigned_dtypes: + return Array._from_tensor( + tf.where( + tensor == 0, + tf.zeros((), dtype=tensor.dtype), + tf.ones((), dtype=tensor.dtype), + ) + ) + return Array._from_tensor(tf.sign(tensor)) + + +def signbit(x: Array, /) -> Array: + tensor = _unwrap(x) + if tensor.dtype in _unsigned_dtypes or tensor.dtype == tf.bool: + return Array._from_tensor(tf.zeros(_shape_tuple(tensor), dtype=tf.bool)) + if tensor.dtype in _integral_dtypes: + return Array._from_tensor(tensor < 0) + if tensor.dtype == tf.float16: + return Array._from_tensor( + tf.bitwise.right_shift(tf.bitcast(tensor, tf.uint16), 15) == 1 + ) + if tensor.dtype == tf.bfloat16: + return Array._from_tensor( + tf.bitwise.right_shift(tf.bitcast(tensor, tf.uint16), 15) == 1 + ) + if tensor.dtype == tf.float32: + return Array._from_tensor( + tf.bitwise.right_shift(tf.bitcast(tensor, tf.uint32), 31) == 1 + ) + if tensor.dtype == tf.float64: + return Array._from_tensor( + tf.bitwise.right_shift(tf.bitcast(tensor, tf.uint64), 63) == 1 + ) + raise TypeError("signbit is only defined for real-valued dtypes") + + +def clip( + x: Array, + /, + min: Array | complex | None = None, + max: Array | complex | None = None, +) -> Array: + tensor = _unwrap(x) + if min is None and max is None: + return Array._from_tensor(tf.identity(tensor)) + if min is None: + max_ = _to_tensor(max, tensor.dtype) + return Array._from_tensor(tf.minimum(tensor, max_)) + if max is None: + min_ = _to_tensor(min, tensor.dtype) + return Array._from_tensor(tf.maximum(tensor, min_)) + min_ = _to_tensor(min, tensor.dtype) + max_ = _to_tensor(max, tensor.dtype) + return Array._from_tensor(tf.minimum(tf.maximum(tensor, min_), max_)) + + +def _as_bool(x: Array) -> tf.Tensor: + tensor = _unwrap(x) + return ( + tensor + if tensor.dtype == tf.bool + else tensor != tf.zeros((), dtype=tensor.dtype) + ) + + +def all( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: py_bool = False, +) -> Array: + return Array._from_tensor(tf.reduce_all(_as_bool(x), axis=axis, keepdims=keepdims)) + + +def any( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: py_bool = False, +) -> Array: + return Array._from_tensor(tf.reduce_any(_as_bool(x), axis=axis, keepdims=keepdims)) + + +def max( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: py_bool = False, +) -> Array: + return Array._from_tensor(tf.reduce_max(_unwrap(x), axis=axis, keepdims=keepdims)) + + +def min( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: py_bool = False, +) -> Array: + return Array._from_tensor(tf.reduce_min(_unwrap(x), axis=axis, keepdims=keepdims)) + + +def mean( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: py_bool = False, +) -> Array: + tensor = _unwrap(x) + if dtype is not None: + tensor = tf.cast(tensor, dtype) + return Array._from_tensor(tf.reduce_mean(tensor, axis=axis, keepdims=keepdims)) + + +def prod( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: py_bool = False, +) -> Array: + tensor = _unwrap(x) + dtype = dtype or _accumulation_dtype(tensor.dtype) + return Array._from_tensor( + tf.reduce_prod(tf.cast(tensor, dtype), axis=axis, keepdims=keepdims) + ) + + +def sum( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + dtype: DType | None = None, + keepdims: py_bool = False, +) -> Array: + tensor = _unwrap(x) + dtype = dtype or _accumulation_dtype(tensor.dtype) + return Array._from_tensor( + tf.reduce_sum(tf.cast(tensor, dtype), axis=axis, keepdims=keepdims) + ) + + +def _axis_size(x: tf.Tensor, axis: int | tuple[int, ...] | None) -> int: + if axis is None: + return math.prod(_shape_tuple(x)) + axes = _normalize_axes(axis, x.shape.rank) + return math.prod(_shape_tuple(x)[a] for a in axes) + + +def var( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: float = 0.0, + dtype: DType | None = None, + keepdims: py_bool = False, +) -> Array: + tensor = _unwrap(x) + dtype = dtype or tensor.dtype + tensor = tf.cast(tensor, dtype) + if axis == (): + return Array._from_tensor(tf.zeros_like(tensor)) + mean_ = tf.reduce_mean(tensor, axis=axis, keepdims=True) + centered = tensor - mean_ + if _is_complex(tensor.dtype): + squared = tf.math.real(centered * tf.math.conj(centered)) + else: + squared = tf.square(centered) + n = _axis_size(tensor, axis) + out = tf.reduce_sum(squared, axis=axis, keepdims=keepdims) / tf.cast( + n - correction, + squared.dtype, + ) + return Array._from_tensor(out) + + +def std( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + correction: float = 0.0, + dtype: DType | None = None, + keepdims: py_bool = False, +) -> Array: + return sqrt( + var(x, axis=axis, correction=correction, dtype=dtype, keepdims=keepdims) + ) + + +def cumulative_sum( + x: Array, + /, + *, + axis: int | None = None, + dtype: DType | None = None, + include_initial: py_bool = False, +) -> Array: + tensor = _unwrap(x) + if axis is None: + if tensor.shape.rank > 1: + raise ValueError( + "axis must be specified in cumulative_sum for more than one dimension" + ) + axis = 0 + axis = _normalize_axis(axis, tensor.shape.rank) + dtype = dtype or _accumulation_dtype(tensor.dtype) + tensor = tf.cast(tensor, dtype) + out = tf.cumsum(tensor, axis=axis) + if include_initial: + shape = list(out.shape.as_list()) + shape[axis] = 1 + out = tf.concat([tf.zeros(shape, dtype=out.dtype), out], axis=axis) + return Array._from_tensor(out) + + +def cumulative_prod( + x: Array, + /, + *, + axis: int | None = None, + dtype: DType | None = None, + include_initial: py_bool = False, +) -> Array: + tensor = _unwrap(x) + if axis is None: + if tensor.shape.rank > 1: + raise ValueError( + "axis must be specified in cumulative_prod for more than one dimension" + ) + axis = 0 + axis = _normalize_axis(axis, tensor.shape.rank) + dtype = dtype or _accumulation_dtype(tensor.dtype) + tensor = tf.cast(tensor, dtype) + out = tf.math.cumprod(tensor, axis=axis) + if include_initial: + shape = list(out.shape.as_list()) + shape[axis] = 1 + out = tf.concat([tf.ones(shape, dtype=out.dtype), out], axis=axis) + return Array._from_tensor(out) + + +def diff( + x: Array, + /, + *, + axis: int = -1, + n: int = 1, + prepend: Array | None = None, + append: Array | None = None, +) -> Array: + tensor = _unwrap(x) + axis = _normalize_axis(axis, tensor.shape.rank) + parts = [] + if prepend is not None: + parts.append(_unwrap(prepend)) + parts.append(tensor) + if append is not None: + parts.append(_unwrap(append)) + tensor = tf.concat(parts, axis=axis) if len(parts) > 1 else tensor + for _ in range(n): + upper = [slice(None)] * tensor.shape.rank + lower = [slice(None)] * tensor.shape.rank + upper[axis] = slice(1, None) + lower[axis] = slice(None, -1) + tensor = tensor[tuple(upper)] - tensor[tuple(lower)] + return Array._from_tensor(tensor) + + +def argsort( + x: Array, + /, + *, + axis: int = -1, + descending: py_bool = False, + stable: py_bool = True, +) -> Array: + del stable + return Array._from_tensor( + tf.argsort( + _unwrap(x), + axis=axis, + direction="DESCENDING" if descending else "ASCENDING", + stable=True, + ) + ) + + +def sort( + x: Array, + /, + *, + axis: int = -1, + descending: py_bool = False, + stable: py_bool = True, +) -> Array: + del stable + return Array._from_tensor( + tf.sort( + _unwrap(x), axis=axis, direction="DESCENDING" if descending else "ASCENDING" + ) + ) + + +def take(x: Array, indices: Array, /, *, axis: int | None = None) -> Array: + tensor = _unwrap(x) + if axis is None: + tensor = tf.reshape(tensor, (-1,)) + axis = 0 + axis = _normalize_axis(axis, tensor.shape.rank) + indices_ = tf.cast(_unwrap(indices), tf.int64) + dim = tensor.shape[axis] + dim = ( + tf.shape(tensor, out_type=tf.int64)[axis] + if dim is None + else tf.cast(dim, tf.int64) + ) + indices_ = tf.where(indices_ < 0, indices_ + dim, indices_) + return Array._from_tensor(tf.gather(tensor, indices_, axis=axis)) + + +def take_along_axis(x: Array, indices: Array, /, *, axis: int = -1) -> Array: + tensor = _unwrap(x) + indices_ = tf.cast(_unwrap(indices), tf.int64) + axis = _normalize_axis(axis, tensor.shape.rank) + dim = tensor.shape[axis] + dim = ( + tf.shape(tensor, out_type=tf.int64)[axis] + if dim is None + else tf.cast(dim, tf.int64) + ) + indices_ = tf.where(indices_ < 0, indices_ + dim, indices_) + out_shape = tf.shape(indices_) + tensor_shape = _shape_tuple(tensor) + coords = [] + for dim_axis, dim_size in enumerate(tensor_shape): + if dim_axis == axis: + coords.append(indices_) + continue + shape = [1] * tensor.shape.rank + shape[dim_axis] = dim_size + coord = tf.reshape(tf.range(dim_size, dtype=tf.int64), shape) + coords.append(tf.broadcast_to(coord, out_shape)) + return Array._from_tensor(tf.gather_nd(tensor, tf.stack(coords, axis=-1))) + + +def matmul(x1: Array, x2: Array, /) -> Array: + x1, x2 = _promote_two(x1, x2) + if ( + x1.shape.rank == 0 + or x2.shape.rank == 0 + or x1.shape.rank == x2.shape.rank == 1 + and _known_unequal(x1.shape[0], x2.shape[0]) + or x1.shape.rank == 1 + and x2.shape.rank >= 2 + and _known_unequal(x1.shape[0], x2.shape[-2]) + or x2.shape.rank == 1 + and x1.shape.rank >= 2 + and _known_unequal(x2.shape[0], x1.shape[-1]) + or x1.shape.rank >= 2 + and x2.shape.rank >= 2 + and _known_unequal(x1.shape[-1], x2.shape[-2]) + ): + raise ValueError("matmul input shapes are incompatible") + out_dtype = x1.dtype + work_dtype = tf.int64 if out_dtype in _integral_dtypes else out_dtype + x1 = tf.cast(x1, work_dtype) + x2 = tf.cast(x2, work_dtype) + x1_was_vector = x1.shape.rank == 1 + x2_was_vector = x2.shape.rank == 1 + if x1_was_vector: + x1 = tf.expand_dims(x1, -2) + if x2_was_vector: + x2 = tf.expand_dims(x2, -1) + batch_shape = tf.broadcast_static_shape(x1.shape[:-2], x2.shape[:-2]) + dynamic_batch_shape = tf.broadcast_dynamic_shape( + tf.shape(x1)[:-2], tf.shape(x2)[:-2] + ) + x1_matrix_shape = x1.shape[-2:] + x2_matrix_shape = x2.shape[-2:] + x1 = tf.broadcast_to( + x1, tf.concat([dynamic_batch_shape, tf.shape(x1)[-2:]], axis=0) + ) + x2 = tf.broadcast_to( + x2, tf.concat([dynamic_batch_shape, tf.shape(x2)[-2:]], axis=0) + ) + x1.set_shape(batch_shape.concatenate(x1_matrix_shape)) + x2.set_shape(batch_shape.concatenate(x2_matrix_shape)) + out = tf.reduce_sum(tf.expand_dims(x1, -1) * tf.expand_dims(x2, -3), axis=-2) + if x1_was_vector: + out = tf.squeeze(out, axis=-2) + if x2_was_vector: + out = tf.squeeze(out, axis=-1) + return Array._from_tensor(tf.cast(out, out_dtype)) + + +def matrix_transpose(x: Array, /) -> Array: + return Array._from_tensor(tf.linalg.matrix_transpose(_unwrap(x))) + + +def tensordot( + x1: Array, + x2: Array, + /, + *, + axes: int | tuple[Sequence[int], Sequence[int]] = 2, +) -> Array: + x1, x2 = _promote_two(x1, x2) + if isinstance(axes, int): + axes1 = tuple(range(x1.shape.rank - axes, x1.shape.rank)) + axes2 = tuple(range(axes)) + else: + axes1, axes2 = tuple(axes[0]), tuple(axes[1]) + axes1 = tuple(_normalize_axis(axis, x1.shape.rank) for axis in axes1) + axes2 = tuple(_normalize_axis(axis, x2.shape.rank) for axis in axes2) + if len(axes1) != len(axes2): + raise ValueError("tensordot axes must have the same length") + for axis1, axis2 in zip(axes1, axes2, strict=True): + if _known_unequal(x1.shape[axis1], x2.shape[axis2]): + raise ValueError("tensordot contraction dimensions must match") + + x1_shape = _shape_tuple(x1) + x2_shape = _shape_tuple(x2) + x1_outer = tuple(axis for axis in range(x1.shape.rank) if axis not in axes1) + x2_outer = tuple(axis for axis in range(x2.shape.rank) if axis not in axes2) + x1_perm = x1_outer + axes1 + x2_perm = axes2 + x2_outer + x1_t = tf.transpose(x1, x1_perm) if x1_perm else x1 + x2_t = tf.transpose(x2, x2_perm) if x2_perm else x2 + x1_outer_shape = tuple(x1_shape[axis] for axis in x1_outer) + x2_outer_shape = tuple(x2_shape[axis] for axis in x2_outer) + contract_shape = tuple(x1_shape[axis] for axis in axes1) + outer1 = _shape_product(x1_outer_shape) + outer2 = _shape_product(x2_outer_shape) + contract = _shape_product(contract_shape) + x1_m = tf.reshape(x1_t, _shape_arg_for_tf((outer1, contract))) + x2_m = tf.reshape(x2_t, _shape_arg_for_tf((contract, outer2))) + out = matmul(Array._from_tensor(x1_m), Array._from_tensor(x2_m)).unwrap() + return Array._from_tensor( + tf.reshape(out, _shape_arg_for_tf(x1_outer_shape + x2_outer_shape)) + ) + + +def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array: + x1, x2 = _promote_two(x1, x2) + shape = tuple(tf.broadcast_static_shape(x1.shape, x2.shape).as_list()) + axis = _normalize_axis(axis, len(shape)) + x1_shape = (1,) * (len(shape) - x1.shape.rank) + _shape_tuple(x1) + x2_shape = (1,) * (len(shape) - x2.shape.rank) + _shape_tuple(x2) + if x1_shape[axis] != x2_shape[axis]: + raise ValueError("vecdot contraction dimensions must match") + x1 = tf.broadcast_to(x1, shape) + x2 = tf.broadcast_to(x2, shape) + if _is_complex(x1.dtype): + x1 = tf.math.conj(x1) + work_dtype = tf.int64 if x1.dtype in _integral_dtypes else x1.dtype + out = tf.reduce_sum(tf.cast(x1, work_dtype) * tf.cast(x2, work_dtype), axis=axis) + return Array._from_tensor(tf.cast(out, x1.dtype)) + + +def broadcast_shapes(*shapes: tuple[int, ...]) -> tuple[int, ...]: + shape = tf.TensorShape(()) + for item in shapes: + shape = tf.broadcast_static_shape(shape, tf.TensorShape(item)) + return tuple(shape.as_list()) + + +def broadcast_to(x: Array, /, shape: tuple[int, ...]) -> Array: + return Array._from_tensor(tf.broadcast_to(_unwrap(x), _shape_arg_for_tf(shape))) + + +def broadcast_arrays(*arrays: Array) -> tuple[Array, ...]: + shape = broadcast_shapes(*(_shape_tuple(x) for x in arrays)) + return tuple(broadcast_to(x, shape) for x in arrays) + + +def concat( + arrays: tuple[Array, ...] | list[Array], + /, + *, + axis: int | None = 0, +) -> Array: + dtype = result_type(*arrays) + tensors = [_to_tensor(x, dtype) for x in arrays] + if axis is None: + tensors = [tf.reshape(x, (-1,)) for x in tensors] + axis = 0 + return Array._from_tensor(tf.concat(tensors, axis=axis)) + + +def expand_dims(x: Array, /, axis: int | tuple[int, ...]) -> Array: + tensor = _unwrap(x) + if isinstance(axis, int): + axis = (axis,) + final_ndim = tensor.shape.rank + len(axis) + axes = tuple(a + final_ndim if a < 0 else a for a in axis) + if len(set(axes)) != len(axes): + raise ValueError("repeated axis") + if py_any(a < 0 or a >= final_ndim for a in axes): + raise IndexError("axis out of bounds") + shape = list(_shape_tuple(tensor)) + for a in sorted(axes): + shape.insert(a, 1) + return Array._from_tensor(tf.reshape(tensor, shape)) + + +def flip(x: Array, /, *, axis: int | tuple[int, ...] | None = None) -> Array: + tensor = _unwrap(x) + axes = _normalize_axes(axis, tensor.shape.rank) + return Array._from_tensor(tf.reverse(tensor, axes)) + + +def meshgrid(*arrays: Array, indexing: Literal["xy", "ij"] = "xy") -> tuple[Array, ...]: + return tuple( + Array._from_tensor(x) + for x in tf.meshgrid(*[_unwrap(a) for a in arrays], indexing=indexing) + ) + + +def moveaxis( + x: Array, + /, + source: int | Sequence[int], + destination: int | Sequence[int], +) -> Array: + return Array._from_tensor(_moveaxis(x, source, destination)) + + +def permute_dims(x: Array, /, axes: tuple[int, ...]) -> Array: + return Array._from_tensor(tf.transpose(_unwrap(x), axes)) + + +def transpose(x: Array, /, axes: tuple[int, ...] | None = None) -> Array: + tensor = _unwrap(x) + if axes is None: + axes = tuple(range(tensor.shape.rank - 1, -1, -1)) + return Array._from_tensor(tf.transpose(tensor, axes)) + + +def einsum(subscripts: str, *operands: Array) -> Array: + return Array._from_tensor(tf.einsum(subscripts, *[_to_tensor(x) for x in operands])) + + +def repeat(x: Array, repeats: int | Array, /, *, axis: int | None = None) -> Array: + tensor = _unwrap(x) + if axis is None: + tensor = tf.reshape(tensor, (-1,)) + axis = 0 + axis = _normalize_axis(axis, tensor.shape.rank) + n = tf.shape(tensor, out_type=tf.int64)[axis] + repeats_ = _unwrap(repeats) + if isinstance(repeats_, tf.Tensor): + repeats_ = tf.cast(repeats_, tf.int64) + indices = tf.repeat(tf.range(n, dtype=tf.int64), repeats_) + return Array._from_tensor(tf.gather(tensor, indices, axis=axis)) + + +def reshape( + x: Array, /, shape: tuple[int, ...], *, copy: py_bool | None = None +) -> Array: + del copy + return Array._from_tensor(tf.reshape(_unwrap(x), _shape_arg_for_tf(shape))) + + +def roll( + x: Array, + /, + shift: int | tuple[int, ...], + axis: int | tuple[int, ...] | None = None, +) -> Array: + tensor = _unwrap(x) + if axis is None: + if tensor.shape.rank == 0: + return Array._from_tensor(tf.identity(tensor)) + shape = _shape_tuple(tensor) + out = tf.roll(tf.reshape(tensor, (-1,)), shift, 0) + return Array._from_tensor(tf.reshape(out, shape)) + return Array._from_tensor(tf.roll(tensor, shift, axis)) + + +def squeeze(x: Array, /, axis: int | tuple[int, ...]) -> Array: + tensor = _unwrap(x) + axes = _normalize_axes(axis, tensor.shape.rank) + if axes == (): + return Array._from_tensor(tf.identity(tensor)) + if py_any(tensor.shape[axis] != 1 for axis in axes): + raise ValueError("cannot squeeze an axis whose size is not 1") + return Array._from_tensor(tf.squeeze(tensor, axis=axes)) + + +def stack(arrays: tuple[Array, ...] | list[Array], /, *, axis: int = 0) -> Array: + dtype = result_type(*arrays) + return Array._from_tensor( + tf.stack([_to_tensor(x, dtype) for x in arrays], axis=axis) + ) + + +def tile(x: Array, repetitions: tuple[int, ...], /) -> Array: + tensor = _unwrap(x) + repetitions = tuple(_unwrap(rep) for rep in repetitions) + if tensor.shape.rank > len(repetitions): + repetitions = (1,) * (tensor.shape.rank - len(repetitions)) + tuple(repetitions) + elif tensor.shape.rank < len(repetitions): + tensor = tf.reshape( + tensor, (1,) * (len(repetitions) - tensor.shape.rank) + _shape_tuple(tensor) + ) + if py_all(isinstance(rep, int) for rep in repetitions): + out = tensor + for axis, rep in enumerate(repetitions): + if rep == 0: + shape = list(_shape_tuple(out)) + shape[axis] = 0 + return Array._from_tensor( + tf.zeros(_shape_arg_for_tf(shape), dtype=out.dtype) + ) + if rep != 1: + out = tf.concat([out] * rep, axis=axis) + return Array._from_tensor(out) + return Array._from_tensor(tf.tile(tensor, _shape_arg_tensor(repetitions))) + + +def unstack(x: Array, /, *, axis: int = 0) -> tuple[Array, ...]: + return tuple(Array._from_tensor(item) for item in tf.unstack(_unwrap(x), axis=axis)) + + +def argmax(x: Array, /, *, axis: int | None = None, keepdims: py_bool = False) -> Array: + tensor = _unwrap(x) + if axis is None: + out = tf.argmax(tf.reshape(tensor, (-1,)), axis=0, output_type=tf.int64) + if keepdims: + out = tf.reshape(out, (1,) * tensor.shape.rank) + return Array._from_tensor(out) + out = tf.argmax(tensor, axis=axis, output_type=tf.int64) + if keepdims: + out = tf.expand_dims(out, axis) + return Array._from_tensor(out) + + +def argmin(x: Array, /, *, axis: int | None = None, keepdims: py_bool = False) -> Array: + tensor = _unwrap(x) + if axis is None: + out = tf.argmin(tf.reshape(tensor, (-1,)), axis=0, output_type=tf.int64) + if keepdims: + out = tf.reshape(out, (1,) * tensor.shape.rank) + return Array._from_tensor(out) + out = tf.argmin(tensor, axis=axis, output_type=tf.int64) + if keepdims: + out = tf.expand_dims(out, axis) + return Array._from_tensor(out) + + +def count_nonzero( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: py_bool = False, +) -> Array: + return Array._from_tensor( + tf.math.count_nonzero(_unwrap(x), axis=axis, keepdims=keepdims, dtype=tf.int64) + ) + + +def nonzero(x: Array, /) -> tuple[Array, ...]: + indices = tf.where(_as_bool(x)) + return tuple( + Array._from_tensor(item) for item in tf.unstack(tf.transpose(indices), axis=0) + ) + + +def searchsorted( + x1: Array, + x2: Array, + /, + *, + side: Literal["left", "right"] = "left", + sorter: Array | None = None, +) -> Array: + if sorter is not None: + x1 = take(x1, sorter) + x1_ = _unwrap(x1) + x2_ = _to_tensor(x2, x1_.dtype) + if x1_.shape.rank == 1: + out = tf.searchsorted(x1_, tf.reshape(x2_, (-1,)), side=side, out_type=tf.int64) + return Array._from_tensor(tf.reshape(out, _shape_tuple(x2_))) + return Array._from_tensor(tf.searchsorted(x1_, x2_, side=side, out_type=tf.int64)) + + +def where(condition: Array, x1: Any, x2: Any, /) -> Array: + x1, x2 = _promote_two(x1, x2) + return Array._from_tensor(tf.where(_unwrap(condition), x1, x2)) + + +def _isnan_tensor(x: tf.Tensor) -> tf.Tensor: + if x.dtype in _real_floating_dtypes: + return tf.math.is_nan(x) + if x.dtype in _complex_floating_dtypes: + return tf.math.is_nan(tf.math.real(x)) | tf.math.is_nan(tf.math.imag(x)) + return tf.zeros(tf.shape(x), dtype=tf.bool) + + +def _unique(x: Array) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]: + flat = tf.reshape(_unwrap(x), (-1,)) + n = tf.shape(flat, out_type=tf.int64)[0] + matrix_shape = tf.stack([n, n]) + idx = tf.range(n, dtype=tf.int64) + equality = tf.equal(tf.expand_dims(flat, 1), tf.expand_dims(flat, 0)) + nan = _isnan_tensor(flat) + same_position = tf.equal(tf.expand_dims(idx, 1), tf.expand_dims(idx, 0)) + equality = equality | ( + tf.expand_dims(nan, 1) & tf.expand_dims(nan, 0) & same_position + ) + + first_indices = tf.reduce_min( + tf.where( + equality, + tf.broadcast_to(idx, matrix_shape), + tf.fill(matrix_shape, tf.cast(n, tf.int64)), + ), + axis=1, + ) + unique_mask = tf.equal(first_indices, idx) + unique_mask.set_shape((None,)) + indices = tf.boolean_mask(idx, unique_mask) + values = tf.gather(flat, indices) + inverse_equality = tf.equal( + tf.expand_dims(first_indices, 1), tf.expand_dims(indices, 0) + ) + inverse = tf.argmax( + tf.cast(inverse_equality, tf.int64), axis=1, output_type=tf.int64 + ) + counts = tf.reduce_sum(tf.cast(inverse_equality, tf.int64), axis=0) + return values, indices, inverse, counts + + +def unique_all(x: Array) -> UniqueAllResult: + values, indices, inverse, counts = _unique(x) + return UniqueAllResult( + Array._from_tensor(values), + Array._from_tensor(indices), + Array._from_tensor(tf.reshape(inverse, tf.shape(_unwrap(x)))), + Array._from_tensor(counts), + ) + + +def unique_counts(x: Array) -> UniqueCountsResult: + values, _, _, counts = _unique(x) + return UniqueCountsResult(Array._from_tensor(values), Array._from_tensor(counts)) + + +def unique_inverse(x: Array) -> UniqueInverseResult: + values, _, inverse, _ = _unique(x) + return UniqueInverseResult( + Array._from_tensor(values), + Array._from_tensor(tf.reshape(inverse, tf.shape(_unwrap(x)))), + ) + + +def unique_values(x: Array) -> Array: + values, _, _, _ = _unique(x) + return Array._from_tensor(values) + + +def isin(x1: Array | int, x2: Array | int, /, *, invert: py_bool = False) -> Array: + dtype = _result_type_with_scalars(x1, x2) + x1_ = _to_tensor(x1, dtype) + x2_ = tf.reshape(_to_tensor(x2, dtype), (-1,)) + out = tf.reduce_any(tf.equal(tf.expand_dims(x1_, -1), x2_), axis=-1) + return Array._from_tensor(tf.logical_not(out) if invert else out) + + +__all__ = [ + "Array", + "DType", + "Device", + "UniqueAllResult", + "UniqueCountsResult", + "UniqueInverseResult", + "abs", + "acos", + "acosh", + "add", + "all", + "any", + "arange", + "argmax", + "argmin", + "argsort", + "asarray", + "asin", + "asinh", + "astype", + "atan", + "atan2", + "atanh", + "bfloat16", + "bitwise_and", + "bitwise_invert", + "bitwise_left_shift", + "bitwise_or", + "bitwise_right_shift", + "bitwise_xor", + "bool", + "broadcast_arrays", + "broadcast_shapes", + "broadcast_to", + "can_cast", + "ceil", + "clip", + "complex64", + "complex128", + "concat", + "conj", + "copysign", + "cos", + "cosh", + "count_nonzero", + "cumulative_prod", + "cumulative_sum", + "diff", + "divide", + "e", + "einsum", + "empty", + "empty_like", + "equal", + "exp", + "expand_dims", + "expm1", + "eye", + "finfo", + "flip", + "float16", + "float32", + "float64", + "floor", + "floor_divide", + "from_dlpack", + "full", + "full_like", + "greater", + "greater_equal", + "hypot", + "iinfo", + "imag", + "inf", + "int8", + "int16", + "int32", + "int64", + "isdtype", + "isfinite", + "isinf", + "isin", + "isnan", + "less", + "less_equal", + "linspace", + "log", + "log1p", + "log2", + "log10", + "logaddexp", + "logical_and", + "logical_not", + "logical_or", + "logical_xor", + "matmul", + "matrix_transpose", + "max", + "maximum", + "mean", + "meshgrid", + "min", + "minimum", + "moveaxis", + "multiply", + "nan", + "negative", + "nextafter", + "newaxis", + "nonzero", + "not_equal", + "ones", + "ones_like", + "permute_dims", + "pi", + "positive", + "pow", + "prod", + "real", + "reciprocal", + "remainder", + "repeat", + "reshape", + "result_type", + "roll", + "round", + "searchsorted", + "sign", + "signbit", + "sin", + "sinh", + "sort", + "sqrt", + "square", + "squeeze", + "stack", + "std", + "subtract", + "sum", + "take", + "take_along_axis", + "tan", + "tanh", + "tensordot", + "tile", + "tril", + "triu", + "trunc", + "transpose", + "uint8", + "uint16", + "uint32", + "uint64", + "unique_all", + "unique_counts", + "unique_inverse", + "unique_values", + "unstack", + "var", + "vecdot", + "where", + "zeros", + "zeros_like", +] + + +def __dir__() -> list[str]: + return __all__ diff --git a/deepmd/_vendors/ndtensorflow/fft.py b/deepmd/_vendors/ndtensorflow/fft.py new file mode 100644 index 0000000000..4c3a44179b --- /dev/null +++ b/deepmd/_vendors/ndtensorflow/fft.py @@ -0,0 +1,379 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +import math +from collections.abc import ( + Sequence, +) +from contextlib import ( + nullcontext, +) +from typing import ( + Literal, +) + +import tensorflow as tf + +from ._array import ( + Array, +) +from ._namespace import ( + Device, + DType, + _moveaxis, + _unwrap, +) + +_Norm = Literal["backward", "ortho", "forward"] + + +def _wrap(x: tf.Tensor) -> Array: + return Array._from_tensor(x) + + +def _complex_dtype(dtype: DType) -> DType: + if dtype in (tf.float64, tf.complex128): + return tf.complex128 + return tf.complex64 + + +def _real_dtype(dtype: DType) -> DType: + if dtype == tf.complex128: + return tf.float64 + return tf.float32 + + +def _as_complex(x: tf.Tensor) -> tf.Tensor: + if x.dtype.is_complex: + return x + return tf.cast(x, _complex_dtype(x.dtype)) + + +def _shape_tuple(x: tf.Tensor) -> tuple[int, ...]: + return tuple(x.shape.as_list()) + + +def _normalize_axis(axis: int, ndim: int) -> int: + if axis < 0: + axis += ndim + if axis < 0 or axis >= ndim: + raise IndexError(f"axis {axis} is out of bounds for array of dimension {ndim}") + return axis + + +def _normalize_axes( + axes: Sequence[int] | None, + ndim: int, + s: Sequence[int] | None, +) -> tuple[int, ...]: + if axes is None: + axes = tuple(range(ndim)) if s is None else tuple(range(ndim - len(s), ndim)) + axes = tuple(_normalize_axis(a, ndim) for a in axes) + if len(set(axes)) != len(axes): + raise ValueError("repeated axis") + return axes + + +def _resize_axis(x: tf.Tensor, n: int | None, axis: int) -> tf.Tensor: + if n is None or x.shape[axis] == n: + return x + shape = list(_shape_tuple(x)) + current = shape[axis] + if current > n: + begin = [0] * x.shape.rank + size = shape + size[axis] = n + return tf.slice(x, begin, size) + paddings = [[0, 0] for _ in range(x.shape.rank)] + paddings[axis][1] = n - current + return tf.pad(x, paddings) + + +def _apply_1d(x: tf.Tensor, func, n: int | None, axis: int) -> tf.Tensor: + x = _moveaxis(x, axis, -1) + x = _resize_axis(x, n, -1) + x = func(x) + return _moveaxis(x, -1, axis) + + +def _norm_size(x: tf.Tensor, axes: tuple[int, ...], s: Sequence[int] | None) -> int: + if s is None: + return math.prod(x.shape[a] for a in axes) + return math.prod(s) + + +def _scale_forward(x: tf.Tensor, n: int, norm: _Norm) -> tf.Tensor: + if norm == "backward": + return x + scale = tf.cast(n if norm == "forward" else math.sqrt(n), x.dtype) + return x / scale + + +def _scale_inverse(x: tf.Tensor, n: int, norm: _Norm) -> tf.Tensor: + if norm == "backward": + return x + scale = tf.cast(n if norm == "forward" else math.sqrt(n), x.dtype) + return x * scale + + +def fft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + tensor = _as_complex(_unwrap(x)) + axis = _normalize_axis(axis, tensor.shape.rank) + out = _apply_1d(tensor, tf.signal.fft, n, axis) + return _wrap(_scale_forward(out, out.shape[axis], norm)) + + +def ifft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + tensor = _as_complex(_unwrap(x)) + axis = _normalize_axis(axis, tensor.shape.rank) + out = _apply_1d(tensor, tf.signal.ifft, n, axis) + return _wrap(_scale_inverse(out, out.shape[axis], norm)) + + +def fftn( + x: Array, + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> Array: + tensor = _as_complex(_unwrap(x)) + axes_ = _normalize_axes(axes, tensor.shape.rank, s) + sizes = [None] * len(axes_) if s is None else list(s) + n = _norm_size(tensor, axes_, s) + out = tensor + for axis, size in zip(axes_, sizes, strict=True): + out = _apply_1d(out, tf.signal.fft, size, axis) + return _wrap(_scale_forward(out, n, norm)) + + +def ifftn( + x: Array, + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> Array: + tensor = _as_complex(_unwrap(x)) + axes_ = _normalize_axes(axes, tensor.shape.rank, s) + sizes = [None] * len(axes_) if s is None else list(s) + n = _norm_size(tensor, axes_, s) + out = tensor + for axis, size in zip(axes_, sizes, strict=True): + out = _apply_1d(out, tf.signal.ifft, size, axis) + return _wrap(_scale_inverse(out, n, norm)) + + +def rfft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + tensor = _unwrap(x) + axis = _normalize_axis(axis, tensor.shape.rank) + out = _apply_1d(tensor, tf.signal.rfft, n, axis) + size = n if n is not None else tensor.shape[axis] + return _wrap(_scale_forward(out, size, norm)) + + +def irfft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + tensor = _unwrap(x) + axis = _normalize_axis(axis, tensor.shape.rank) + out = _apply_1d( + tensor, + lambda y: tf.signal.irfft(y, fft_length=[n] if n is not None else None), + None, + axis, + ) + size = n if n is not None else out.shape[axis] + return _wrap(_scale_inverse(out, size, norm)) + + +def rfftn( + x: Array, + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> Array: + tensor = _unwrap(x) + axes_ = _normalize_axes(axes, tensor.shape.rank, s) + sizes = [None] * len(axes_) if s is None else list(s) + n = _norm_size(tensor, axes_, s) + out = _apply_1d(tensor, tf.signal.rfft, sizes[-1], axes_[-1]) + for axis, size in zip(axes_[:-1], sizes[:-1], strict=True): + out = _apply_1d(_as_complex(out), tf.signal.fft, size, axis) + return _wrap(_scale_forward(out, n, norm)) + + +def irfftn( + x: Array, + /, + *, + s: Sequence[int] | None = None, + axes: Sequence[int] | None = None, + norm: _Norm = "backward", +) -> Array: + tensor = _unwrap(x) + axes_ = _normalize_axes(axes, tensor.shape.rank, s) + sizes = [None] * len(axes_) if s is None else list(s) + if s is None: + last_size = 2 * (tensor.shape[axes_[-1]] - 1) + n = math.prod([*(tensor.shape[a] for a in axes_[:-1]), last_size]) + else: + n = math.prod(s) + out = tensor + for axis, size in zip(axes_[:-1], sizes[:-1], strict=True): + out = _apply_1d(out, tf.signal.ifft, size, axis) + out = _apply_1d( + out, + lambda y: tf.signal.irfft( + y, fft_length=[sizes[-1]] if sizes[-1] is not None else None + ), + None, + axes_[-1], + ) + return _wrap(_scale_inverse(out, n, norm)) + + +def hfft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + tensor = _unwrap(x) + size = n if n is not None else 2 * (tensor.shape[axis] - 1) + return _wrap( + _unwrap(irfft(_wrap(tf.math.conj(tensor)), n=size, axis=axis, norm=norm)) + * tf.cast(size, _real_dtype(tensor.dtype)) + ) + + +def ihfft( + x: Array, + /, + *, + n: int | None = None, + axis: int = -1, + norm: _Norm = "backward", +) -> Array: + tensor = _unwrap(x) + size = n if n is not None else tensor.shape[axis] + return _wrap( + tf.math.conj(_unwrap(rfft(x, n=size, axis=axis, norm=norm))) + / tf.cast(size, _complex_dtype(tensor.dtype)) + ) + + +def fftfreq( + n: int, + /, + *, + d: float = 1.0, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + with tf.device(device) if device is not None else nullcontext(): + dtype = dtype or tf.float32 + positive = tf.range(0, (n - 1) // 2 + 1, dtype=dtype) + negative = tf.range(-(n // 2), 0, dtype=dtype) + return _wrap(tf.concat([positive, negative], axis=0) / tf.cast(n * d, dtype)) + + +def rfftfreq( + n: int, + /, + *, + d: float = 1.0, + dtype: DType | None = None, + device: Device | None = None, +) -> Array: + with tf.device(device) if device is not None else nullcontext(): + dtype = dtype or tf.float32 + return _wrap(tf.range(0, n // 2 + 1, dtype=dtype) / tf.cast(n * d, dtype)) + + +def fftshift( + x: Array, + /, + *, + axes: int | Sequence[int] | None = None, +) -> Array: + tensor = _unwrap(x) + axes_ = _normalize_axes( + None if axes is None else (axes if isinstance(axes, Sequence) else (axes,)), + tensor.shape.rank, + None, + ) + shifts = tuple(tensor.shape[axis] // 2 for axis in axes_) + return _wrap(tf.roll(tensor, shifts, axes_)) + + +def ifftshift( + x: Array, + /, + *, + axes: int | Sequence[int] | None = None, +) -> Array: + tensor = _unwrap(x) + axes_ = _normalize_axes( + None if axes is None else (axes if isinstance(axes, Sequence) else (axes,)), + tensor.shape.rank, + None, + ) + shifts = tuple(-(tensor.shape[axis] // 2) for axis in axes_) + return _wrap(tf.roll(tensor, shifts, axes_)) + + +__all__ = [ + "fft", + "ifft", + "fftn", + "ifftn", + "rfft", + "irfft", + "rfftn", + "irfftn", + "hfft", + "ihfft", + "fftfreq", + "rfftfreq", + "fftshift", + "ifftshift", +] + + +def __dir__() -> list[str]: + return __all__ diff --git a/deepmd/_vendors/ndtensorflow/linalg.py b/deepmd/_vendors/ndtensorflow/linalg.py new file mode 100644 index 0000000000..9b5c4222aa --- /dev/null +++ b/deepmd/_vendors/ndtensorflow/linalg.py @@ -0,0 +1,365 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from __future__ import ( + annotations, +) + +import math +from collections import ( + namedtuple, +) + +import tensorflow as tf + +from ._array import ( + Array, +) +from ._namespace import ( + _is_complex, + _moveaxis, + _promote_two, + _real_dtype_for, + _shape_tuple, + _unwrap, + count_nonzero, + finfo, + matmul, + matrix_transpose, +) +from ._namespace import sum as xp_sum +from ._namespace import ( + tensordot, + vecdot, +) + +EighResult = namedtuple("EighResult", ["eigenvalues", "eigenvectors"]) +EigResult = namedtuple("EigResult", ["eigenvalues", "eigenvectors"]) +QRResult = namedtuple("QRResult", ["Q", "R"]) +SlogdetResult = namedtuple("SlogdetResult", ["sign", "logabsdet"]) +SVDResult = namedtuple("SVDResult", ["U", "S", "Vh"]) + + +def _wrap(x: tf.Tensor) -> Array: + return Array._from_tensor(x) + + +def _replace_nonfinite(x: tf.Tensor) -> tf.Tensor: + if _is_complex(x.dtype): + finite = tf.math.is_finite(tf.math.real(x)) & tf.math.is_finite(tf.math.imag(x)) + else: + finite = tf.math.is_finite(x) + return tf.where(finite, x, tf.zeros((), dtype=x.dtype)) + + +def outer(x1: Array, x2: Array, /) -> Array: + x1_, x2_ = _promote_two(x1, x2) + return _wrap(tf.reshape(x1_, (-1, 1)) * tf.reshape(x2_, (1, -1))) + + +def cross(x1: Array, x2: Array, /, *, axis: int = -1) -> Array: + x1_, x2_ = _promote_two(x1, x2) + if (x1_.shape[axis] is not None and x1_.shape[axis] != 3) or ( + x2_.shape[axis] is not None and x2_.shape[axis] != 3 + ): + raise ValueError("cross product axis must have size 3") + x1_ = _moveaxis(x1_, axis, -1) + x2_ = _moveaxis(x2_, axis, -1) + shape = tf.broadcast_static_shape(x1_.shape, x2_.shape) + shape = ( + tf.broadcast_dynamic_shape(tf.shape(x1_), tf.shape(x2_)) + if not shape.is_fully_defined() + else shape + ) + x1_, x2_ = tf.broadcast_to(x1_, shape), tf.broadcast_to(x2_, shape) + return _wrap(_moveaxis(tf.linalg.cross(x1_, x2_), -1, axis)) + + +def eigh(x: Array, /) -> EighResult: + values, vectors = tf.linalg.eigh(_unwrap(x)) + return EighResult(_wrap(values), _wrap(vectors)) + + +def eig(x: Array, /) -> EigResult: + values, vectors = tf.linalg.eig(_unwrap(x)) + return EigResult(_wrap(values), _wrap(vectors)) + + +def eigvals(x: Array, /) -> Array: + values, _ = tf.linalg.eig(_unwrap(x)) + return _wrap(values) + + +def eigvalsh(x: Array, /) -> Array: + return _wrap(tf.linalg.eigvalsh(_unwrap(x))) + + +def det(x: Array, /) -> Array: + return _wrap(tf.linalg.det(_unwrap(x))) + + +def inv(x: Array, /) -> Array: + return _wrap(tf.linalg.inv(_unwrap(x))) + + +def qr(x: Array, /, *, mode: str = "reduced") -> QRResult: + if mode not in ("reduced", "complete"): + raise ValueError("mode must be 'reduced' or 'complete'") + res = tf.linalg.qr(_replace_nonfinite(_unwrap(x)), full_matrices=mode == "complete") + return QRResult(_wrap(_replace_nonfinite(res.q)), _wrap(_replace_nonfinite(res.r))) + + +def slogdet(x: Array, /) -> SlogdetResult: + res = tf.linalg.slogdet(_unwrap(x)) + return SlogdetResult(_wrap(res.sign), _wrap(tf.math.real(res.log_abs_determinant))) + + +def svd(x: Array, /, *, full_matrices: bool = True) -> SVDResult: + s, u, v = tf.linalg.svd(_unwrap(x), full_matrices=full_matrices, compute_uv=True) + vh = tf.linalg.matrix_transpose(tf.math.conj(v)) + return SVDResult(_wrap(u), _wrap(s), _wrap(vh)) + + +def cholesky(x: Array, /, *, upper: bool = False) -> Array: + out = tf.linalg.cholesky(_unwrap(x)) + if upper: + out = tf.linalg.matrix_transpose(out) + if _is_complex(out.dtype): + out = tf.math.conj(out) + return _wrap(out) + + +def matrix_rank(x: Array, /, *, rtol: float | Array | None = None) -> Array: + tensor = _unwrap(x) + if tensor.shape.rank < 2: + raise ValueError( + "1-dimensional array given. Array must be at least two-dimensional" + ) + s = _unwrap(svdvals(x)) + if rtol is None: + tol = ( + tf.reduce_max(s, axis=-1, keepdims=True) + * max(tensor.shape[-2:]) + * finfo(s.dtype).eps + ) + else: + tol = ( + tf.reduce_max(s, axis=-1, keepdims=True) + * tf.cast(_unwrap(rtol), s.dtype)[..., tf.newaxis] + ) + return count_nonzero(Array._from_tensor(s > tol), axis=-1) + + +def pinv(x: Array, /, *, rtol: float | Array | None = None) -> Array: + tensor = _unwrap(x) + s, u, v = tf.linalg.svd(tensor, full_matrices=False, compute_uv=True) + if rtol is None: + rtol = max(tensor.shape[-2:]) * finfo(tensor.dtype).eps + rtol_ = tf.cast(_unwrap(rtol), s.dtype) + if rtol_.shape.rank != 0: + rtol_ = rtol_[..., tf.newaxis] + cutoff = tf.reduce_max(s, axis=-1, keepdims=True) * rtol_ + s_inv = tf.where(s > cutoff, tf.math.reciprocal(s), tf.zeros((), dtype=s.dtype)) + v_scaled = v * tf.cast(s_inv[..., tf.newaxis, :], v.dtype) + return _wrap( + tf.linalg.matmul(v_scaled, tf.linalg.matrix_transpose(tf.math.conj(u))) + ) + + +def matrix_norm( + x: Array, + /, + *, + keepdims: bool = False, + ord: int | float | str | None = "fro", +) -> Array: + tensor = _unwrap(x) + out_dtype = _real_dtype_for(tensor.dtype) + abs_x = tf.cast(tf.abs(tensor), out_dtype) + + if ord in (None, "fro"): + out = tf.sqrt(tf.reduce_sum(tf.square(abs_x), axis=(-2, -1))) + elif ord == 1: + out = tf.reduce_max(tf.reduce_sum(abs_x, axis=-2), axis=-1) + elif ord == -1: + out = tf.reduce_min(tf.reduce_sum(abs_x, axis=-2), axis=-1) + elif ord == float("inf"): + out = tf.reduce_max(tf.reduce_sum(abs_x, axis=-1), axis=-1) + elif ord == -float("inf"): + out = tf.reduce_min(tf.reduce_sum(abs_x, axis=-1), axis=-1) + elif ord in (2, -2, "nuc"): + s = tf.linalg.svd(tensor, compute_uv=False) + if ord == 2: + out = tf.reduce_max(s, axis=-1) + elif ord == -2: + out = tf.reduce_min(s, axis=-1) + else: + out = tf.reduce_sum(s, axis=-1) + else: + raise ValueError(f"unsupported matrix norm order: {ord!r}") + + out = tf.cast(out, out_dtype) + if keepdims: + out = tf.reshape(out, _shape_tuple(tensor)[:-2] + (1, 1)) + return _wrap(out) + + +def matrix_power(x: Array, n: int, /) -> Array: + tensor = _unwrap(x) + if n == 0: + eye = tf.eye(tensor.shape[-1], dtype=tensor.dtype) + return _wrap(tf.broadcast_to(eye, _shape_tuple(tensor))) + if n < 0: + tensor = tf.linalg.inv(tensor) + n = -n + result = tensor + for _ in range(n - 1): + result = tf.linalg.matmul(result, tensor) + return _wrap(result) + + +def solve(x1: Array, x2: Array, /) -> Array: + x1_, x2_ = _promote_two(x1, x2) + squeeze = False + if x2_.shape.rank == 1: + stack_shape = _shape_tuple(x1_)[:-2] + x2_ = tf.reshape(x2_, (1,) * len(stack_shape) + _shape_tuple(x2_) + (1,)) + x2_ = tf.broadcast_to(x2_, stack_shape + (_shape_tuple(x1_)[-1], 1)) + squeeze = True + else: + stack_shape = tuple( + tf.broadcast_static_shape( + tf.TensorShape(_shape_tuple(x1_)[:-2]), + tf.TensorShape(_shape_tuple(x2_)[:-2]), + ).as_list() + ) + x1_ = tf.broadcast_to(x1_, stack_shape + _shape_tuple(x1_)[-2:]) + x2_ = tf.broadcast_to(x2_, stack_shape + _shape_tuple(x2_)[-2:]) + out = tf.linalg.solve(x1_, x2_) + return _wrap(tf.squeeze(out, axis=-1) if squeeze else out) + + +def svdvals(x: Array, /) -> Array: + return _wrap(tf.linalg.svd(_unwrap(x), compute_uv=False)) + + +def diagonal(x: Array, /, *, offset: int = 0) -> Array: + tensor = _unwrap(x) + if tensor.shape.rank < 2: + raise ValueError("x must be at least 2-dimensional for diagonal") + if offset >= tensor.shape[-1] or offset <= -tensor.shape[-2]: + return _wrap(tf.zeros(_shape_tuple(tensor)[:-2] + (0,), dtype=tensor.dtype)) + return _wrap(tf.linalg.diag_part(tensor, k=offset)) + + +def trace(x: Array, /, *, offset: int = 0, dtype: tf.DType | None = None) -> Array: + return xp_sum(diagonal(x, offset=offset), axis=-1, dtype=dtype) + + +def vector_norm( + x: Array, + /, + *, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, + ord: int | float = 2, +) -> Array: + tensor = _unwrap(x) + out_dtype = _real_dtype_for(tensor.dtype) + if axis == (): + return _wrap( + tf.cast(tensor != 0, out_dtype) + if ord == 0 + else tf.cast(tf.abs(tensor), out_dtype) + ) + + if axis is None: + x_ = tf.reshape(tensor, (-1,)) + axis_ = 0 + elif isinstance(axis, tuple): + axes = tuple(a + tensor.shape.rank if a < 0 else a for a in axis) + rest = tuple(i for i in range(tensor.shape.rank) if i not in axes) + x_ = tf.transpose(tensor, axes + rest) + axis_size = math.prod(tensor.shape[a] for a in axes) + x_ = tf.reshape(x_, (axis_size, *[tensor.shape[i] for i in rest])) + axis_ = 0 + else: + x_ = tensor + axis_ = axis + + abs_x = tf.cast(tf.abs(x_), out_dtype) + if ord == 0: + out = tf.cast( + count_nonzero(Array._from_tensor(x_), axis=axis_).unwrap(), out_dtype + ) + elif ord == 1: + out = tf.reduce_sum(abs_x, axis=axis_) + elif ord == 2: + out = tf.sqrt(tf.reduce_sum(tf.square(abs_x), axis=axis_)) + elif ord == float("inf"): + out = tf.reduce_max(abs_x, axis=axis_) + elif ord == -float("inf"): + out = tf.reduce_min(abs_x, axis=axis_) + else: + p = tf.cast(ord, out_dtype) + out = tf.pow(tf.reduce_sum(tf.pow(abs_x, p), axis=axis_), 1 / p) + + if keepdims: + shape = list(_shape_tuple(tensor)) + axes = ( + range(tensor.shape.rank) + if axis is None + else (axis if isinstance(axis, tuple) else (axis,)) + ) + for a in axes: + shape[a] = 1 + out = tf.reshape(out, shape) + return _wrap(out) + + +def norm( + x: Array, + ord: int | float | str | None = None, + axis: int | tuple[int, ...] | None = None, + keepdims: bool = False, +) -> Array: + if isinstance(axis, tuple) and len(axis) == 2: + return matrix_norm(x, ord="fro" if ord is None else ord, keepdims=keepdims) + return vector_norm(x, ord=2 if ord is None else ord, axis=axis, keepdims=keepdims) + + +__all__ = [ + "EigResult", + "EighResult", + "QRResult", + "SVDResult", + "SlogdetResult", + "cholesky", + "cross", + "diagonal", + "det", + "eig", + "eigh", + "eigvals", + "eigvalsh", + "inv", + "matmul", + "matrix_norm", + "norm", + "matrix_power", + "matrix_rank", + "matrix_transpose", + "outer", + "pinv", + "qr", + "slogdet", + "solve", + "svd", + "svdvals", + "tensordot", + "trace", + "vecdot", + "vector_norm", +] + + +def __dir__() -> list[str]: + return __all__ diff --git a/deepmd/_vendors/ndtensorflow/py.typed b/deepmd/_vendors/ndtensorflow/py.typed new file mode 100644 index 0000000000..e69de29bb2 diff --git a/deepmd/backend/tf2.py b/deepmd/backend/tf2.py new file mode 100644 index 0000000000..e0c750cb70 --- /dev/null +++ b/deepmd/backend/tf2.py @@ -0,0 +1,76 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) +from importlib.util import ( + find_spec, +) +from typing import ( + TYPE_CHECKING, + ClassVar, +) + +from deepmd.backend.backend import ( + Backend, +) + +if TYPE_CHECKING: + from argparse import ( + Namespace, + ) + + from deepmd.infer.deep_eval import ( + DeepEvalBackend, + ) + from deepmd.utils.neighbor_stat import ( + NeighborStat, + ) + + +@Backend.register("tf2") +@Backend.register("tensorflow2") +class TensorFlow2Backend(Backend): + """TensorFlow 2 eager backend.""" + + name = "TensorFlow2" + features: ClassVar[Backend.Feature] = Backend.Feature.DEEP_EVAL | Backend.Feature.IO + suffixes: ClassVar[list[str]] = [".savedmodeltf"] + + @classmethod + def match_filename(cls, filename: str) -> int: + return 2 if str(filename).lower().endswith(".savedmodeltf") else 0 + + def is_available(self) -> bool: + return find_spec("tensorflow") is not None + + @property + def entry_point_hook(self) -> Callable[["Namespace"], None]: + raise NotImplementedError("Training entry point is not implemented for TF2") + + @property + def deep_eval(self) -> type["DeepEvalBackend"]: + from deepmd.tf2.infer.deep_eval import ( + DeepEval, + ) + + return DeepEval + + @property + def neighbor_stat(self) -> type["NeighborStat"]: + raise NotImplementedError("Neighbor statistics are not implemented for TF2") + + @property + def serialize_hook(self) -> Callable[[str], dict]: + from deepmd.tf2.utils.serialization import ( + serialize_from_file, + ) + + return serialize_from_file + + @property + def deserialize_hook(self) -> Callable[[str, dict], None]: + from deepmd.tf2.utils.serialization import ( + deserialize_to_file, + ) + + return deserialize_to_file diff --git a/deepmd/dpmodel/array_api.py b/deepmd/dpmodel/array_api.py index c5547af79a..454ffc41a4 100644 --- a/deepmd/dpmodel/array_api.py +++ b/deepmd/dpmodel/array_api.py @@ -133,6 +133,37 @@ def xp_scatter_sum(input: Array, dim: int, index: Array, src: Array) -> Array: # Generic array_api implementation (works for JAX, NumPy, array-api-strict, etc.) xp = array_api_compat.array_namespace(input) + if getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow": + import tensorflow as tf + + input_tensor = input.unwrap() + index_tensor = tf.cast(index.unwrap(), tf.int64) + src_tensor = src.unwrap() + rank = input_tensor.shape.rank + if rank is None: + raise ValueError("xp_scatter_sum requires a statically known rank") + dim = dim + rank if dim < 0 else dim + src_shape = tf.shape(src_tensor, out_type=tf.int64) + coords = [] + for axis in range(rank): + if axis == dim: + coord = index_tensor + else: + view_shape = [1] * rank + view_shape[axis] = src_shape[axis] + coord = tf.broadcast_to( + tf.reshape(tf.range(src_shape[axis], dtype=tf.int64), view_shape), + src_shape, + ) + coords.append(coord) + scatter_indices = tf.reshape(tf.stack(coords, axis=-1), (-1, rank)) + scatter_updates = tf.reshape(src_tensor, (-1,)) + scattered = tf.scatter_nd( + scatter_indices, + scatter_updates, + tf.shape(input_tensor, out_type=tf.int64), + ) + return xp.asarray(input_tensor + scattered) # Create flat index array matching input shape idx = xp.arange(input.size, dtype=xp.int64, device=array_api_compat.device(input)) diff --git a/deepmd/dpmodel/atomic_model/polar_atomic_model.py b/deepmd/dpmodel/atomic_model/polar_atomic_model.py index 76a221de46..f1234102e7 100644 --- a/deepmd/dpmodel/atomic_model/polar_atomic_model.py +++ b/deepmd/dpmodel/atomic_model/polar_atomic_model.py @@ -52,7 +52,10 @@ def apply_out_stat( for kk in self.bias_keys: ntypes = out_bias[kk].shape[0] temp = xp.mean( - xp.diagonal(out_bias[kk].reshape(ntypes, 3, 3), 0, 1, 2), + xp.linalg.diagonal( + out_bias[kk].reshape(ntypes, 3, 3), + offset=0, + ), axis=1, ) modified_bias = temp[atype] diff --git a/deepmd/dpmodel/descriptor/dpa2.py b/deepmd/dpmodel/descriptor/dpa2.py index 928c58a9b4..08f68849d2 100644 --- a/deepmd/dpmodel/descriptor/dpa2.py +++ b/deepmd/dpmodel/descriptor/dpa2.py @@ -936,7 +936,12 @@ def call( mapping_ext = xp.tile( xp.expand_dims(mapping, axis=-1), (1, 1, g1.shape[-1]) ) + mapping_mask = mapping_ext >= 0 + mapping_ext = xp.where( + mapping_mask, mapping_ext, xp.zeros_like(mapping_ext) + ) g1_ext = xp_take_along_axis(g1, mapping_ext, axis=1) + g1_ext = xp.where(mapping_mask, g1_ext, xp.zeros_like(g1_ext)) else: # parallel mode: hand the local-only g1 to the repformer block; # its per-layer override fills ghosts via the MPI exchange. diff --git a/deepmd/dpmodel/descriptor/repflows.py b/deepmd/dpmodel/descriptor/repflows.py index c7a7f4a8d6..5d1ee9fd24 100644 --- a/deepmd/dpmodel/descriptor/repflows.py +++ b/deepmd/dpmodel/descriptor/repflows.py @@ -55,6 +55,7 @@ _cal_grrg, _cal_hg, _make_nei_g1, + _statically_compatible_shape, get_residual, symmetrization_op, ) @@ -554,7 +555,13 @@ def _exchange_ghosts( "`mapping` is required when use_loc_mapping=False unless " "`_exchange_ghosts` is overridden for parallel comm handling." ) - return xp_take_along_axis(node_ebd, mapping_tiled, axis=1) + xp = array_api_compat.array_namespace(node_ebd, mapping_tiled) + mapping_mask = mapping_tiled >= 0 + mapping_tiled = xp.where( + mapping_mask, mapping_tiled, xp.zeros_like(mapping_tiled) + ) + node_ebd_ext = xp_take_along_axis(node_ebd, mapping_tiled, axis=1) + return xp.where(mapping_mask, node_ebd_ext, xp.zeros_like(node_ebd_ext)) def call( self, @@ -615,7 +622,9 @@ def call( # get node embedding # nb x nloc x tebd_dim atype_embd = xp_take_first_n(atype_embd_ext, 1, nloc) - assert list(atype_embd.shape) == [nframes, nloc, self.n_dim] + assert _statically_compatible_shape( + atype_embd.shape, (nframes, nloc, self.n_dim) + ) node_ebd = self.act(atype_embd) @@ -1622,11 +1631,11 @@ def call( else 0 ) node_ebd = xp_take_first_n(node_ebd_ext, 1, nloc) - assert (nb, nloc) == node_ebd.shape[:2] + assert _statically_compatible_shape(node_ebd.shape[:2], (nb, nloc)) if not self.use_dynamic_sel: - assert (nb, nloc, nnei) == h2.shape[:3] + assert _statically_compatible_shape(h2.shape[:3], (nb, nloc, nnei)) else: - assert (n_edge, 3) == h2.shape + assert _statically_compatible_shape(h2.shape, (n_edge, 3)) del a_nlist # may be used in the future n2e_index, n_ext2e_index = edge_index[0, :], edge_index[1, :] diff --git a/deepmd/dpmodel/descriptor/repformers.py b/deepmd/dpmodel/descriptor/repformers.py index 799ab0c3c3..1207d87baa 100644 --- a/deepmd/dpmodel/descriptor/repformers.py +++ b/deepmd/dpmodel/descriptor/repformers.py @@ -79,6 +79,21 @@ def xp_transpose_01342(x: Array) -> Array: return x +def _statically_compatible_shape( + actual: tuple, + expected: tuple, +) -> bool: + if len(actual) != len(expected): + return False + static_int = (int, np.integer) + return all( + not isinstance(actual_dim, static_int) + or not isinstance(expected_dim, static_int) + or actual_dim == expected_dim + for actual_dim, expected_dim in zip(actual, expected, strict=True) + ) + + @DescriptorBlock.register("se_repformer") @DescriptorBlock.register("se_uni") class DescrptBlockRepformers(NativeOP, DescriptorBlock): @@ -504,7 +519,13 @@ def _exchange_ghosts( "implementation; pass a valid mapping or override the method " "for parallel comm handling." ) - return xp_take_along_axis(g1, mapping_tiled, axis=1) + xp = array_api_compat.array_namespace(g1, mapping_tiled) + mapping_mask = mapping_tiled >= 0 + mapping_tiled = xp.where( + mapping_mask, mapping_tiled, xp.zeros_like(mapping_tiled) + ) + g1_ext = xp_take_along_axis(g1, mapping_tiled, axis=1) + return xp.where(mapping_mask, g1_ext, xp.zeros_like(g1_ext)) def call( self, @@ -536,7 +557,7 @@ def call( sw = xp.where(nlist_mask, sw, xp.zeros_like(sw)) # nf x nloc x tebd_dim atype_embd = xp_take_first_n(atype_embd_ext, 1, nloc) - assert list(atype_embd.shape) == [nf, nloc, self.g1_dim] + assert _statically_compatible_shape(atype_embd.shape, (nf, nloc, self.g1_dim)) g1 = self.act(atype_embd) # nf x nloc x nnei x 1, nf x nloc x nnei x 3 @@ -1827,8 +1848,8 @@ def call( nf, nloc, nnei, _ = g2.shape # g1, _ = xp.split(g1_ext, [nloc], axis=1) g1 = xp_take_first_n(g1_ext, 1, nloc) - assert (nf, nloc) == g1.shape[:2] - assert (nf, nloc, nnei) == h2.shape[:3] + assert _statically_compatible_shape(g1.shape[:2], (nf, nloc)) + assert _statically_compatible_shape(h2.shape[:3], (nf, nloc, nnei)) g2_update: list[Array] = [g2] h2_update: list[Array] = [h2] diff --git a/deepmd/dpmodel/model/make_model.py b/deepmd/dpmodel/model/make_model.py index aba6b9fd48..7ee69b0132 100644 --- a/deepmd/dpmodel/model/make_model.py +++ b/deepmd/dpmodel/model/make_model.py @@ -17,8 +17,6 @@ from deepmd.dpmodel.array_api import ( Array, - xp_take_along_axis, - xp_take_first_n, ) from deepmd.dpmodel.atomic_model.base_atomic_model import ( BaseAtomicModel, @@ -40,6 +38,7 @@ from deepmd.dpmodel.utils import ( DefaultNeighborList, NeighborList, + format_nlist, nlist_distinguish_types, ) from deepmd.utils.path import ( @@ -611,72 +610,13 @@ def _format_nlist( nnei: int, extra_nlist_sort: bool = False, ) -> Array: - xp = array_api_compat.array_namespace(extended_coord, nlist) - n_nf, n_nloc, n_nnei = nlist.shape - extended_coord = extended_coord.reshape([n_nf, -1, 3]) - rcut = self.get_rcut() - - if n_nnei < nnei: - # make a copy before revise - ret = xp.concat( - [ - nlist, - -1 - * xp.ones( - [n_nf, n_nloc, nnei - n_nnei], - dtype=nlist.dtype, - device=array_api_compat.device(nlist), - ), - ], - axis=-1, - ) - - # Order matters for torch.export: Python evaluates `or` left-to-right - # with short-circuit. When `extra_nlist_sort=True` (Python bool) is - # on the left, the right-hand `n_nnei > nnei` is not evaluated, so no - # symbolic guard is registered on the dynamic `n_nnei` dimension. - # Swapping the operands would force the SymInt comparison to run and - # emit an `_assert_scalar` node in the exported graph. - if extra_nlist_sort or n_nnei > nnei: - n_nf, n_nloc, n_nnei = nlist.shape - # make a copy before revise - m_real_nei = nlist >= 0 - ret = xp.where(m_real_nei, nlist, 0) - coord0 = xp_take_first_n(extended_coord, 1, n_nloc) - index = xp.tile(ret.reshape(n_nf, n_nloc * n_nnei, 1), (1, 1, 3)) - coord1 = xp_take_along_axis(extended_coord, index, axis=1) - coord1 = coord1.reshape(n_nf, n_nloc, n_nnei, 3) - rr = xp.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1) - rr = xp.where(m_real_nei, rr, float("inf")) - rr, ret_mapping = xp.sort(rr, axis=-1), xp.argsort(rr, axis=-1) - ret = xp_take_along_axis(ret, ret_mapping, axis=2) - ret = xp.where(rr > rcut, -1, ret) - ret = ret[..., :nnei] - else: - # not extra_nlist_sort and n_nnei <= nnei: no reordering is - # needed (these descriptors reduce over neighbors order- - # independently), but we must still drop neighbors beyond rcut. - # The C++/LAMMPS neighbor list is built with rcut+skin and is - # NOT rcut-filtered before forward_lower; without this, out-of- - # rcut neighbors leak into the descriptor whenever the per-atom - # neighbor count <= nnei (this branch), making the result - # order-dependent (see discussion #5438). - if n_nnei == nnei: - ret = nlist - # else (n_nnei < nnei): `ret` is already padded to nnei above. - n_nf, n_nloc, n_pad = ret.shape - m_real_nei = ret >= 0 - coord0 = xp_take_first_n(extended_coord, 1, n_nloc) - index = xp.tile( - xp.where(m_real_nei, ret, 0).reshape(n_nf, n_nloc * n_pad, 1), - (1, 1, 3), - ) - coord1 = xp_take_along_axis(extended_coord, index, axis=1) - coord1 = coord1.reshape(n_nf, n_nloc, n_pad, 3) - rr = xp.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1) - ret = xp.where(m_real_nei & (rr > rcut), -1, ret) - assert ret.shape[-1] == nnei - return ret + return format_nlist( + extended_coord, + nlist, + nnei, + self.get_rcut(), + extra_nlist_sort=extra_nlist_sort, + ) def do_grad_r( self, diff --git a/deepmd/dpmodel/utils/__init__.py b/deepmd/dpmodel/utils/__init__.py index 0179543dd4..28aee2a4d6 100644 --- a/deepmd/dpmodel/utils/__init__.py +++ b/deepmd/dpmodel/utils/__init__.py @@ -48,6 +48,7 @@ build_multiple_neighbor_list, build_neighbor_list, extend_coord_with_ghosts, + format_nlist, get_multiple_nlist_key, nlist_distinguish_types, ) @@ -93,6 +94,7 @@ "compute_total_numb_batch", "edge_force_virial", "extend_coord_with_ghosts", + "format_nlist", "from_dense_quartet", "get_graph_index", "get_multiple_nlist_key", diff --git a/deepmd/dpmodel/utils/nlist.py b/deepmd/dpmodel/utils/nlist.py index b7b493f342..59e68a64a0 100644 --- a/deepmd/dpmodel/utils/nlist.py +++ b/deepmd/dpmodel/utils/nlist.py @@ -1,5 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + import array_api_compat from deepmd.dpmodel.array_api import ( @@ -14,6 +18,27 @@ ) +def _is_ndtensorflow_namespace(xp: Any) -> bool: + return getattr(xp, "__name__", "") == "deepmd._vendors.ndtensorflow" + + +def _arange_nbuff(nbuff: Array, index: int, xp: Any, device: Any) -> Array: + bound = nbuff[index] + if not _is_ndtensorflow_namespace(xp): + bound = int(bound) + return xp.arange(-bound, bound + 1, 1, dtype=xp.int64, device=device) + + +def _size(x: Array, xp: Any) -> Any: + if _is_ndtensorflow_namespace(xp): + return x.size + return array_api_compat.size(x) + + +def _is_static_shape(shape: Any) -> bool: + return all(isinstance(dim, int) for dim in shape) + + def extend_input_and_build_neighbor_list( coord: Array, atype: Array, @@ -97,10 +122,17 @@ def build_neighbor_list( nall = coord.shape[1] // 3 # fill virtual atoms with large coords so they are not neighbors of any # real atom. - if array_api_compat.size(coord) > 0: + if _size(coord, xp) > 0: xmax = xp.max(coord) + 2.0 * rcut else: - xmax = 2.0 * rcut + if _is_ndtensorflow_namespace(xp): + xmax = xp.asarray( + 2.0 * rcut, + dtype=coord.dtype, + device=array_api_compat.device(coord), + ) + else: + xmax = 2.0 * rcut # nf x nall is_vir = atype < 0 coord1 = xp.where( @@ -115,13 +147,14 @@ def build_neighbor_list( xp.reshape(coord1, (batch_size, -1, 3))[:, None, :, :] - xp.reshape(coord0, (batch_size, -1, 3))[:, :, None, :] ) - assert list(diff.shape) == [batch_size, nloc, nall, 3] + if _is_static_shape(diff.shape): + assert list(diff.shape) == [batch_size, nloc, nall, 3] rr = xp.linalg.vector_norm(diff, axis=-1) # if central atom has two zero distances, sorting sometimes can not exclude itself rr -= xp.eye(nloc, nall, dtype=diff.dtype, device=array_api_compat.device(diff))[ xp.newaxis, :, : ] - nlist = xp.argsort(rr, axis=-1) + nlist = xp.astype(xp.argsort(rr, axis=-1), xp.int64) rr = xp.sort(rr, axis=-1) rr = rr[:, :, 1:] nlist = nlist[:, :, 1:] @@ -130,7 +163,7 @@ def build_neighbor_list( rr = rr[:, :, :nsel] nlist = nlist[:, :, :nsel] else: - rr = xp.concatenate( + rr = xp.concat( [ rr, xp.ones( @@ -142,7 +175,7 @@ def build_neighbor_list( ], axis=-1, ) - nlist = xp.concatenate( + nlist = xp.concat( [ nlist, xp.ones( @@ -153,7 +186,8 @@ def build_neighbor_list( ], axis=-1, ) - assert list(nlist.shape) == [batch_size, nloc, nsel] + if _is_static_shape(nlist.shape): + assert list(nlist.shape) == [batch_size, nloc, nsel] nlist = xp.where( xp.logical_or((rr > rcut), is_vir[:, :nloc, None]), xp.full_like(nlist, -1), @@ -196,6 +230,113 @@ def nlist_distinguish_types( return ret +def format_nlist( + extended_coord: Array, + nlist: Array, + nnei: int, + rcut: float, + extra_nlist_sort: bool = False, +) -> Array: + """Format a neighbor list to a fixed neighbor count. + + If the input neighbor axis is shorter than ``nnei``, pad it with ``-1``. + If the input neighbor axis is longer than ``nnei``, sort neighbors by + distance, mask neighbors outside ``rcut`` with ``-1``, and truncate to + ``nnei`` entries. Otherwise, preserve the input order and mask neighbors + outside ``rcut`` with ``-1``. When ``extra_nlist_sort`` is true, use the + sort-and-truncate path even when the input neighbor axis is not longer than + ``nnei``. + + Parameters + ---------- + extended_coord : Array + Extended coordinates of shape ``[nf, nall, 3]`` or + ``[nf, nall * 3]``. + nlist : Array + Neighbor list of shape ``[nf, nloc, n_nnei]``. Invalid neighbor + entries are marked with ``-1``. + nnei : int + Target number of selected neighbors. + rcut : float + Cutoff radius. Neighbors farther than ``rcut`` are marked with ``-1``. + extra_nlist_sort : bool, optional + Whether to force distance sorting and truncation even when the input + neighbor axis is not larger than ``nnei``. This is needed by models + whose lower-level forward path requires a sorted neighbor list. + + Returns + ------- + Array + Formatted neighbor list of shape ``[nf, nloc, nnei]``. Missing or + out-of-cutoff neighbors are marked with ``-1``. + """ + xp = array_api_compat.array_namespace(extended_coord, nlist) + n_nf, n_nloc, n_nnei = nlist.shape + extended_coord = extended_coord.reshape([n_nf, -1, 3]) + ret = nlist + + if n_nnei < nnei: + ret = xp.concat( + [ + nlist, + -1 + * xp.ones( + [n_nf, n_nloc, nnei - n_nnei], + dtype=nlist.dtype, + device=array_api_compat.device(nlist), + ), + ], + axis=-1, + ) + + # Order matters for torch.export: Python evaluates `or` left-to-right + # with short-circuit. When `extra_nlist_sort=True` (Python bool) is + # on the left, the right-hand `n_nnei > nnei` is not evaluated, so no + # symbolic guard is registered on the dynamic `n_nnei` dimension. + # Swapping the operands would force the SymInt comparison to run and + # emit an `_assert_scalar` node in the exported graph. + if extra_nlist_sort or n_nnei > nnei: + n_nf, n_nloc, n_nnei = nlist.shape + m_real_nei = nlist >= 0 + ret = xp.where(m_real_nei, nlist, 0) + coord0 = xp_take_first_n(extended_coord, 1, n_nloc) + index = xp.tile(ret.reshape(n_nf, n_nloc * n_nnei, 1), (1, 1, 3)) + coord1 = xp_take_along_axis(extended_coord, index, axis=1) + coord1 = coord1.reshape(n_nf, n_nloc, n_nnei, 3) + rr = xp.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1) + rr = xp.where(m_real_nei, rr, float("inf")) + rr, ret_mapping = xp.sort(rr, axis=-1), xp.argsort(rr, axis=-1) + ret = xp_take_along_axis(ret, ret_mapping, axis=2) + ret = xp.where(rr > rcut, -1, ret) + ret = ret[..., :nnei] + else: + # not extra_nlist_sort and n_nnei <= nnei: no reordering is + # needed (these descriptors reduce over neighbors order- + # independently), but we must still drop neighbors beyond rcut. + # The C++/LAMMPS neighbor list is built with rcut+skin and is + # NOT rcut-filtered before forward_lower; without this, out-of- + # rcut neighbors leak into the descriptor whenever the per-atom + # neighbor count <= nnei (this branch), making the result + # order-dependent (see discussion #5438). + if n_nnei == nnei: + ret = nlist + # else (n_nnei < nnei): `ret` is already padded to nnei above. + n_nf, n_nloc, n_pad = ret.shape + m_real_nei = ret >= 0 + coord0 = xp_take_first_n(extended_coord, 1, n_nloc) + index = xp.tile( + xp.where(m_real_nei, ret, 0).reshape(n_nf, n_nloc * n_pad, 1), + (1, 1, 3), + ) + coord1 = xp_take_along_axis(extended_coord, index, axis=1) + coord1 = coord1.reshape(n_nf, n_nloc, n_pad, 3) + rr = xp.linalg.norm(coord0[:, :, None, :] - coord1, axis=-1) + ret = xp.where(m_real_nei & (rr > rcut), -1, ret) + if isinstance(ret.shape[-1], int): + assert ret.shape[-1] == nnei + return ret + + def get_multiple_nlist_key(rcut: float, nsel: int) -> str: return str(rcut) + "_" + str(nsel) @@ -314,27 +455,10 @@ def extend_coord_with_ghosts( to_face = to_face_distance(cell) nbuff = xp.astype(xp.ceil(rcut / to_face), xp.int64) nbuff = xp.max(nbuff, axis=0) - xi = xp.arange( - -int(nbuff[0]), - int(nbuff[0]) + 1, - 1, - dtype=xp.int64, - device=array_api_compat.device(coord), - ) - yi = xp.arange( - -int(nbuff[1]), - int(nbuff[1]) + 1, - 1, - dtype=xp.int64, - device=array_api_compat.device(coord), - ) - zi = xp.arange( - -int(nbuff[2]), - int(nbuff[2]) + 1, - 1, - dtype=xp.int64, - device=array_api_compat.device(coord), - ) + device = array_api_compat.device(coord) + xi = _arange_nbuff(nbuff, 0, xp, device) + yi = _arange_nbuff(nbuff, 1, xp, device) + zi = _arange_nbuff(nbuff, 2, xp, device) xyz = xp.linalg.outer( xi, xp.asarray([1, 0, 0], device=array_api_compat.device(xi)) )[:, xp.newaxis, xp.newaxis, :] diff --git a/deepmd/jax/jax2tf/__init__.py b/deepmd/jax/jax2tf/__init__.py index c2cda24bd7..0742bc3384 100644 --- a/deepmd/jax/jax2tf/__init__.py +++ b/deepmd/jax/jax2tf/__init__.py @@ -1,14 +1,12 @@ # SPDX-License-Identifier: LGPL-3.0-or-later import tensorflow as tf -import tensorflow.experimental.numpy as tnp if not tf.executing_eagerly(): # TF disallow temporary eager execution raise RuntimeError( - "Unfortunatly, jax2tf (requires eager execution) cannot be used with the " - "TensorFlow backend (disables eager execution). " + "The TensorFlow SavedModel compatibility layer requires eager execution. " + "It cannot be used with the TensorFlow v1 backend after eager execution " + "has been disabled. " "If you are converting a model between different backends, " - "considering converting to the `.dp` format first." + "consider converting to the `.dp` format first." ) - -tnp.experimental_enable_numpy_behavior() diff --git a/deepmd/jax/jax2tf/format_nlist.py b/deepmd/jax/jax2tf/format_nlist.py index 5cf93610e7..f9b216fb27 100644 --- a/deepmd/jax/jax2tf/format_nlist.py +++ b/deepmd/jax/jax2tf/format_nlist.py @@ -1,71 +1,24 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import tensorflow as tf -import tensorflow.experimental.numpy as tnp - +"""Compatibility wrappers for TensorFlow neighbor-list formatting.""" -@tf.function(autograph=True) -def format_nlist( - extended_coord: tnp.ndarray, - nlist: tnp.ndarray, - nsel: int, - rcut: float, -) -> tnp.ndarray: - """Format neighbor list. +from typing import ( + Any, +) - If nnei == nsel, do nothing; - If nnei < nsel, pad -1; - If nnei > nsel, sort by distance and truncate. +import tensorflow as tf - Parameters - ---------- - extended_coord - The extended coordinates of the atoms. - shape: nf x nall x 3 - nlist - The neighbor list. - shape: nf x nloc x nnei - nsel - The number of selected neighbors. - rcut - The cutoff radius. +from deepmd.tf2.common import ( + to_tf_tensor, +) +from deepmd.tf2.utils._dpmodel import format_nlist as tf2_format_nlist - Returns - ------- - nlist - The formatted neighbor list. - shape: nf x nloc x nsel - """ - nlist_shape = tf.shape(nlist) - n_nf, n_nloc, n_nsel = nlist_shape[0], nlist_shape[1], nlist_shape[2] - extended_coord = extended_coord.reshape([n_nf, -1, 3]) +__all__ = ["format_nlist"] - if n_nsel < nsel: - # make a copy before revise - ret = tnp.concatenate( - [ - nlist, - tnp.full([n_nf, n_nloc, nsel - n_nsel], -1, dtype=nlist.dtype), - ], - axis=-1, - ) - elif n_nsel > nsel: - # make a copy before revise - m_real_nei = nlist >= 0 - ret = tnp.where(m_real_nei, nlist, 0) - coord0 = extended_coord[:, :n_nloc, :] - index = ret.reshape(n_nf, n_nloc * n_nsel, 1) - index = tnp.repeat(index, 3, axis=2) - coord1 = tnp.take_along_axis(extended_coord, index, axis=1) - coord1 = coord1.reshape(n_nf, n_nloc, n_nsel, 3) - rr2 = tnp.sum(tnp.square(coord0[:, :, None, :] - coord1), axis=-1) - rr2 = tnp.where(m_real_nei, rr2, float("inf")) - rr2, ret_mapping = tnp.sort(rr2, axis=-1), tnp.argsort(rr2, axis=-1) - ret = tnp.take_along_axis(ret, ret_mapping, axis=2) - ret = tnp.where(rr2 > rcut * rcut, -1, ret) - ret = ret[..., :nsel] - else: # n_nsel == nsel: - ret = nlist - # do a reshape any way; this will tell the xla the shape without any dynamic shape - ret = tnp.reshape(ret, [n_nf, n_nloc, nsel]) - return ret +def format_nlist( + extended_coord: Any, + nlist: Any, + nsel: int, + rcut: float, +) -> tf.Tensor: + return to_tf_tensor(tf2_format_nlist(extended_coord, nlist, nsel, rcut)) diff --git a/deepmd/jax/jax2tf/make_model.py b/deepmd/jax/jax2tf/make_model.py index 3cd30de85a..ebb93492a9 100644 --- a/deepmd/jax/jax2tf/make_model.py +++ b/deepmd/jax/jax2tf/make_model.py @@ -1,112 +1,77 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +"""Compatibility wrappers for TensorFlow model-call helpers.""" + from collections.abc import ( Callable, ) +from typing import ( + Any, +) import tensorflow as tf -import tensorflow.experimental.numpy as tnp from deepmd.dpmodel.output_def import ( ModelOutputDef, ) -from deepmd.jax.jax2tf.nlist import ( - build_neighbor_list, - extend_coord_with_ghosts, +from deepmd.tf2.common import ( + to_tf_tensor, + unwrap_value, + wrap_value, ) -from deepmd.jax.jax2tf.region import ( - normalize_coord, -) -from deepmd.jax.jax2tf.transform_output import ( - communicate_extended_output, +from deepmd.tf2.make_model import ( + model_call_from_call_lower as tf2_model_call_from_call_lower, ) +__all__ = ["model_call_from_call_lower"] + + +def _wrap_call_lower(call_lower: Callable[..., dict[str, Any]]) -> Callable: + def wrapped_call_lower( + extended_coord: Any, + extended_atype: Any, + nlist: Any, + mapping: Any, + **kwargs: Any, + ) -> dict[str, Any]: + return wrap_value( + call_lower( + to_tf_tensor(extended_coord), + to_tf_tensor(extended_atype), + to_tf_tensor(nlist), + to_tf_tensor(mapping), + **{kk: to_tf_tensor(vv) for kk, vv in kwargs.items()}, + ) + ) + + return wrapped_call_lower + def model_call_from_call_lower( *, # enforce keyword-only arguments - call_lower: Callable[ - [ - tnp.ndarray, - tnp.ndarray, - tnp.ndarray, - tnp.ndarray, - tnp.ndarray, - bool, - ], - dict[str, tnp.ndarray], - ], + call_lower: Callable[..., dict[str, Any]], rcut: float, sel: list[int], mixed_types: bool, model_output_def: ModelOutputDef, - coord: tnp.ndarray, - atype: tnp.ndarray, - box: tnp.ndarray, - fparam: tnp.ndarray, - aparam: tnp.ndarray, + coord: tf.Tensor, + atype: tf.Tensor, + box: tf.Tensor, + fparam: tf.Tensor, + aparam: tf.Tensor, do_atomic_virial: bool = False, -) -> dict[str, tnp.ndarray]: - """Return model prediction from lower interface. - - Parameters - ---------- - coord - The coordinates of the atoms. - shape: nf x (nloc x 3) - atype - The type of atoms. shape: nf x nloc - box - The simulation box. shape: nf x 9 - fparam - frame parameter. nf x ndf - aparam - atomic parameter. nf x nloc x nda - do_atomic_virial - If calculate the atomic virial. - - Returns - ------- - ret_dict - The result dict of type dict[str,tnp.ndarray]. - The keys are defined by the `ModelOutputDef`. - - """ - atype_shape = tf.shape(atype) - nframes, nloc = atype_shape[0], atype_shape[1] - cc, bb, fp, ap = coord, box, fparam, aparam - del coord, box, fparam, aparam - if tf.shape(bb)[-1] != 0: - coord_normalized = normalize_coord( - cc.reshape(nframes, nloc, 3), - bb.reshape(nframes, 3, 3), +) -> dict[str, tf.Tensor]: + return unwrap_value( + tf2_model_call_from_call_lower( + call_lower=_wrap_call_lower(call_lower), + rcut=rcut, + sel=sel, + mixed_types=mixed_types, + model_output_def=model_output_def, + coord=coord, + atype=atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, ) - else: - coord_normalized = cc - extended_coord, extended_atype, mapping = extend_coord_with_ghosts( - coord_normalized, atype, bb, rcut - ) - nlist = build_neighbor_list( - extended_coord, - extended_atype, - nloc, - rcut, - sel, - # types will be distinguished in the lower interface, - # so it doesn't need to be distinguished here - distinguish_types=False, - ) - extended_coord = extended_coord.reshape(nframes, -1, 3) - model_predict_lower = call_lower( - extended_coord, - extended_atype, - nlist, - mapping, - fparam=fp, - aparam=ap, - ) - model_predict = communicate_extended_output( - model_predict_lower, - model_output_def, - mapping, - do_atomic_virial=do_atomic_virial, ) - return model_predict diff --git a/deepmd/jax/jax2tf/nlist.py b/deepmd/jax/jax2tf/nlist.py index c44a1196c8..9ba8fa78cc 100644 --- a/deepmd/jax/jax2tf/nlist.py +++ b/deepmd/jax/jax2tf/nlist.py @@ -1,214 +1,53 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +"""Compatibility wrappers for TensorFlow neighbor-list helpers.""" + +from typing import ( + Any, +) import tensorflow as tf -import tensorflow.experimental.numpy as tnp -from .region import ( - to_face_distance, +from deepmd.tf2.common import ( + to_tf_tensor, +) +from deepmd.tf2.utils._dpmodel import build_neighbor_list as tf2_build_neighbor_list +from deepmd.tf2.utils._dpmodel import ( + extend_coord_with_ghosts as tf2_extend_coord_with_ghosts, ) +__all__ = [ + "build_neighbor_list", + "extend_coord_with_ghosts", +] + -## translated from torch implementation by chatgpt def build_neighbor_list( - coord: tnp.ndarray, - atype: tnp.ndarray, + coord: Any, + atype: Any, nloc: int, rcut: float, sel: int | list[int], distinguish_types: bool = True, -) -> tnp.ndarray: - """Build neighbor list for a single frame. keeps nsel neighbors. - - Parameters - ---------- - coord : tnp.ndarray - exptended coordinates of shape [batch_size, nall x 3] - atype : tnp.ndarray - extended atomic types of shape [batch_size, nall] - type < 0 the atom is treat as virtual atoms. - nloc : int - number of local atoms. - rcut : float - cut-off radius - sel : int or list[int] - maximal number of neighbors (of each type). - if distinguish_types==True, nsel should be list and - the length of nsel should be equal to number of - types. - distinguish_types : bool - distinguish different types. - - Returns - ------- - neighbor_list : tnp.ndarray - Neighbor list of shape [batch_size, nloc, nsel], the neighbors - are stored in an ascending order. If the number of - neighbors is less than nsel, the positions are masked - with -1. The neighbor list of an atom looks like - |------ nsel ------| - xx xx xx xx -1 -1 -1 - if distinguish_types==True and we have two types - |---- nsel[0] -----| |---- nsel[1] -----| - xx xx xx xx -1 -1 -1 xx xx xx -1 -1 -1 -1 - For virtual atoms all neighboring positions are filled with -1. - - """ - batch_size = tf.shape(coord)[0] - coord = tnp.reshape(coord, (batch_size, -1)) - nall = tf.shape(coord)[1] // 3 - # fill virtual atoms with large coords so they are not neighbors of any - # real atom. - if tf.size(coord) > 0: - xmax = tnp.max(coord) + 2.0 * rcut - else: - xmax = tf.cast(2.0 * rcut, coord.dtype) - # nf x nall - is_vir = atype < 0 - coord1 = tnp.where( - is_vir[:, :, None], xmax, tnp.reshape(coord, (batch_size, nall, 3)) - ) - coord1 = tnp.reshape(coord1, (batch_size, nall * 3)) - if isinstance(sel, int): - sel = [sel] - nsel = sum(sel) - coord0 = coord1[:, : nloc * 3] - diff = ( - tnp.reshape(coord1, [batch_size, -1, 3])[:, None, :, :] - - tnp.reshape(coord0, [batch_size, -1, 3])[:, :, None, :] - ) - rr = tf.linalg.norm(diff, axis=-1) - # if central atom has two zero distances, sorting sometimes can not exclude itself - rr -= tf.eye(nloc, nall, dtype=diff.dtype)[tnp.newaxis, :, :] - nlist = tnp.argsort(rr, axis=-1) - rr = tnp.sort(rr, axis=-1) - rr = rr[:, :, 1:] - nlist = nlist[:, :, 1:] - nnei = tf.shape(rr)[2] - if nsel <= nnei: - rr = rr[:, :, :nsel] - nlist = nlist[:, :, :nsel] - else: - rr = tnp.concatenate( - [rr, tnp.ones([batch_size, nloc, nsel - nnei], dtype=rr.dtype) + rcut], - axis=-1, +) -> tf.Tensor: + return to_tf_tensor( + tf2_build_neighbor_list( + coord, + atype, + nloc, + rcut, + sel, + distinguish_types=distinguish_types, ) - nlist = tnp.concatenate( - [nlist, tnp.ones([batch_size, nloc, nsel - nnei], dtype=nlist.dtype)], - axis=-1, - ) - nlist = tnp.where( - tnp.logical_or((rr > rcut), is_vir[:, :nloc, None]), - tnp.full_like(nlist, -1), - nlist, ) - if distinguish_types: - return nlist_distinguish_types(nlist, atype, sel) - else: - return nlist - -def nlist_distinguish_types( - nlist: tnp.ndarray, - atype: tnp.ndarray, - sel: list[int], -) -> tnp.ndarray: - """Given a nlist that does not distinguish atom types, return a nlist that - distinguish atom types. - - """ - nloc = tf.shape(nlist)[1] - ret_nlist = [] - tmp_atype = tnp.tile(atype[:, None, :], (1, nloc, 1)) - mask = nlist == -1 - tnlist_0 = tnp.where(mask, tnp.zeros_like(nlist), nlist) - tnlist = tnp.take_along_axis(tmp_atype, tnlist_0, axis=2) - tnlist = tnp.where(mask, tnp.full_like(tnlist, -1), tnlist) - for ii, ss in enumerate(sel): - pick_mask = tf.cast(tnlist == ii, tnp.int32) - sorted_indices = tnp.argsort(-pick_mask, kind="stable", axis=-1) - pick_mask_sorted = -tnp.sort(-pick_mask, axis=-1) - inlist = tnp.take_along_axis(nlist, sorted_indices, axis=2) - inlist = tnp.where( - ~tf.cast(pick_mask_sorted, tf.bool), tnp.full_like(inlist, -1), inlist - ) - ret_nlist.append(inlist[..., :ss]) - ret = tf.concat(ret_nlist, axis=-1) - return ret - - -def tf_outer(a: tnp.ndarray, b: tnp.ndarray) -> tnp.ndarray: - return tf.einsum("i,j->ij", a, b) - - -## translated from torch implementation by chatgpt def extend_coord_with_ghosts( - coord: tnp.ndarray, - atype: tnp.ndarray, - cell: tnp.ndarray, + coord: Any, + atype: Any, + cell: Any | None, rcut: float, -) -> tuple[tnp.ndarray, tnp.ndarray, tnp.ndarray]: - """Extend the coordinates of the atoms by appending peridoc images. - The number of images is large enough to ensure all the neighbors - within rcut are appended. - - Parameters - ---------- - coord : tnp.ndarray - original coordinates of shape [-1, nloc*3]. - atype : tnp.ndarray - atom type of shape [-1, nloc]. - cell : tnp.ndarray - simulation cell tensor of shape [-1, 9]. - rcut : float - the cutoff radius - - Returns - ------- - extended_coord: tnp.ndarray - extended coordinates of shape [-1, nall*3]. - extended_atype: tnp.ndarray - extended atom type of shape [-1, nall]. - index_mapping: tnp.ndarray - mapping extended index to the local index - - """ - atype_shape = tf.shape(atype) - nf, nloc = atype_shape[0], atype_shape[1] - # int64 for index - aidx = tf.range(nloc, dtype=tnp.int64) - aidx = tnp.tile(aidx[tnp.newaxis, :], (nf, 1)) - if tf.shape(cell)[-1] == 0: - nall = nloc - extend_coord = coord - extend_atype = atype - extend_aidx = aidx - else: - coord = tnp.reshape(coord, (nf, nloc, 3)) - cell = tnp.reshape(cell, (nf, 3, 3)) - to_face = to_face_distance(cell) - nbuff = tf.cast(tnp.ceil(rcut / to_face), tnp.int64) - nbuff = tnp.max(nbuff, axis=0) - xi = tf.range(-nbuff[0], nbuff[0] + 1, 1, dtype=tnp.int64) - yi = tf.range(-nbuff[1], nbuff[1] + 1, 1, dtype=tnp.int64) - zi = tf.range(-nbuff[2], nbuff[2] + 1, 1, dtype=tnp.int64) - xyz = tf_outer(xi, tnp.asarray([1, 0, 0]))[:, tnp.newaxis, tnp.newaxis, :] - xyz = xyz + tf_outer(yi, tnp.asarray([0, 1, 0]))[tnp.newaxis, :, tnp.newaxis, :] - xyz = xyz + tf_outer(zi, tnp.asarray([0, 0, 1]))[tnp.newaxis, tnp.newaxis, :, :] - xyz = tnp.reshape(xyz, (-1, 3)) - xyz = tf.cast(xyz, coord.dtype) - shift_idx = tnp.take(xyz, tnp.argsort(tf.linalg.norm(xyz, axis=1)), axis=0) - ns = tf.shape(shift_idx)[0] - nall = ns * nloc - shift_vec = tnp.einsum("sd,fdk->fsk", shift_idx, cell) - # shift_vec = tnp.tensordot(shift_idx, cell, axes=([1], [1])) - # shift_vec = tnp.transpose(shift_vec, (1, 0, 2)) - extend_coord = coord[:, None, :, :] + shift_vec[:, :, None, :] - extend_atype = tnp.tile(atype[:, :, tnp.newaxis], (1, ns, 1)) - extend_aidx = tnp.tile(aidx[:, :, tnp.newaxis], (1, ns, 1)) - - return ( - tnp.reshape(extend_coord, (nf, nall * 3)), - tnp.reshape(extend_atype, (nf, nall)), - tnp.reshape(extend_aidx, (nf, nall)), +) -> tuple[tf.Tensor, tf.Tensor, tf.Tensor]: + return tuple( + to_tf_tensor(value) + for value in tf2_extend_coord_with_ghosts(coord, atype, cell, rcut) ) diff --git a/deepmd/jax/jax2tf/region.py b/deepmd/jax/jax2tf/region.py index a90e693478..3c80277f2e 100644 --- a/deepmd/jax/jax2tf/region.py +++ b/deepmd/jax/jax2tf/region.py @@ -1,104 +1,33 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import tensorflow as tf -import tensorflow.experimental.numpy as tnp - - -def phys2inter( - coord: tnp.ndarray, - cell: tnp.ndarray, -) -> tnp.ndarray: - """Convert physical coordinates to internal(direct) coordinates. - - Parameters - ---------- - coord : tnp.ndarray - physical coordinates of shape [*, na, 3]. - cell : tnp.ndarray - simulation cell tensor of shape [*, 3, 3]. - - Returns - ------- - inter_coord: tnp.ndarray - the internal coordinates - - """ - rec_cell = tf.linalg.inv(cell) - return tnp.matmul(coord, rec_cell) - +"""Compatibility wrappers for TensorFlow region helpers.""" -def inter2phys( - coord: tnp.ndarray, - cell: tnp.ndarray, -) -> tnp.ndarray: - """Convert internal(direct) coordinates to physical coordinates. +from typing import ( + Any, +) - Parameters - ---------- - coord : tnp.ndarray - internal coordinates of shape [*, na, 3]. - cell : tnp.ndarray - simulation cell tensor of shape [*, 3, 3]. - - Returns - ------- - phys_coord: tnp.ndarray - the physical coordinates - - """ - return tnp.matmul(coord, cell) - - -def normalize_coord( - coord: tnp.ndarray, - cell: tnp.ndarray, -) -> tnp.ndarray: - """Apply PBC according to the atomic coordinates. - - Parameters - ---------- - coord : tnp.ndarray - original coordinates of shape [*, na, 3]. - cell : tnp.ndarray - simulation cell shape [*, 3, 3]. - - Returns - ------- - wrapped_coord: tnp.ndarray - wrapped coordinates of shape [*, na, 3]. +import tensorflow as tf - """ - icoord = phys2inter(coord, cell) - icoord = tnp.remainder(icoord, 1.0) - return inter2phys(icoord, cell) +from deepmd.tf2.common import ( + to_tf_tensor, +) +from deepmd.tf2.utils._dpmodel import inter2phys as tf2_inter2phys +from deepmd.tf2.utils._dpmodel import normalize_coord as tf2_normalize_coord +from deepmd.tf2.utils._dpmodel import to_face_distance as tf2_to_face_distance +__all__ = [ + "inter2phys", + "normalize_coord", + "to_face_distance", +] -def to_face_distance( - cell: tnp.ndarray, -) -> tnp.ndarray: - """Compute the to-face-distance of the simulation cell. - Parameters - ---------- - cell : tnp.ndarray - simulation cell tensor of shape [*, 3, 3]. +def inter2phys(coord: Any, cell: Any) -> tf.Tensor: + return to_tf_tensor(tf2_inter2phys(coord, cell)) - Returns - ------- - dist: tnp.ndarray - the to face distances of shape [*, 3] - """ - cshape = tf.shape(cell) - dist = b_to_face_distance(tnp.reshape(cell, [-1, 3, 3])) - return tnp.reshape(dist, tf.concat([cshape[:-2], [3]], axis=0)) +def normalize_coord(coord: Any, cell: Any) -> tf.Tensor: + return to_tf_tensor(tf2_normalize_coord(coord, cell)) -def b_to_face_distance(cell: tnp.ndarray) -> tnp.ndarray: - volume = tf.linalg.det(cell) - c_yz = tf.linalg.cross(cell[:, 1, ...], cell[:, 2, ...]) - _h2yz = volume / tf.linalg.norm(c_yz, axis=-1) - c_zx = tf.linalg.cross(cell[:, 2, ...], cell[:, 0, ...]) - _h2zx = volume / tf.linalg.norm(c_zx, axis=-1) - c_xy = tf.linalg.cross(cell[:, 0, ...], cell[:, 1, ...]) - _h2xy = volume / tf.linalg.norm(c_xy, axis=-1) - return tnp.stack([_h2yz, _h2zx, _h2xy], axis=1) +def to_face_distance(cell: Any) -> tf.Tensor: + return to_tf_tensor(tf2_to_face_distance(cell)) diff --git a/deepmd/jax/jax2tf/serialization.py b/deepmd/jax/jax2tf/serialization.py index 4881ca98f8..b26aeb2c40 100644 --- a/deepmd/jax/jax2tf/serialization.py +++ b/deepmd/jax/jax2tf/serialization.py @@ -1,343 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import json -from collections.abc import ( - Callable, -) - -import tensorflow as tf -import tensorflow.experimental.numpy as tnp -from jax.experimental import ( - jax2tf, -) +"""Compatibility wrapper for the TF2 SavedModel exporter.""" -from deepmd.jax.jax2tf.format_nlist import ( - format_nlist, -) -from deepmd.jax.jax2tf.make_model import ( - model_call_from_call_lower, +from deepmd.tf2.utils.serialization import ( + deserialize_to_file, ) -from deepmd.jax.model.base_model import ( - BaseModel, -) - - -def deserialize_to_file(model_file: str, data: dict) -> None: - """Deserialize the dictionary to a model file. - - Parameters - ---------- - model_file : str - The model file to be saved. - data : dict - The dictionary to be deserialized. - """ - if model_file.endswith(".savedmodel"): - model = BaseModel.deserialize(data["model"]) - model_def_script = data["model_def_script"] - call_lower = model.call_common_lower - - tf_model = tf.Module() - - def exported_whether_do_atomic_virial( - do_atomic_virial: bool, has_ghost_atoms: bool - ) -> Callable: - def call_lower_with_fixed_do_atomic_virial( - coord: tnp.ndarray, - atype: tnp.ndarray, - nlist: tnp.ndarray, - mapping: tnp.ndarray, - fparam: tnp.ndarray, - aparam: tnp.ndarray, - ) -> dict[str, tnp.ndarray]: - return call_lower( - coord, - atype, - nlist, - mapping, - fparam, - aparam, - do_atomic_virial=do_atomic_virial, - ) - - # nghost >= 1 is assumed if there is - # other workaround does not work, such as - # nall; nloc + nghost - 1 - if has_ghost_atoms: - nghost = "nghost" - else: - nghost = "0" - return jax2tf.convert( - call_lower_with_fixed_do_atomic_virial, - polymorphic_shapes=[ - f"(nf, nloc + {nghost}, 3)", - f"(nf, nloc + {nghost})", - f"(nf, nloc, {model.get_nnei()})", - f"(nf, nloc + {nghost})", - f"(nf, {model.get_dim_fparam()})", - f"(nf, nloc, {model.get_dim_aparam()})", - ], - with_gradient=True, - ) - - # Save a function that can take scalar inputs. - # We need to explicit set the function name, so C++ can find it. - @tf.function( - autograph=False, - input_signature=[ - tf.TensorSpec([None, None, 3], tf.float64), - tf.TensorSpec([None, None], tf.int32), - tf.TensorSpec([None, None, None], tf.int64), - tf.TensorSpec([None, None], tf.int64), - tf.TensorSpec([None, model.get_dim_fparam()], tf.float64), - tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), - ], - ) - def call_lower_without_atomic_virial( - coord: tnp.ndarray, - atype: tnp.ndarray, - nlist: tnp.ndarray, - mapping: tnp.ndarray, - fparam: tnp.ndarray, - aparam: tnp.ndarray, - ) -> dict[str, tnp.ndarray]: - nlist = format_nlist(coord, nlist, model.get_nnei(), model.get_rcut()) - return tf.cond( - tf.shape(coord)[1] == tf.shape(nlist)[1], - lambda: exported_whether_do_atomic_virial( - do_atomic_virial=False, has_ghost_atoms=False - )(coord, atype, nlist, mapping, fparam, aparam), - lambda: exported_whether_do_atomic_virial( - do_atomic_virial=False, has_ghost_atoms=True - )(coord, atype, nlist, mapping, fparam, aparam), - ) - - tf_model.call_lower = call_lower_without_atomic_virial - - @tf.function( - autograph=False, - input_signature=[ - tf.TensorSpec([None, None, 3], tf.float64), - tf.TensorSpec([None, None], tf.int32), - tf.TensorSpec([None, None, None], tf.int64), - tf.TensorSpec([None, None], tf.int64), - tf.TensorSpec([None, model.get_dim_fparam()], tf.float64), - tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), - ], - ) - def call_lower_with_atomic_virial( - coord: tnp.ndarray, - atype: tnp.ndarray, - nlist: tnp.ndarray, - mapping: tnp.ndarray, - fparam: tnp.ndarray, - aparam: tnp.ndarray, - ) -> dict[str, tnp.ndarray]: - nlist = format_nlist(coord, nlist, model.get_nnei(), model.get_rcut()) - return tf.cond( - tf.shape(coord)[1] == tf.shape(nlist)[1], - lambda: exported_whether_do_atomic_virial( - do_atomic_virial=True, has_ghost_atoms=False - )(coord, atype, nlist, mapping, fparam, aparam), - lambda: exported_whether_do_atomic_virial( - do_atomic_virial=True, has_ghost_atoms=True - )(coord, atype, nlist, mapping, fparam, aparam), - ) - - tf_model.call_lower_atomic_virial = call_lower_with_atomic_virial - - def make_call_whether_do_atomic_virial(do_atomic_virial: bool) -> Callable: - if do_atomic_virial: - call_lower = call_lower_with_atomic_virial - else: - call_lower = call_lower_without_atomic_virial - - def call( - coord: tnp.ndarray, - atype: tnp.ndarray, - box: tnp.ndarray | None = None, - fparam: tnp.ndarray | None = None, - aparam: tnp.ndarray | None = None, - ) -> dict[str, tnp.ndarray]: - """Return model prediction. - - Parameters - ---------- - coord - The coordinates of the atoms. - shape: nf x (nloc x 3) - atype - The type of atoms. shape: nf x nloc - box - The simulation box. shape: nf x 9 - fparam - frame parameter. nf x ndf - aparam - atomic parameter. nf x nloc x nda - - Returns - ------- - ret_dict - The result dict of type dict[str,jnp.ndarray]. - The keys are defined by the `ModelOutputDef`. - - """ - return model_call_from_call_lower( - call_lower=call_lower, - rcut=model.get_rcut(), - sel=model.get_sel(), - mixed_types=model.mixed_types(), - model_output_def=model.model_output_def(), - coord=coord, - atype=atype, - box=box, - fparam=fparam, - aparam=aparam, - do_atomic_virial=do_atomic_virial, - ) - - return call - - @tf.function( - autograph=True, - input_signature=[ - tf.TensorSpec([None, None, 3], tf.float64), - tf.TensorSpec([None, None], tf.int32), - tf.TensorSpec([None, None, None], tf.float64), - tf.TensorSpec([None, model.get_dim_fparam()], tf.float64), - tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), - ], - ) - def call_with_atomic_virial( - coord: tnp.ndarray, - atype: tnp.ndarray, - box: tnp.ndarray, - fparam: tnp.ndarray, - aparam: tnp.ndarray, - ) -> dict[str, tnp.ndarray]: - return make_call_whether_do_atomic_virial(do_atomic_virial=True)( - coord, atype, box, fparam, aparam - ) - - tf_model.call_atomic_virial = call_with_atomic_virial - - @tf.function( - autograph=True, - input_signature=[ - tf.TensorSpec([None, None, 3], tf.float64), - tf.TensorSpec([None, None], tf.int32), - tf.TensorSpec([None, None, None], tf.float64), - tf.TensorSpec([None, model.get_dim_fparam()], tf.float64), - tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), - ], - ) - def call_without_atomic_virial( - coord: tnp.ndarray, - atype: tnp.ndarray, - box: tnp.ndarray, - fparam: tnp.ndarray, - aparam: tnp.ndarray, - ) -> dict[str, tnp.ndarray]: - return make_call_whether_do_atomic_virial(do_atomic_virial=False)( - coord, atype, box, fparam, aparam - ) - - tf_model.call = call_without_atomic_virial - - # set functions to export other attributes - @tf.function - def get_type_map() -> tf.Tensor: - return tf.constant(model.get_type_map(), dtype=tf.string) - - tf_model.get_type_map = get_type_map - - @tf.function - def get_rcut() -> tf.Tensor: - return tf.constant(model.get_rcut(), dtype=tf.double) - - tf_model.get_rcut = get_rcut - - @tf.function - def get_dim_fparam() -> tf.Tensor: - return tf.constant(model.get_dim_fparam(), dtype=tf.int64) - - tf_model.get_dim_fparam = get_dim_fparam - - @tf.function - def get_dim_aparam() -> tf.Tensor: - return tf.constant(model.get_dim_aparam(), dtype=tf.int64) - - tf_model.get_dim_aparam = get_dim_aparam - - @tf.function - def get_sel_type() -> tf.Tensor: - return tf.constant(model.get_sel_type(), dtype=tf.int64) - - tf_model.get_sel_type = get_sel_type - - @tf.function - def is_aparam_nall() -> tf.Tensor: - return tf.constant(model.is_aparam_nall(), dtype=tf.bool) - - tf_model.is_aparam_nall = is_aparam_nall - - @tf.function - def model_output_type() -> tf.Tensor: - return tf.constant(model.model_output_type(), dtype=tf.string) - - tf_model.model_output_type = model_output_type - - @tf.function - def mixed_types() -> tf.Tensor: - return tf.constant(model.mixed_types(), dtype=tf.bool) - - tf_model.mixed_types = mixed_types - - if model.get_min_nbor_dist() is not None: - - @tf.function - def get_min_nbor_dist() -> tf.Tensor: - return tf.constant(model.get_min_nbor_dist(), dtype=tf.double) - - tf_model.get_min_nbor_dist = get_min_nbor_dist - - @tf.function - def get_sel() -> tf.Tensor: - return tf.constant(model.get_sel(), dtype=tf.int64) - - tf_model.get_sel = get_sel - - @tf.function - def get_model_def_script() -> tf.Tensor: - return tf.constant( - json.dumps(model_def_script, separators=(",", ":")), dtype=tf.string - ) - - tf_model.get_model_def_script = get_model_def_script - - @tf.function - def has_message_passing() -> tf.Tensor: - return tf.constant(model.has_message_passing(), dtype=tf.bool) - - tf_model.has_message_passing = has_message_passing - - @tf.function - def has_default_fparam() -> tf.Tensor: - return tf.constant(model.has_default_fparam(), dtype=tf.bool) - - tf_model.has_default_fparam = has_default_fparam - - @tf.function - def get_default_fparam() -> tf.Tensor: - default_fparam = model.get_default_fparam() - if default_fparam is None: - return tf.constant([], dtype=tf.double) - else: - return tf.constant(default_fparam, dtype=tf.double) - - tf_model.get_default_fparam = get_default_fparam - tf.saved_model.save( - tf_model, - model_file, - options=tf.saved_model.SaveOptions(experimental_custom_gradients=True), - ) +__all__ = ["deserialize_to_file"] diff --git a/deepmd/jax/jax2tf/transform_output.py b/deepmd/jax/jax2tf/transform_output.py index f853744c02..e2548d07a7 100644 --- a/deepmd/jax/jax2tf/transform_output.py +++ b/deepmd/jax/jax2tf/transform_output.py @@ -1,113 +1,9 @@ # SPDX-License-Identifier: LGPL-3.0-or-later -import tensorflow as tf -import tensorflow.experimental.numpy as tnp +"""Compatibility exports for TensorFlow output transforms.""" -from deepmd.dpmodel.output_def import ( - ModelOutputDef, - OutputVariableDef, - get_deriv_name, - get_reduce_name, +from deepmd.tf2.transform_output import ( + communicate_extended_output, + get_leading_dims, ) - -def get_leading_dims( - vv: tnp.ndarray, - vdef: OutputVariableDef, -) -> tnp.ndarray: - """Get the dimensions of nf x nloc. - - Parameters - ---------- - vv : np.ndarray - The input array from which to compute the leading dimensions. - vdef : OutputVariableDef - The output variable definition containing the shape to exclude from `vv`. - - Returns - ------- - list - A list of leading dimensions of `vv`, excluding the last `len(vdef.shape)` dimensions. - """ - vshape = tf.shape(vv) - return vshape[: (len(vshape) - len(vdef.shape))] - - -def communicate_extended_output( - model_ret: dict[str, tnp.ndarray], - model_output_def: ModelOutputDef, - mapping: tnp.ndarray, # nf x nloc - do_atomic_virial: bool = False, -) -> dict[str, tnp.ndarray]: - """Transform the output of the model network defined on - local and ghost (extended) atoms to local atoms. - - """ - new_ret = {} - for kk in model_output_def.keys_outp(): - vv = model_ret[kk] - vdef = model_output_def[kk] - new_ret[kk] = vv - if vdef.reducible: - kk_redu = get_reduce_name(kk) - new_ret[kk_redu] = model_ret[kk_redu] - kk_derv_r, kk_derv_c = get_deriv_name(kk) - mldims = tf.shape(mapping) - vldims = get_leading_dims(vv, vdef) - if vdef.r_differentiable: - if model_ret[kk_derv_r] is not None: - derv_r_ext_dims = list(vdef.shape) + [3] # noqa:RUF005 - indices = mapping.reshape(tf.shape(mapping)[0], -1, 1) - # concat frame idx - indices = tf.concat( - [ - tf.repeat( - tf.range(tf.shape(indices)[0], dtype=indices.dtype), - tf.shape(mapping)[1], - ).reshape(tf.shape(indices)), - indices, - ], - axis=-1, - ) - force = tf.scatter_nd( - indices, - model_ret[kk_derv_r], - tf.cast(tf.concat([vldims, derv_r_ext_dims], axis=0), tf.int64), - ) - new_ret[kk_derv_r] = force.reshape( - tf.concat([tf.shape(force)[:2], list(vdef.shape), [3]], axis=0) - ) - else: - # name holders - new_ret[kk_derv_r] = None - if vdef.c_differentiable: - assert vdef.r_differentiable - if model_ret[kk_derv_c] is not None: - derv_c_ext_dims = list(vdef.shape) + [9] # noqa:RUF005 - indices = mapping.reshape(tf.shape(mapping)[0], -1, 1) - # concat frame idx - indices = tf.concat( - [ - tf.repeat( - tf.range(tf.shape(indices)[0], dtype=indices.dtype), - tf.shape(mapping)[1], - ).reshape(tf.shape(indices)), - indices, - ], - axis=-1, - ) - virial = tf.scatter_nd( - indices, - model_ret[kk_derv_c], - tf.cast(tf.concat([vldims, derv_c_ext_dims], axis=0), tf.int64), - ) - new_ret[kk_derv_c] = virial.reshape( - tf.concat([tf.shape(virial)[:2], list(vdef.shape), [9]], axis=0) - ) - new_ret[kk_derv_c + "_redu"] = tnp.sum(new_ret[kk_derv_c], axis=1) - else: - new_ret[kk_derv_c] = None - new_ret[kk_derv_c + "_redu"] = None - if not do_atomic_virial: - # pop atomic virial, because it is not correctly calculated. - new_ret.pop(kk_derv_c) - return new_ret +__all__ = ["communicate_extended_output", "get_leading_dims"] diff --git a/deepmd/jax/utils/serialization.py b/deepmd/jax/utils/serialization.py index 14386d9f3d..d324fe8701 100644 --- a/deepmd/jax/utils/serialization.py +++ b/deepmd/jax/utils/serialization.py @@ -138,8 +138,8 @@ def call_lower_with_fixed_do_atomic_virial( } save_dp_model(filename=model_file, model_dict=data) elif model_file.endswith(".savedmodel"): - from deepmd.jax.jax2tf.serialization import ( - deserialize_to_file as deserialize_to_savedmodel, + from deepmd.tf2.utils.serialization import ( + deserialize_to_savedmodel, ) return deserialize_to_savedmodel(model_file, data) diff --git a/deepmd/tf2/__init__.py b/deepmd/tf2/__init__.py new file mode 100644 index 0000000000..1b643dcc30 --- /dev/null +++ b/deepmd/tf2/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""TensorFlow eager backend based on dpmodel Array API code.""" diff --git a/deepmd/tf2/atomic_model/__init__.py b/deepmd/tf2/atomic_model/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/deepmd/tf2/atomic_model/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/deepmd/tf2/atomic_model/base_atomic_model.py b/deepmd/tf2/atomic_model/base_atomic_model.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/deepmd/tf2/atomic_model/base_atomic_model.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/deepmd/tf2/atomic_model/dipole_atomic_model.py b/deepmd/tf2/atomic_model/dipole_atomic_model.py new file mode 100644 index 0000000000..e77bcaa597 --- /dev/null +++ b/deepmd/tf2/atomic_model/dipole_atomic_model.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.atomic_model.dipole_atomic_model import ( + DPDipoleAtomicModel as DPAtomicModelDipoleDP, +) +from deepmd.tf2.atomic_model.dp_atomic_model import ( + make_tf2_dp_atomic_model_from_dpmodel, +) + + +class DPAtomicModelDipole(make_tf2_dp_atomic_model_from_dpmodel(DPAtomicModelDipoleDP)): + pass diff --git a/deepmd/tf2/atomic_model/dos_atomic_model.py b/deepmd/tf2/atomic_model/dos_atomic_model.py new file mode 100644 index 0000000000..206ce4cf25 --- /dev/null +++ b/deepmd/tf2/atomic_model/dos_atomic_model.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.atomic_model.dos_atomic_model import ( + DPDOSAtomicModel as DPAtomicModelDOSDP, +) +from deepmd.tf2.atomic_model.dp_atomic_model import ( + make_tf2_dp_atomic_model_from_dpmodel, +) + + +class DPAtomicModelDOS(make_tf2_dp_atomic_model_from_dpmodel(DPAtomicModelDOSDP)): + pass diff --git a/deepmd/tf2/atomic_model/dp_atomic_model.py b/deepmd/tf2/atomic_model/dp_atomic_model.py new file mode 100644 index 0000000000..604ea5b6de --- /dev/null +++ b/deepmd/tf2/atomic_model/dp_atomic_model.py @@ -0,0 +1,70 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import deepmd.tf2.descriptor as _tf2_descriptor # noqa: F401 +import deepmd.tf2.fitting.fitting as _tf2_fitting # noqa: F401 +import deepmd.tf2.utils.exclude_mask as _tf2_exclude_mask # noqa: F401 +from deepmd.dpmodel.atomic_model.dp_atomic_model import DPAtomicModel as DPAtomicModelDP +from deepmd.tf2.common import ( + tf2_module, +) +from deepmd.tf2.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.tf2.env import ( + stop_gradient, + xp, +) +from deepmd.tf2.fitting.base_fitting import ( + BaseFitting, +) + + +def make_tf2_dp_atomic_model_from_dpmodel( + dpmodel_atomic_model: type[DPAtomicModelDP], +) -> type[DPAtomicModelDP]: + """Make a tf2 backend DP atomic model from a DPModel backend DP atomic model. + + Parameters + ---------- + dpmodel_atomic_model : type[DPAtomicModelDP] + The DPModel backend DP atomic model. + + Returns + ------- + type[DPAtomicModel] + The tf2 backend DP atomic model. + """ + + @tf2_module + class tf2_atomic_model(dpmodel_atomic_model): + base_descriptor_cls = BaseDescriptor + """The base descriptor class.""" + base_fitting_cls = BaseFitting + """The base fitting class.""" + + def forward_common_atomic( + self, + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + mapping: xp.ndarray | None = None, + fparam: xp.ndarray | None = None, + aparam: xp.ndarray | None = None, + comm_dict: dict | None = None, + charge_spin: xp.ndarray | None = None, + ) -> dict[str, xp.ndarray]: + del comm_dict # tf2 path has no MPI ghost exchange + return super().forward_common_atomic( + extended_coord, + extended_atype, + stop_gradient(nlist), + mapping=mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + + return tf2_atomic_model + + +class DPAtomicModel(make_tf2_dp_atomic_model_from_dpmodel(DPAtomicModelDP)): + pass diff --git a/deepmd/tf2/atomic_model/energy_atomic_model.py b/deepmd/tf2/atomic_model/energy_atomic_model.py new file mode 100644 index 0000000000..23b87e8b04 --- /dev/null +++ b/deepmd/tf2/atomic_model/energy_atomic_model.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.atomic_model.energy_atomic_model import ( + DPEnergyAtomicModel as DPAtomicModelEnergyDP, +) +from deepmd.tf2.atomic_model.dp_atomic_model import ( + make_tf2_dp_atomic_model_from_dpmodel, +) + + +class DPAtomicModelEnergy(make_tf2_dp_atomic_model_from_dpmodel(DPAtomicModelEnergyDP)): + pass diff --git a/deepmd/tf2/atomic_model/linear_atomic_model.py b/deepmd/tf2/atomic_model/linear_atomic_model.py new file mode 100644 index 0000000000..f6030291f1 --- /dev/null +++ b/deepmd/tf2/atomic_model/linear_atomic_model.py @@ -0,0 +1,50 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import deepmd.tf2.atomic_model.dp_atomic_model as _tf2_dp_atomic_model # noqa: F401 +import deepmd.tf2.atomic_model.pairtab_atomic_model as _tf2_pairtab_model # noqa: F401 +import deepmd.tf2.utils.exclude_mask as _tf2_exclude_mask # noqa: F401 +from deepmd.dpmodel.atomic_model.linear_atomic_model import ( + DPZBLLinearEnergyAtomicModel as DPZBLLinearEnergyAtomicModelDP, +) +from deepmd.tf2.common import ( + tf2_module, +) +from deepmd.tf2.env import ( + stop_gradient, + xp, +) + + +@tf2_module +class DPZBLLinearEnergyAtomicModel(DPZBLLinearEnergyAtomicModelDP): + def __setattr__(self, name: str, value: Any) -> None: + if name == "zbl_weight": + # discard since it's only used in tests + # to fix TensorFlow tracing mutation error: Cannot mutate 'FlaxModule' from different trace level + return + return super().__setattr__(name, value) + + def forward_common_atomic( + self, + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + mapping: xp.ndarray | None = None, + fparam: xp.ndarray | None = None, + aparam: xp.ndarray | None = None, + comm_dict: dict | None = None, + charge_spin: xp.ndarray | None = None, + ) -> dict[str, xp.ndarray]: + del comm_dict # tf2 path has no MPI ghost exchange + return super().forward_common_atomic( + extended_coord, + extended_atype, + stop_gradient(nlist), + mapping=mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) diff --git a/deepmd/tf2/atomic_model/pairtab_atomic_model.py b/deepmd/tf2/atomic_model/pairtab_atomic_model.py new file mode 100644 index 0000000000..7a7e939e03 --- /dev/null +++ b/deepmd/tf2/atomic_model/pairtab_atomic_model.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import deepmd.tf2.utils.exclude_mask as _tf2_exclude_mask # noqa: F401 +from deepmd.dpmodel.atomic_model.pairtab_atomic_model import ( + PairTabAtomicModel as PairTabAtomicModelDP, +) +from deepmd.tf2.common import ( + tf2_module, +) +from deepmd.tf2.env import ( + stop_gradient, + xp, +) + + +@tf2_module +class PairTabAtomicModel(PairTabAtomicModelDP): + def forward_common_atomic( + self, + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + mapping: xp.ndarray | None = None, + fparam: xp.ndarray | None = None, + aparam: xp.ndarray | None = None, + comm_dict: dict | None = None, + charge_spin: xp.ndarray | None = None, + ) -> dict[str, xp.ndarray]: + del comm_dict # tf2 path has no MPI ghost exchange + return super().forward_common_atomic( + extended_coord, + extended_atype, + stop_gradient(nlist), + mapping=mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) diff --git a/deepmd/tf2/atomic_model/polar_atomic_model.py b/deepmd/tf2/atomic_model/polar_atomic_model.py new file mode 100644 index 0000000000..48f494d2d4 --- /dev/null +++ b/deepmd/tf2/atomic_model/polar_atomic_model.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.atomic_model.polar_atomic_model import ( + DPPolarAtomicModel as DPAtomicModelPolarDP, +) +from deepmd.tf2.atomic_model.dp_atomic_model import ( + make_tf2_dp_atomic_model_from_dpmodel, +) + + +class DPAtomicModelPolar(make_tf2_dp_atomic_model_from_dpmodel(DPAtomicModelPolarDP)): + pass diff --git a/deepmd/tf2/atomic_model/property_atomic_model.py b/deepmd/tf2/atomic_model/property_atomic_model.py new file mode 100644 index 0000000000..aef020c2c5 --- /dev/null +++ b/deepmd/tf2/atomic_model/property_atomic_model.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.atomic_model.property_atomic_model import ( + DPPropertyAtomicModel as DPAtomicModelPropertyDP, +) +from deepmd.tf2.atomic_model.dp_atomic_model import ( + make_tf2_dp_atomic_model_from_dpmodel, +) + + +class DPAtomicModelProperty( + make_tf2_dp_atomic_model_from_dpmodel(DPAtomicModelPropertyDP) +): + pass diff --git a/deepmd/tf2/common.py b/deepmd/tf2/common.py new file mode 100644 index 0000000000..bb8155a38c --- /dev/null +++ b/deepmd/tf2/common.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from collections.abc import ( + Callable, +) +from functools import ( + wraps, +) +from importlib import ( + import_module, +) +from typing import ( + Any, + TypeVar, +) + +import numpy as np +import tensorflow as tf + +from deepmd._vendors import ndtensorflow as xp +from deepmd.dpmodel.common import ( + NativeOP, +) + + +def to_tensorflow_array(array: Any | None) -> Any: + """Convert an object to an ndtensorflow Array. + + Parameters + ---------- + array + The object to convert. + + Returns + ------- + ndtensorflow.Array + The TensorFlow-backed array. + """ + if array is None: + return None + if isinstance(array, np.ndarray): + return xp.asarray(tf.convert_to_tensor(array)) + return xp.asarray(array) + + +def to_tf_tensor(array: Any | None) -> tf.Tensor | None: + """Unwrap a TensorFlow-backed Array to a TensorFlow tensor.""" + if array is None: + return None + if isinstance(array, xp.Array): + return array.unwrap() + if isinstance(array, tf.Tensor): + return array + return tf.convert_to_tensor(array) + + +def wrap_tensor(tensor: Any | None) -> Any | None: + """Wrap a TensorFlow tensor as an ndtensorflow Array.""" + if tensor is None: + return None + return xp.asarray(tensor) + + +def wrap_value(value: Any) -> Any: + """Recursively wrap TensorFlow tensors as ndtensorflow Arrays.""" + if isinstance(value, dict): + return {kk: wrap_value(vv) for kk, vv in value.items()} + if isinstance(value, tuple): + return tuple(wrap_value(vv) for vv in value) + if isinstance(value, list): + return [wrap_value(vv) for vv in value] + return wrap_tensor(value) + + +def unwrap_value(value: Any) -> Any: + """Recursively unwrap ndtensorflow Arrays for TensorFlow SavedModel returns.""" + if isinstance(value, xp.Array): + return value.unwrap() + if isinstance(value, dict): + return {kk: unwrap_value(vv) for kk, vv in value.items()} + if isinstance(value, tuple): + return tuple(unwrap_value(vv) for vv in value) + if isinstance(value, list): + return [unwrap_value(vv) for vv in value] + return value + + +_PACKAGE_ROOT = __name__.rsplit(".", 1)[0] +_DPMODEL_TO_TF2: dict[type[Any], Callable[[Any], Any]] = {} +_AUTO_WRAPPED_CLASSES: dict[type[NativeOP], type[Any]] = {} +_REGISTRATIONS_READY = False +_REGISTRATIONS_IN_PROGRESS = False +_REGISTRATION_MODULES = ( + f"{_PACKAGE_ROOT}.utils.network", + f"{_PACKAGE_ROOT}.utils.exclude_mask", + f"{_PACKAGE_ROOT}.utils.type_embed", + f"{_PACKAGE_ROOT}.descriptor.dpa1", + f"{_PACKAGE_ROOT}.descriptor.se_atten_v2", + f"{_PACKAGE_ROOT}.descriptor.se_e2_a", + f"{_PACKAGE_ROOT}.descriptor.se_e2_r", + f"{_PACKAGE_ROOT}.descriptor.se_t", + f"{_PACKAGE_ROOT}.descriptor.se_t_tebd", + f"{_PACKAGE_ROOT}.descriptor.repformers", + f"{_PACKAGE_ROOT}.descriptor.dpa2", + f"{_PACKAGE_ROOT}.descriptor.repflows", + f"{_PACKAGE_ROOT}.descriptor.dpa3", + f"{_PACKAGE_ROOT}.descriptor.hybrid", + f"{_PACKAGE_ROOT}.fitting", + f"{_PACKAGE_ROOT}.atomic_model.dp_atomic_model", + f"{_PACKAGE_ROOT}.atomic_model.energy_atomic_model", + f"{_PACKAGE_ROOT}.atomic_model.dipole_atomic_model", + f"{_PACKAGE_ROOT}.atomic_model.dos_atomic_model", + f"{_PACKAGE_ROOT}.atomic_model.polar_atomic_model", + f"{_PACKAGE_ROOT}.atomic_model.property_atomic_model", + f"{_PACKAGE_ROOT}.atomic_model.pairtab_atomic_model", + f"{_PACKAGE_ROOT}.atomic_model.linear_atomic_model", + f"{_PACKAGE_ROOT}.model", +) + + +class TF2List(list): + def append(self, item: Any) -> None: + return super().append(convert_tf2_value(item)) + + def extend(self, items: list[Any]) -> None: + return super().extend(convert_tf2_value(item) for item in items) + + def insert(self, index: int, item: Any) -> None: + return super().insert(index, convert_tf2_value(item)) + + def __setitem__(self, index: Any, item: Any) -> None: + if isinstance(index, slice): + item = [convert_tf2_value(ii) for ii in item] + else: + item = convert_tf2_value(item) + return super().__setitem__(index, item) + + +def register_dpmodel_mapping( + dpmodel_cls: type[Any], converter: Callable[[Any], Any] +) -> None: + """Register how to convert a dpmodel object to its tf2 wrapper.""" + _DPMODEL_TO_TF2[dpmodel_cls] = converter + + +def _looks_like_dpmodel_class(cls: type[Any]) -> bool: + module = cls.__module__ + return module == "deepmd.dpmodel" or module.startswith("deepmd.dpmodel.") + + +def _looks_like_dpmodel_object(value: Any) -> bool: + return _looks_like_dpmodel_class(type(value)) + + +def _looks_like_tf2_object(value: Any) -> bool: + module = type(value).__module__ + return module == _PACKAGE_ROOT or module.startswith(f"{_PACKAGE_ROOT}.") + + +def _ensure_registrations() -> None: + global _REGISTRATIONS_IN_PROGRESS, _REGISTRATIONS_READY + + if _REGISTRATIONS_READY or _REGISTRATIONS_IN_PROGRESS: + return + + _REGISTRATIONS_IN_PROGRESS = True + try: + for module in _REGISTRATION_MODULES: + import_module(module) + _REGISTRATIONS_READY = True + finally: + _REGISTRATIONS_IN_PROGRESS = False + + +def try_convert_module(value: Any) -> Any | None: + """Convert a registered dpmodel object to its tf2 wrapper.""" + if _looks_like_tf2_object(value): + return None + converter = _DPMODEL_TO_TF2.get(type(value)) + if converter is not None: + return converter(value) + if _looks_like_dpmodel_object(value): + _ensure_registrations() + converter = _DPMODEL_TO_TF2.get(type(value)) + if converter is not None: + return converter(value) + if isinstance(value, NativeOP): + return _auto_wrap_native_op(value) + return None + + +def _auto_wrap_native_op(value: NativeOP) -> Any: + cls = type(value) + if cls not in _AUTO_WRAPPED_CLASSES: + wrapped_cls = type( + cls.__name__, + (cls,), + { + "__module__": __name__, + "__qualname__": cls.__qualname__, + }, + ) + _AUTO_WRAPPED_CLASSES[cls] = tf2_module(wrapped_cls) + wrapped_cls = _AUTO_WRAPPED_CLASSES[cls] + if not (hasattr(value, "serialize") and hasattr(wrapped_cls, "deserialize")): + raise TypeError( + f"Cannot auto-wrap {cls.__name__}: " + "it must implement serialize()/deserialize() or be explicitly " + "registered via register_dpmodel_mapping()." + ) + return wrapped_cls.deserialize(value.serialize()) + + +def _try_convert_list(value: list[Any], *, keep_converting: bool = False) -> list[Any]: + converted = TF2List() if keep_converting else [] + changed = keep_converting + for item in value: + converted_item = convert_tf2_value(item) + converted.append(converted_item) + changed = changed or converted_item is not item + return converted if changed else value + + +def convert_tf2_value(value: Any) -> Any: + if isinstance(value, np.ndarray): + return to_tensorflow_array(value) + + if isinstance(value, list): + return _try_convert_list(value) + + converted = try_convert_module(value) + if converted is not None: + return converted + + return value + + +def tf2_setattr(obj: Any, name: str, value: Any) -> Any: + if name in getattr(obj, "_tf2_skip_auto_convert_attrs", ()): + return value + + if isinstance(value, list) and name in getattr(obj, "_tf2_data_list_attrs", ()): + return _try_convert_list(value, keep_converting=True) + + return convert_tf2_value(value) + + +T = TypeVar("T") + + +def tf2_module(module: type[T]) -> type[T]: + """Wrap a dpmodel subclass as a TensorFlow ``tf.Module``.""" + + @wraps(module, updated=()) + class TF2Module(module, tf.Module): # type: ignore[misc, valid-type] + def __init__(self, *args: Any, **kwargs: Any) -> None: + tf.Module.__init__(self) + super().__init__(*args, **kwargs) + for name in list(self.__dict__): + value = self.__dict__[name] + if isinstance(value, list): + converted = _try_convert_list( + value, + keep_converting=name + in getattr(self, "_tf2_data_list_attrs", ()), + ) + if converted is not value: + setattr(self, name, converted) + + def __setattr__(self, name: str, value: Any) -> None: + value = tf2_setattr(self, name, value) + return super().__setattr__(name, value) + + if hasattr(TF2Module, "deserialize"): + for base in module.__bases__: + if base in (object, NativeOP): + continue + if ( + _looks_like_dpmodel_class(base) + and hasattr(base, "serialize") + and base not in _DPMODEL_TO_TF2 + ): + + def _converter(v: Any, _cls: type[Any] = TF2Module) -> Any: + return _cls.deserialize(v.serialize()) + + _DPMODEL_TO_TF2[base] = _converter + + return TF2Module diff --git a/deepmd/tf2/descriptor/__init__.py b/deepmd/tf2/descriptor/__init__.py new file mode 100644 index 0000000000..1bbefbea6f --- /dev/null +++ b/deepmd/tf2/descriptor/__init__.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .dpa1 import ( + DescrptDPA1, +) +from .dpa2 import ( + DescrptDPA2, +) +from .dpa3 import ( + DescrptDPA3, +) +from .hybrid import ( + DescrptHybrid, +) +from .se_atten_v2 import ( + DescrptSeAttenV2, +) +from .se_e2_a import ( + DescrptSeA, +) +from .se_e2_r import ( + DescrptSeR, +) +from .se_t import ( + DescrptSeT, +) +from .se_t_tebd import ( + DescrptSeTTebd, +) + +__all__ = [ + "DescrptDPA1", + "DescrptDPA2", + "DescrptDPA3", + "DescrptHybrid", + "DescrptSeA", + "DescrptSeAttenV2", + "DescrptSeR", + "DescrptSeT", + "DescrptSeTTebd", +] diff --git a/deepmd/tf2/descriptor/base_descriptor.py b/deepmd/tf2/descriptor/base_descriptor.py new file mode 100644 index 0000000000..2a31895f55 --- /dev/null +++ b/deepmd/tf2/descriptor/base_descriptor.py @@ -0,0 +1,11 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +from deepmd.dpmodel.descriptor.make_base_descriptor import ( + make_base_descriptor, +) + +# no type annotations standard in array api +BaseDescriptor = make_base_descriptor(Any) diff --git a/deepmd/tf2/descriptor/dpa1.py b/deepmd/tf2/descriptor/dpa1.py new file mode 100644 index 0000000000..8bb17f66d6 --- /dev/null +++ b/deepmd/tf2/descriptor/dpa1.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.dpa1 import DescrptBlockSeAtten as DescrptBlockSeAttenDP +from deepmd.dpmodel.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1DP +from deepmd.dpmodel.descriptor.dpa1 import GatedAttentionLayer as GatedAttentionLayerDP +from deepmd.dpmodel.descriptor.dpa1 import ( + NeighborGatedAttention as NeighborGatedAttentionDP, +) +from deepmd.dpmodel.descriptor.dpa1 import ( + NeighborGatedAttentionLayer as NeighborGatedAttentionLayerDP, +) + +from ..common import ( + tf2_module, +) +from ..utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from ..utils import network as _tf2_network # noqa: F401 +from ..utils import type_embed as _tf2_type_embed # noqa: F401 +from .base_descriptor import ( + BaseDescriptor, +) + + +@tf2_module +class GatedAttentionLayer(GatedAttentionLayerDP): + pass + + +@tf2_module +class NeighborGatedAttentionLayer(NeighborGatedAttentionLayerDP): + pass + + +@tf2_module +class NeighborGatedAttention(NeighborGatedAttentionDP): + pass + + +@tf2_module +class DescrptBlockSeAtten(DescrptBlockSeAttenDP): + pass + + +@BaseDescriptor.register("dpa1") +@BaseDescriptor.register("se_atten") +@tf2_module +class DescrptDPA1(DescrptDPA1DP): + pass diff --git a/deepmd/tf2/descriptor/dpa2.py b/deepmd/tf2/descriptor/dpa2.py new file mode 100644 index 0000000000..3b9f3c1709 --- /dev/null +++ b/deepmd/tf2/descriptor/dpa2.py @@ -0,0 +1,20 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.dpa2 import DescrptDPA2 as DescrptDPA2DP + +from ..common import ( + tf2_module, +) +from ..utils import network as _tf2_network # noqa: F401 +from ..utils import type_embed as _tf2_type_embed # noqa: F401 +from . import dpa1 as _tf2_dpa1 # noqa: F401 +from . import repformers as _tf2_repformers # noqa: F401 +from . import se_t_tebd as _tf2_se_t_tebd # noqa: F401 +from .base_descriptor import ( + BaseDescriptor, +) + + +@BaseDescriptor.register("dpa2") +@tf2_module +class DescrptDPA2(DescrptDPA2DP): + pass diff --git a/deepmd/tf2/descriptor/dpa3.py b/deepmd/tf2/descriptor/dpa3.py new file mode 100644 index 0000000000..68e68bdd76 --- /dev/null +++ b/deepmd/tf2/descriptor/dpa3.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.dpa3 import DescrptDPA3 as DescrptDPA3DP + +from ..common import ( + tf2_module, +) +from ..utils import network as _tf2_network # noqa: F401 +from ..utils import type_embed as _tf2_type_embed # noqa: F401 +from . import repflows as _tf2_repflows # noqa: F401 +from .base_descriptor import ( + BaseDescriptor, +) + + +@BaseDescriptor.register("dpa3") +@tf2_module +class DescrptDPA3(DescrptDPA3DP): + pass diff --git a/deepmd/tf2/descriptor/hybrid.py b/deepmd/tf2/descriptor/hybrid.py new file mode 100644 index 0000000000..6c9f63d228 --- /dev/null +++ b/deepmd/tf2/descriptor/hybrid.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.hybrid import DescrptHybrid as DescrptHybridDP + +from ..common import ( + tf2_module, +) +from .base_descriptor import ( + BaseDescriptor, +) + + +@BaseDescriptor.register("hybrid") +@tf2_module +class DescrptHybrid(DescrptHybridDP): + pass diff --git a/deepmd/tf2/descriptor/repflows.py b/deepmd/tf2/descriptor/repflows.py new file mode 100644 index 0000000000..c3a9244691 --- /dev/null +++ b/deepmd/tf2/descriptor/repflows.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + ClassVar, +) + +from deepmd.dpmodel.descriptor.repflows import ( + DescrptBlockRepflows as DescrptBlockRepflowsDP, +) +from deepmd.dpmodel.descriptor.repflows import RepFlowLayer as RepFlowLayerDP + +from ..common import ( + tf2_module, +) +from ..utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from ..utils import network as _tf2_network # noqa: F401 + + +@tf2_module +class DescrptBlockRepflows(DescrptBlockRepflowsDP): + pass + + +@tf2_module +class RepFlowLayer(RepFlowLayerDP): + _tf2_data_list_attrs: ClassVar[set[str]] = { + "n_residual", + "e_residual", + "a_residual", + } diff --git a/deepmd/tf2/descriptor/repformers.py b/deepmd/tf2/descriptor/repformers.py new file mode 100644 index 0000000000..d35d889756 --- /dev/null +++ b/deepmd/tf2/descriptor/repformers.py @@ -0,0 +1,57 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + ClassVar, +) + +from deepmd.dpmodel.descriptor.repformers import ( + Atten2EquiVarApply as Atten2EquiVarApplyDP, +) +from deepmd.dpmodel.descriptor.repformers import Atten2Map as Atten2MapDP +from deepmd.dpmodel.descriptor.repformers import ( + Atten2MultiHeadApply as Atten2MultiHeadApplyDP, +) +from deepmd.dpmodel.descriptor.repformers import ( + DescrptBlockRepformers as DescrptBlockRepformersDP, +) +from deepmd.dpmodel.descriptor.repformers import LocalAtten as LocalAttenDP +from deepmd.dpmodel.descriptor.repformers import RepformerLayer as RepformerLayerDP + +from ..common import ( + tf2_module, +) +from ..utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from ..utils import network as _tf2_network # noqa: F401 + + +@tf2_module +class DescrptBlockRepformers(DescrptBlockRepformersDP): + pass + + +@tf2_module +class Atten2Map(Atten2MapDP): + pass + + +@tf2_module +class Atten2MultiHeadApply(Atten2MultiHeadApplyDP): + pass + + +@tf2_module +class Atten2EquiVarApply(Atten2EquiVarApplyDP): + pass + + +@tf2_module +class LocalAtten(LocalAttenDP): + pass + + +@tf2_module +class RepformerLayer(RepformerLayerDP): + _tf2_data_list_attrs: ClassVar[set[str]] = { + "g1_residual", + "g2_residual", + "h2_residual", + } diff --git a/deepmd/tf2/descriptor/se_atten_v2.py b/deepmd/tf2/descriptor/se_atten_v2.py new file mode 100644 index 0000000000..84db19fc59 --- /dev/null +++ b/deepmd/tf2/descriptor/se_atten_v2.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.se_atten_v2 import DescrptSeAttenV2 as DescrptSeAttenV2DP + +from ..common import ( + register_dpmodel_mapping, +) +from .dpa1 import ( + DescrptDPA1, +) + + +class DescrptSeAttenV2(DescrptDPA1, DescrptSeAttenV2DP): + pass + + +register_dpmodel_mapping( + DescrptSeAttenV2DP, + lambda v: DescrptSeAttenV2.deserialize(v.serialize()), +) diff --git a/deepmd/tf2/descriptor/se_e2_a.py b/deepmd/tf2/descriptor/se_e2_a.py new file mode 100644 index 0000000000..c36614112e --- /dev/null +++ b/deepmd/tf2/descriptor/se_e2_a.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.se_e2_a import DescrptSeAArrayAPI as DescrptSeADP + +from ..common import ( + tf2_module, +) +from ..utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from ..utils import network as _tf2_network # noqa: F401 +from .base_descriptor import ( + BaseDescriptor, +) + + +@BaseDescriptor.register("se_e2_a") +@BaseDescriptor.register("se_a") +@tf2_module +class DescrptSeA(DescrptSeADP): + pass diff --git a/deepmd/tf2/descriptor/se_e2_r.py b/deepmd/tf2/descriptor/se_e2_r.py new file mode 100644 index 0000000000..884889f1c8 --- /dev/null +++ b/deepmd/tf2/descriptor/se_e2_r.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.se_r import DescrptSeR as DescrptSeRDP + +from ..common import ( + tf2_module, +) +from ..utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from ..utils import network as _tf2_network # noqa: F401 +from .base_descriptor import ( + BaseDescriptor, +) + + +@BaseDescriptor.register("se_e2_r") +@BaseDescriptor.register("se_r") +@tf2_module +class DescrptSeR(DescrptSeRDP): + pass diff --git a/deepmd/tf2/descriptor/se_t.py b/deepmd/tf2/descriptor/se_t.py new file mode 100644 index 0000000000..98e142de6c --- /dev/null +++ b/deepmd/tf2/descriptor/se_t.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.se_t import DescrptSeT as DescrptSeTDP + +from ..common import ( + tf2_module, +) +from ..utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from ..utils import network as _tf2_network # noqa: F401 + + +@tf2_module +class DescrptSeT(DescrptSeTDP): + pass diff --git a/deepmd/tf2/descriptor/se_t_tebd.py b/deepmd/tf2/descriptor/se_t_tebd.py new file mode 100644 index 0000000000..4b7e65c50f --- /dev/null +++ b/deepmd/tf2/descriptor/se_t_tebd.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.descriptor.se_t_tebd import ( + DescrptBlockSeTTebd as DescrptBlockSeTTebdDP, +) +from deepmd.dpmodel.descriptor.se_t_tebd import DescrptSeTTebd as DescrptSeTTebdDP + +from ..common import ( + tf2_module, +) +from ..utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from ..utils import network as _tf2_network # noqa: F401 +from ..utils import type_embed as _tf2_type_embed # noqa: F401 + + +@tf2_module +class DescrptBlockSeTTebd(DescrptBlockSeTTebdDP): + pass + + +@tf2_module +class DescrptSeTTebd(DescrptSeTTebdDP): + pass diff --git a/deepmd/tf2/env.py b/deepmd/tf2/env.py new file mode 100644 index 0000000000..a032001348 --- /dev/null +++ b/deepmd/tf2/env.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""TensorFlow eager backend environment.""" + +from typing import ( + Any, +) + +import tensorflow as tf + +from deepmd._vendors import ndtensorflow as xp + +if not tf.executing_eagerly(): + raise RuntimeError( + "The tf2 backend requires TensorFlow eager execution. " + "It cannot be imported after eager execution has been disabled." + ) + +Array = xp.Array +xp.ndarray = xp.Array + + +def stop_gradient(value: Any) -> Any: + """Stop gradients on TensorFlow-backed Array objects.""" + if isinstance(value, xp.Array): + return xp.asarray(tf.stop_gradient(value.unwrap())) + return xp.asarray(tf.stop_gradient(value)) + + +__all__ = ["Array", "stop_gradient", "tf", "xp"] diff --git a/deepmd/tf2/fitting/__init__.py b/deepmd/tf2/fitting/__init__.py new file mode 100644 index 0000000000..2041f600ea --- /dev/null +++ b/deepmd/tf2/fitting/__init__.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from .fitting import ( + DipoleFittingNet, + DOSFittingNet, + EnergyFittingNet, + PolarFittingNet, + PropertyFittingNet, +) + +__all__ = [ + "DOSFittingNet", + "DipoleFittingNet", + "EnergyFittingNet", + "PolarFittingNet", + "PropertyFittingNet", +] diff --git a/deepmd/tf2/fitting/base_fitting.py b/deepmd/tf2/fitting/base_fitting.py new file mode 100644 index 0000000000..18fb3c8e79 --- /dev/null +++ b/deepmd/tf2/fitting/base_fitting.py @@ -0,0 +1,9 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.fitting.make_base_fitting import ( + make_base_fitting, +) +from deepmd.tf2.env import ( + xp, +) + +BaseFitting = make_base_fitting(xp.ndarray) diff --git a/deepmd/tf2/fitting/fitting.py b/deepmd/tf2/fitting/fitting.py new file mode 100644 index 0000000000..7cb2a89a28 --- /dev/null +++ b/deepmd/tf2/fitting/fitting.py @@ -0,0 +1,49 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.fitting.dipole_fitting import DipoleFitting as DipoleFittingNetDP +from deepmd.dpmodel.fitting.dos_fitting import DOSFittingNet as DOSFittingNetDP +from deepmd.dpmodel.fitting.ener_fitting import EnergyFittingNet as EnergyFittingNetDP +from deepmd.dpmodel.fitting.polarizability_fitting import ( + PolarFitting as PolarFittingNetDP, +) +from deepmd.dpmodel.fitting.property_fitting import ( + PropertyFittingNet as PropertyFittingNetDP, +) + +from ..common import ( + tf2_module, +) +from ..utils import exclude_mask as _tf2_exclude_mask # noqa: F401 +from ..utils import network as _tf2_network # noqa: F401 +from .base_fitting import ( + BaseFitting, +) + + +@BaseFitting.register("ener") +@tf2_module +class EnergyFittingNet(EnergyFittingNetDP): + pass + + +@BaseFitting.register("property") +@tf2_module +class PropertyFittingNet(PropertyFittingNetDP): + pass + + +@BaseFitting.register("dos") +@tf2_module +class DOSFittingNet(DOSFittingNetDP): + pass + + +@BaseFitting.register("dipole") +@tf2_module +class DipoleFittingNet(DipoleFittingNetDP): + pass + + +@BaseFitting.register("polar") +@tf2_module +class PolarFittingNet(PolarFittingNetDP): + pass diff --git a/deepmd/tf2/infer/__init__.py b/deepmd/tf2/infer/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/deepmd/tf2/infer/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/deepmd/tf2/infer/deep_eval.py b/deepmd/tf2/infer/deep_eval.py new file mode 100644 index 0000000000..4d6a8c5e2f --- /dev/null +++ b/deepmd/tf2/infer/deep_eval.py @@ -0,0 +1,386 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +from collections.abc import ( + Callable, +) +from typing import ( + TYPE_CHECKING, + Any, + Optional, +) + +import numpy as np +import tensorflow as tf + +from deepmd.dpmodel.output_def import ( + ModelOutputDef, + OutputVariableCategory, + OutputVariableDef, +) +from deepmd.dpmodel.utils.batch_size import ( + AutoBatchSize, +) +from deepmd.env import ( + GLOBAL_NP_FLOAT_PRECISION, +) +from deepmd.infer.deep_dipole import ( + DeepDipole, +) +from deepmd.infer.deep_dos import ( + DeepDOS, +) +from deepmd.infer.deep_eval import DeepEval as DeepEvalWrapper +from deepmd.infer.deep_eval import ( + DeepEvalBackend, +) +from deepmd.infer.deep_polar import ( + DeepPolar, +) +from deepmd.infer.deep_pot import ( + DeepPot, +) +from deepmd.infer.deep_wfc import ( + DeepWFC, +) + +if TYPE_CHECKING: + import ase.neighborlist + + +def _decode_list_of_bytes(list_of_bytes: list[bytes]) -> list[str]: + return [item.decode() for item in list_of_bytes] + + +def _to_numpy_dict(ret: dict[str, Any]) -> dict[str, np.ndarray]: + return { + key: value.numpy() if isinstance(value, tf.Tensor) else value + for key, value in ret.items() + } + + +class TF2SavedModelWrapper(tf.Module): + """Small Python wrapper around the exported TensorFlow SavedModel.""" + + def __init__(self, model: str) -> None: + super().__init__() + self.model = tf.saved_model.load(model) + self.type_map = _decode_list_of_bytes( + self.model.get_type_map().numpy().tolist() + ) + self.rcut = self.model.get_rcut().numpy().item() + self.dim_fparam = self.model.get_dim_fparam().numpy().item() + self.dim_aparam = self.model.get_dim_aparam().numpy().item() + self.sel_type = self.model.get_sel_type().numpy().tolist() + self._is_aparam_nall = self.model.is_aparam_nall().numpy().item() + self._model_output_type = _decode_list_of_bytes( + self.model.model_output_type().numpy().tolist() + ) + self._mixed_types = self.model.mixed_types().numpy().item() + self.min_nbor_dist = ( + self.model.get_min_nbor_dist().numpy().item() + if hasattr(self.model, "get_min_nbor_dist") + else None + ) + self.sel = self.model.get_sel().numpy().tolist() + self.model_def_script = self.model.get_model_def_script().numpy().decode() + self._has_default_fparam = ( + self.model.has_default_fparam().numpy().item() + if hasattr(self.model, "has_default_fparam") + else False + ) + self.default_fparam = ( + self.model.get_default_fparam().numpy().tolist() + if hasattr(self.model, "get_default_fparam") + else None + ) + + def __call__( + self, + coord: np.ndarray, + atype: np.ndarray, + box: np.ndarray | None = None, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + do_atomic_virial: bool = False, + ) -> dict[str, np.ndarray]: + call = self.model.call_atomic_virial if do_atomic_virial else self.model.call + coord = tf.convert_to_tensor(coord, dtype=tf.float64) + atype = tf.convert_to_tensor(atype, dtype=tf.int32) + if box is None: + box = np.empty((coord.shape[0], 0, 0), dtype=np.float64) + if fparam is None: + fparam = np.empty((coord.shape[0], self.get_dim_fparam()), dtype=np.float64) + if aparam is None: + aparam = np.empty( + (coord.shape[0], coord.shape[1], self.get_dim_aparam()), + dtype=np.float64, + ) + ret = call( + coord, + atype, + tf.convert_to_tensor(box, dtype=tf.float64), + tf.convert_to_tensor(fparam, dtype=tf.float64), + tf.convert_to_tensor(aparam, dtype=tf.float64), + ) + return _to_numpy_dict(ret) + + def get_type_map(self) -> list[str]: + return self.type_map + + def get_rcut(self) -> float: + return self.rcut + + def get_dim_fparam(self) -> int: + return self.dim_fparam + + def get_dim_aparam(self) -> int: + return self.dim_aparam + + def get_sel_type(self) -> list[int]: + return self.sel_type + + def is_aparam_nall(self) -> bool: + return self._is_aparam_nall + + def model_output_type(self) -> list[str]: + return self._model_output_type + + def mixed_types(self) -> bool: + return self._mixed_types + + def get_min_nbor_dist(self) -> float | None: + return self.min_nbor_dist + + def get_sel(self) -> list[int]: + return self.sel + + def get_model_def_script(self) -> str: + return self.model_def_script + + def has_default_fparam(self) -> bool: + return self._has_default_fparam + + def get_default_fparam(self) -> list[float] | None: + return self.default_fparam + + +class DeepEval(DeepEvalBackend): + """TensorFlow 2 SavedModel backend implementation of DeepEval.""" + + def __init__( + self, + model_file: str, + output_def: ModelOutputDef, + *args: Any, + auto_batch_size: bool | int | AutoBatchSize = True, + neighbor_list: Optional["ase.neighborlist.NewPrimitiveNeighborList"] = None, + **kwargs: Any, + ) -> None: + if not model_file.endswith(".savedmodeltf"): + raise ValueError("TF2 backend only supports .savedmodeltf files") + self.output_def = output_def + self.model_path = model_file + self.dp = TF2SavedModelWrapper(model_file) + self.rcut = self.dp.get_rcut() + self.type_map = self.dp.get_type_map() + if isinstance(auto_batch_size, bool): + self.auto_batch_size = AutoBatchSize() if auto_batch_size else None + elif isinstance(auto_batch_size, int): + self.auto_batch_size = AutoBatchSize(auto_batch_size) + elif isinstance(auto_batch_size, AutoBatchSize): + self.auto_batch_size = auto_batch_size + else: + raise TypeError("auto_batch_size should be bool, int, or AutoBatchSize") + + def get_rcut(self) -> float: + return self.rcut + + def get_ntypes(self) -> int: + return len(self.type_map) + + def get_type_map(self) -> list[str]: + return self.type_map + + def get_dim_fparam(self) -> int: + return self.dp.get_dim_fparam() + + def get_dim_aparam(self) -> int: + return self.dp.get_dim_aparam() + + def has_default_fparam(self) -> bool: + return self.dp.has_default_fparam() + + @property + def model_type(self) -> type["DeepEvalWrapper"]: + model_output_type = self.dp.model_output_type() + if "energy" in model_output_type: + return DeepPot + if "dos" in model_output_type: + return DeepDOS + if "dipole" in model_output_type: + return DeepDipole + if "polar" in model_output_type or "polarizability" in model_output_type: + return DeepPolar + if "wfc" in model_output_type: + return DeepWFC + raise RuntimeError("Unknown model type") + + def get_sel_type(self) -> list[int]: + return self.dp.get_sel_type() + + def get_numb_dos(self) -> int: + return 0 + + def get_has_efield(self) -> bool: + return False + + def get_ntypes_spin(self) -> int: + return 0 + + def eval( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + atomic: bool = False, + fparam: np.ndarray | None = None, + aparam: np.ndarray | None = None, + **kwargs: Any, + ) -> dict[str, np.ndarray]: + atom_types = np.array(atom_types, dtype=np.int32) + coords = np.array(coords) + if cells is not None: + cells = np.array(cells) + natoms, numb_test = self._get_natoms_and_nframes( + coords, atom_types, len(atom_types.shape) > 1 + ) + request_defs = self._get_request_defs(atomic) + out = self._eval_func(self._eval_model, numb_test, natoms)( + coords, cells, atom_types, fparam, aparam, request_defs + ) + return dict(zip([x.name for x in request_defs], out, strict=True)) + + def _get_request_defs(self, atomic: bool) -> list[OutputVariableDef]: + if atomic: + return list(self.output_def.var_defs.values()) + return [ + x + for x in self.output_def.var_defs.values() + if x.category + in ( + OutputVariableCategory.REDU, + OutputVariableCategory.DERV_R, + OutputVariableCategory.DERV_C_REDU, + ) + ] + + def _eval_func(self, inner_func: Callable, numb_test: int, natoms: int) -> Callable: + if self.auto_batch_size is not None: + + def eval_func(*args: Any, **kwargs: Any) -> Any: + return self.auto_batch_size.execute_all( + inner_func, numb_test, natoms, *args, **kwargs + ) + + else: + eval_func = inner_func + return eval_func + + def _get_natoms_and_nframes( + self, + coords: np.ndarray, + atom_types: np.ndarray, + mixed_type: bool = False, + ) -> tuple[int, int]: + if mixed_type: + natoms = len(atom_types[0]) + else: + natoms = len(atom_types) + if natoms == 0: + assert coords.size == 0 + else: + coords = np.reshape(np.array(coords), [-1, natoms * 3]) + return natoms, coords.shape[0] + + def _eval_model( + self, + coords: np.ndarray, + cells: np.ndarray | None, + atom_types: np.ndarray, + fparam: np.ndarray | None, + aparam: np.ndarray | None, + request_defs: list[OutputVariableDef], + ) -> tuple[np.ndarray, ...]: + nframes = coords.shape[0] + if len(atom_types.shape) == 1: + natoms = len(atom_types) + atom_types = np.tile(atom_types, nframes).reshape(nframes, -1) + else: + natoms = len(atom_types[0]) + + coord_input = coords.reshape([-1, natoms, 3]) + type_input = atom_types + box_input = cells.reshape([-1, 3, 3]) if cells is not None else None + if fparam is not None: + fparam_input = fparam.reshape(nframes, self.get_dim_fparam()) + elif self.dp.has_default_fparam(): + default_fparam = self.dp.get_default_fparam() + assert default_fparam is not None + fparam_input = np.tile( + np.array(default_fparam, dtype=GLOBAL_NP_FLOAT_PRECISION), + (nframes, 1), + ) + else: + fparam_input = None + aparam_input = ( + aparam.reshape(nframes, natoms, self.get_dim_aparam()) + if aparam is not None + else None + ) + + do_atomic_virial = any( + x.category == OutputVariableCategory.DERV_C_REDU for x in request_defs + ) + batch_output = self.dp( + coord_input, + type_input, + box=box_input, + fparam=fparam_input, + aparam=aparam_input, + do_atomic_virial=do_atomic_virial, + ) + + results = [] + for odef in request_defs: + dp_name = odef.name + shape = self._get_output_shape(odef, nframes, natoms) + if dp_name in batch_output and batch_output[dp_name] is not None: + results.append(batch_output[dp_name].reshape(shape)) + else: + results.append( + np.full(np.abs(shape), np.nan, dtype=GLOBAL_NP_FLOAT_PRECISION) + ) + return tuple(results) + + def _get_output_shape( + self, odef: OutputVariableDef, nframes: int, natoms: int + ) -> list[int]: + if odef.category == OutputVariableCategory.DERV_C_REDU: + return [nframes, *odef.shape[:-1], 9] + if odef.category == OutputVariableCategory.REDU: + return [nframes, *odef.shape, 1] + if odef.category == OutputVariableCategory.DERV_C: + return [nframes, *odef.shape[:-1], natoms, 9] + if odef.category == OutputVariableCategory.DERV_R: + return [nframes, *odef.shape[:-1], natoms, 3] + if odef.category == OutputVariableCategory.OUT: + return [nframes, natoms, *odef.shape, 1] + if odef.category == OutputVariableCategory.DERV_R_DERV_R: + return [nframes, 3 * natoms, 3 * natoms] + raise RuntimeError("unknown category") + + def get_model_def_script(self) -> dict: + return json.loads(self.dp.get_model_def_script()) + + def get_model(self) -> TF2SavedModelWrapper: + return self.dp diff --git a/deepmd/tf2/make_model.py b/deepmd/tf2/make_model.py new file mode 100644 index 0000000000..027e41fa4a --- /dev/null +++ b/deepmd/tf2/make_model.py @@ -0,0 +1,146 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import tensorflow as tf + +from deepmd.dpmodel.array_api import ( + Array, +) +from deepmd.dpmodel.model.transform_output import ( + communicate_extended_output, +) +from deepmd.dpmodel.output_def import ( + ModelOutputDef, +) +from deepmd.tf2.common import ( + to_tensorflow_array, + to_tf_tensor, + wrap_value, +) +from deepmd.tf2.env import ( + xp, +) +from deepmd.tf2.utils._dpmodel import ( + build_neighbor_list, + extend_coord_with_ghosts, + normalize_coord, +) + + +def _unwrap_tuple(values: tuple[Array, ...]) -> tuple[tf.Tensor, ...]: + return tuple(to_tf_tensor(value) for value in values) + + +def _box_has_pbc(box: Array | None) -> bool | None: + if box is None: + return False + last_dim = box.shape[-1] + return (last_dim != 0) if isinstance(last_dim, int) else None + + +def model_call_from_call_lower( + *, # enforce keyword-only arguments + call_lower: Callable[..., dict[str, Any]], + rcut: float, + sel: list[int], + mixed_types: bool, + model_output_def: ModelOutputDef, + coord: Array, + atype: Array, + box: Array | None, + fparam: Array | None, + aparam: Array | None, + do_atomic_virial: bool = False, +) -> dict[str, Array]: + """Return model prediction from lower interface. + + Parameters + ---------- + coord + The coordinates of the atoms. + shape: nf x (nloc x 3) + atype + The type of atoms. shape: nf x nloc + box + The simulation box. shape: nf x 9 + fparam + frame parameter. nf x ndf + aparam + atomic parameter. nf x nloc x nda + do_atomic_virial + If calculate the atomic virial. + + Returns + ------- + ret_dict + The result dict of type dict[str, Array]. + The keys are defined by the `ModelOutputDef`. + + """ + cc = to_tensorflow_array(coord) + atype = to_tensorflow_array(atype) + bb = to_tensorflow_array(box) + fp = to_tensorflow_array(fparam) + ap = to_tensorflow_array(aparam) + del coord, box, fparam, aparam + nframes, nloc = atype.shape[:2] + + def with_pbc() -> tuple[Array, Array, Array]: + assert bb is not None + coord_normalized = normalize_coord( + xp.reshape(cc, (nframes, nloc, 3)), + xp.reshape(bb, (nframes, 3, 3)), + ) + return extend_coord_with_ghosts(coord_normalized, atype, bb, rcut) + + def no_pbc() -> tuple[Array, Array, Array]: + return extend_coord_with_ghosts(cc, atype, None, rcut) + + has_pbc = _box_has_pbc(bb) + if has_pbc is True: + extended_coord, extended_atype, mapping = with_pbc() + elif has_pbc is False: + extended_coord, extended_atype, mapping = no_pbc() + else: + assert bb is not None + extended_coord_tensor, extended_atype_tensor, mapping_tensor = tf.cond( + tf.shape(to_tf_tensor(bb))[-1] != 0, + lambda: _unwrap_tuple(with_pbc()), + lambda: _unwrap_tuple(no_pbc()), + ) + extended_coord = to_tensorflow_array(extended_coord_tensor) + extended_atype = to_tensorflow_array(extended_atype_tensor) + mapping = to_tensorflow_array(mapping_tensor) + nlist = build_neighbor_list( + extended_coord, + extended_atype, + nloc, + rcut, + sel, + # types will be distinguished in the lower interface, + # so it doesn't need to be distinguished here + distinguish_types=False, + ) + extended_coord = xp.reshape(extended_coord, (nframes, -1, 3)) + model_predict_lower = wrap_value( + call_lower( + extended_coord, + extended_atype, + nlist, + mapping, + fparam=fp, + aparam=ap, + ) + ) + model_predict = communicate_extended_output( + model_predict_lower, + model_output_def, + mapping, + do_atomic_virial=do_atomic_virial, + ) + return model_predict diff --git a/deepmd/tf2/model/__init__.py b/deepmd/tf2/model/__init__.py new file mode 100644 index 0000000000..32ec725300 --- /dev/null +++ b/deepmd/tf2/model/__init__.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.tf2.atomic_model.linear_atomic_model import ( + DPZBLLinearEnergyAtomicModel, +) + +from .dipole_model import ( + DipoleModel, +) +from .dos_model import ( + DOSModel, +) +from .ener_model import ( + EnergyModel, +) +from .polar_model import ( + PolarModel, +) +from .property_model import ( + PropertyModel, +) + +__all__ = [ + "DOSModel", + "DPZBLLinearEnergyAtomicModel", + "DipoleModel", + "EnergyModel", + "PolarModel", + "PropertyModel", +] diff --git a/deepmd/tf2/model/base_model.py b/deepmd/tf2/model/base_model.py new file mode 100644 index 0000000000..dbb22b6f76 --- /dev/null +++ b/deepmd/tf2/model/base_model.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from deepmd.dpmodel.model.base_model import ( + make_base_model, +) +from deepmd.dpmodel.output_def import ( + get_deriv_name, + get_hessian_name, + get_reduce_name, +) +from deepmd.tf2.common import ( + to_tf_tensor, + wrap_tensor, +) +from deepmd.tf2.env import ( + tf, + xp, +) + +BaseModel = make_base_model() + + +def forward_common_atomic( + self: "BaseModel", + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + mapping: xp.ndarray | None = None, + fparam: xp.ndarray | None = None, + aparam: xp.ndarray | None = None, + do_atomic_virial: bool = False, + extended_coord_corr: xp.ndarray | None = None, + comm_dict: dict | None = None, + charge_spin: xp.ndarray | None = None, +) -> dict[str, xp.ndarray]: + del comm_dict # tf2 path has no MPI ghost exchange + + coord_tensor = to_tf_tensor(extended_coord) + assert coord_tensor is not None + coord_array = wrap_tensor(coord_tensor) + atomic_ret = self.atomic_model.forward_common_atomic( + coord_array, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + atomic_output_def = self.atomic_output_def() + model_predict = {} + for kk, vv in atomic_ret.items(): + model_predict[kk] = vv + vdef = atomic_output_def[kk] + atom_axis = -(len(vdef.shape) + 1) + if not vdef.reducible: + continue + + kk_redu = get_reduce_name(kk) + if vdef.intensive: + mask = atomic_ret["mask"] if "mask" in atomic_ret else None + if mask is not None: + model_predict[kk_redu] = xp.sum(vv, axis=atom_axis) / xp.sum( + mask, axis=-1, keepdims=True + ) + else: + model_predict[kk_redu] = xp.mean(vv, axis=atom_axis) + else: + model_predict[kk_redu] = xp.sum(vv, axis=atom_axis) + + kk_derv_r, kk_derv_c = get_deriv_name(kk) + if vdef.r_differentiable: + with tf.GradientTape() as tape: + tape.watch(coord_tensor) + grad_atomic_ret = self.atomic_model.forward_common_atomic( + wrap_tensor(coord_tensor), + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + reduced_output = xp.sum(grad_atomic_ret[kk], axis=atom_axis) + reduced_output_tensor = to_tf_tensor(reduced_output) + assert reduced_output_tensor is not None + ff_tensor = -tape.batch_jacobian(reduced_output_tensor, coord_tensor) + ff = wrap_tensor(ff_tensor) + + # extended_force: [nf, nall, *def, 3] + def_ndim = len(vdef.shape) + model_predict[kk_derv_r] = xp.transpose( + ff, [0, def_ndim + 1, *range(1, def_ndim + 1), def_ndim + 2] + ) + if vdef.r_hessian: + kk_hessian = get_hessian_name(kk) + model_predict[kk_hessian] = None + + if vdef.c_differentiable: + assert vdef.r_differentiable + # avr: [nf, *def, nall, 3, 3] + avr = xp.einsum("f...ai,faj->f...aij", ff, extended_coord) + if extended_coord_corr is not None: + avr = avr + xp.einsum("f...ai,faj->f...aij", ff, extended_coord_corr) + if do_atomic_virial: + with tf.GradientTape() as virial_tape: + virial_tape.watch(coord_tensor) + virial_atomic_ret = self.atomic_model.forward_common_atomic( + wrap_tensor(coord_tensor), + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + charge_spin=charge_spin, + ) + virial_atomic_tensor = to_tf_tensor(virial_atomic_ret[kk]) + assert virial_atomic_tensor is not None + nloc = tf.shape(nlist)[1] + loc_coord = tf.stop_gradient(coord_tensor[:, :nloc, :]) + loc_coord = tf.reshape( + loc_coord, + [ + tf.shape(loc_coord)[0], + tf.shape(loc_coord)[1], + *([1] * def_ndim), + 3, + ], + ) + corr_output = tf.reduce_sum( + virial_atomic_tensor[..., tf.newaxis] * loc_coord, + axis=1, + ) + virial_corr = virial_tape.batch_jacobian(corr_output, coord_tensor) + virial_corr = tf.transpose( + virial_corr, + [ + 0, + *range(1, def_ndim + 1), + def_ndim + 2, + def_ndim + 3, + def_ndim + 1, + ], + ) + avr = avr + wrap_tensor(virial_corr) + avr = xp.reshape(avr, [*ff.shape[:-1], 9]) + # extended_virial: [nf, nall, *def, 9] + extended_virial = xp.transpose( + avr, [0, def_ndim + 1, *range(1, def_ndim + 1), def_ndim + 2] + ) + model_predict[kk_derv_c] = extended_virial + # [nf, *def, 9] + model_predict[kk_derv_c + "_redu"] = xp.sum(extended_virial, axis=1) + return model_predict diff --git a/deepmd/tf2/model/dipole_model.py b/deepmd/tf2/model/dipole_model.py new file mode 100644 index 0000000000..6694883a0e --- /dev/null +++ b/deepmd/tf2/model/dipole_model.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from deepmd.dpmodel.model.dipole_model import DipoleModel as DipoleModelDP +from deepmd.tf2.atomic_model.dipole_atomic_model import ( + DPAtomicModelDipole, +) +from deepmd.tf2.model.base_model import ( + BaseModel, +) +from deepmd.tf2.model.dp_model import ( + make_tf2_dp_model_from_dpmodel, +) + + +@BaseModel.register("dipole") +class DipoleModel(make_tf2_dp_model_from_dpmodel(DipoleModelDP, DPAtomicModelDipole)): + pass diff --git a/deepmd/tf2/model/dos_model.py b/deepmd/tf2/model/dos_model.py new file mode 100644 index 0000000000..36859fd58e --- /dev/null +++ b/deepmd/tf2/model/dos_model.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.model.dos_model import DOSModel as DOSModelDP +from deepmd.tf2.atomic_model.dos_atomic_model import ( + DPAtomicModelDOS, +) +from deepmd.tf2.model.base_model import ( + BaseModel, +) +from deepmd.tf2.model.dp_model import ( + make_tf2_dp_model_from_dpmodel, +) + + +@BaseModel.register("dos") +class DOSModel(make_tf2_dp_model_from_dpmodel(DOSModelDP, DPAtomicModelDOS)): + pass diff --git a/deepmd/tf2/model/dp_model.py b/deepmd/tf2/model/dp_model.py new file mode 100644 index 0000000000..146e399050 --- /dev/null +++ b/deepmd/tf2/model/dp_model.py @@ -0,0 +1,136 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.model import ( + DPModelCommon, +) +from deepmd.dpmodel.utils.neighbor_list import ( + NeighborList, +) +from deepmd.tf2.atomic_model.dp_atomic_model import ( + DPAtomicModel, +) +from deepmd.tf2.common import ( + tf2_module, + to_tensorflow_array, +) +from deepmd.tf2.env import ( + stop_gradient, + xp, +) +from deepmd.tf2.model.base_model import ( + forward_common_atomic, +) + + +def make_tf2_dp_model_from_dpmodel( + dpmodel_model: type[DPModelCommon], tf2_atomicmodel: type[DPAtomicModel] +) -> type[DPModelCommon]: + """Make a tf2 backend DP model from a DPModel backend DP model. + + Parameters + ---------- + dpmodel_model : type[DPModelCommon] + The DPModel backend DP model. + tf2_atomicmodel : type[DPAtomicModel] + The tf2 backend DP atomic model. + + Returns + ------- + type[DPModelCommon] + The tf2 backend DP model. + """ + + @tf2_module + class tf2_model(dpmodel_model): + def call_common( + self, + coord: xp.ndarray, + atype: xp.ndarray, + box: xp.ndarray | None = None, + fparam: xp.ndarray | None = None, + aparam: xp.ndarray | None = None, + do_atomic_virial: bool = False, + coord_corr_for_virial: xp.ndarray | None = None, + charge_spin: xp.ndarray | None = None, + neighbor_list: NeighborList | None = None, + ) -> dict[str, xp.ndarray]: + return super().call_common( + to_tensorflow_array(coord), + to_tensorflow_array(atype), + box=to_tensorflow_array(box), + fparam=to_tensorflow_array(fparam), + aparam=to_tensorflow_array(aparam), + do_atomic_virial=do_atomic_virial, + coord_corr_for_virial=to_tensorflow_array(coord_corr_for_virial), + charge_spin=to_tensorflow_array(charge_spin), + neighbor_list=neighbor_list, + ) + + def call_common_lower( + self, + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + mapping: xp.ndarray | None = None, + fparam: xp.ndarray | None = None, + aparam: xp.ndarray | None = None, + do_atomic_virial: bool = False, + extended_coord_corr: xp.ndarray | None = None, + comm_dict: dict | None = None, + charge_spin: xp.ndarray | None = None, + ) -> dict[str, xp.ndarray]: + return super().call_common_lower( + to_tensorflow_array(extended_coord), + to_tensorflow_array(extended_atype), + to_tensorflow_array(nlist), + mapping=to_tensorflow_array(mapping), + fparam=to_tensorflow_array(fparam), + aparam=to_tensorflow_array(aparam), + do_atomic_virial=do_atomic_virial, + extended_coord_corr=to_tensorflow_array(extended_coord_corr), + comm_dict=comm_dict, + charge_spin=to_tensorflow_array(charge_spin), + ) + + def forward_common_atomic( + self, + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + mapping: xp.ndarray | None = None, + fparam: xp.ndarray | None = None, + aparam: xp.ndarray | None = None, + do_atomic_virial: bool = False, + extended_coord_corr: xp.ndarray | None = None, + comm_dict: dict | None = None, + charge_spin: xp.ndarray | None = None, + ) -> dict[str, xp.ndarray]: + del comm_dict # tf2 path has no MPI ghost exchange + return forward_common_atomic( + self, + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + extended_coord_corr=extended_coord_corr, + charge_spin=charge_spin, + ) + + def format_nlist( + self, + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + extra_nlist_sort: bool = False, + ) -> xp.ndarray: + return dpmodel_model.format_nlist( + self, + stop_gradient(extended_coord), + extended_atype, + nlist, + extra_nlist_sort=extra_nlist_sort, + ) + + return tf2_model diff --git a/deepmd/tf2/model/dp_zbl_model.py b/deepmd/tf2/model/dp_zbl_model.py new file mode 100644 index 0000000000..3bace216dc --- /dev/null +++ b/deepmd/tf2/model/dp_zbl_model.py @@ -0,0 +1,62 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.model.dp_zbl_model import DPZBLModel as DPZBLModelDP +from deepmd.tf2.atomic_model.linear_atomic_model import ( # noqa: F401 + DPZBLLinearEnergyAtomicModel as _DPZBLLinearEnergyAtomicModel, +) +from deepmd.tf2.common import ( + tf2_module, +) +from deepmd.tf2.env import ( + stop_gradient, + xp, +) +from deepmd.tf2.model.base_model import ( + BaseModel, + forward_common_atomic, +) + + +@BaseModel.register("zbl") +@tf2_module +class DPZBLModel(DPZBLModelDP): + def forward_common_atomic( + self, + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + mapping: xp.ndarray | None = None, + fparam: xp.ndarray | None = None, + aparam: xp.ndarray | None = None, + do_atomic_virial: bool = False, + extended_coord_corr: xp.ndarray | None = None, + comm_dict: dict | None = None, + charge_spin: xp.ndarray | None = None, + ) -> dict[str, xp.ndarray]: + del comm_dict # tf2 path has no MPI ghost exchange + return forward_common_atomic( + self, + extended_coord, + extended_atype, + nlist, + mapping=mapping, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + extended_coord_corr=extended_coord_corr, + charge_spin=charge_spin, + ) + + def format_nlist( + self, + extended_coord: xp.ndarray, + extended_atype: xp.ndarray, + nlist: xp.ndarray, + extra_nlist_sort: bool = False, + ) -> xp.ndarray: + return DPZBLModelDP.format_nlist( + self, + stop_gradient(extended_coord), + extended_atype, + nlist, + extra_nlist_sort=extra_nlist_sort, + ) diff --git a/deepmd/tf2/model/ener_model.py b/deepmd/tf2/model/ener_model.py new file mode 100644 index 0000000000..ca019bbaaa --- /dev/null +++ b/deepmd/tf2/model/ener_model.py @@ -0,0 +1,16 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.model import EnergyModel as EnergyModelDP +from deepmd.tf2.atomic_model.energy_atomic_model import ( + DPAtomicModelEnergy, +) +from deepmd.tf2.model.base_model import ( + BaseModel, +) +from deepmd.tf2.model.dp_model import ( + make_tf2_dp_model_from_dpmodel, +) + + +@BaseModel.register("ener") +class EnergyModel(make_tf2_dp_model_from_dpmodel(EnergyModelDP, DPAtomicModelEnergy)): + pass diff --git a/deepmd/tf2/model/model.py b/deepmd/tf2/model/model.py new file mode 100644 index 0000000000..2877cd60f6 --- /dev/null +++ b/deepmd/tf2/model/model.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from copy import ( + deepcopy, +) + +from deepmd.tf2.atomic_model.dp_atomic_model import ( + DPAtomicModel, +) +from deepmd.tf2.atomic_model.pairtab_atomic_model import ( + PairTabAtomicModel, +) +from deepmd.tf2.descriptor.base_descriptor import ( + BaseDescriptor, +) +from deepmd.tf2.fitting.base_fitting import ( + BaseFitting, +) +from deepmd.tf2.fitting.fitting import ( + EnergyFittingNet, +) +from deepmd.tf2.model.base_model import ( + BaseModel, +) +from deepmd.tf2.model.dp_zbl_model import ( + DPZBLModel, +) + + +def get_standard_model(data: dict) -> BaseModel: + """Get a Model from a dictionary. + + Parameters + ---------- + data : dict + The data to construct the model. + """ + data = deepcopy(data) + if "type_embedding" in data: + raise ValueError( + "In the tf2 backend, type_embedding is not at the model level, but within the descriptor. See type embedding documentation for details." + ) + descriptor_type = data["descriptor"].pop("type") + data["descriptor"]["type_map"] = data["type_map"] + data["descriptor"]["ntypes"] = len(data["type_map"]) + data["fitting_net"] = data.get("fitting_net", {}) + fitting_type = data["fitting_net"].pop("type", "ener") + data["fitting_net"]["type_map"] = data["type_map"] + descriptor = BaseDescriptor.get_class_by_type(descriptor_type)( + **data["descriptor"], + ) + if fitting_type in {"dipole", "polar"}: + data["fitting_net"]["embedding_width"] = descriptor.get_dim_emb() + fitting = BaseFitting.get_class_by_type(fitting_type)( + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + mixed_types=descriptor.mixed_types(), + **data["fitting_net"], + ) + return BaseModel.get_class_by_type(fitting_type)( + descriptor=descriptor, + fitting=fitting, + type_map=data["type_map"], + atom_exclude_types=data.get("atom_exclude_types", []), + pair_exclude_types=data.get("pair_exclude_types", []), + ) + + +def get_zbl_model(data: dict) -> DPZBLModel: + data["descriptor"]["ntypes"] = len(data["type_map"]) + descriptor_type = data["descriptor"].pop("type") + descriptor = BaseDescriptor.get_class_by_type(descriptor_type)(**data["descriptor"]) + fitting_type = data["fitting_net"].pop("type") + if fitting_type == "ener": + fitting = EnergyFittingNet( + ntypes=descriptor.get_ntypes(), + dim_descrpt=descriptor.get_dim_out(), + mixed_types=descriptor.mixed_types(), + **data["fitting_net"], + ) + else: + raise ValueError(f"Unknown fitting type {fitting_type}") + + dp_model = DPAtomicModel(descriptor, fitting, type_map=data["type_map"]) + # pairtab + filepath = data["use_srtab"] + pt_model = PairTabAtomicModel( + filepath, + data["descriptor"]["rcut"], + data["descriptor"]["sel"], + type_map=data["type_map"], + ) + rmin = data["sw_rmin"] + rmax = data["sw_rmax"] + atom_exclude_types = data.get("atom_exclude_types", []) + pair_exclude_types = data.get("pair_exclude_types", []) + return DPZBLModel( + dp_model, + pt_model, + rmin, + rmax, + type_map=data["type_map"], + atom_exclude_types=atom_exclude_types, + pair_exclude_types=pair_exclude_types, + ) + + +def get_model(data: dict) -> BaseModel: + """Get a model from a dictionary. + + Parameters + ---------- + data : dict + The data to construct the model. + """ + model_type = data.get("type", "standard") + if model_type == "standard": + if "spin" in data: + raise NotImplementedError("Spin model is not implemented yet.") + elif "use_srtab" in data: + return get_zbl_model(data) + else: + return get_standard_model(data) + else: + return BaseModel.get_class_by_type(model_type).get_model(data) diff --git a/deepmd/tf2/model/polar_model.py b/deepmd/tf2/model/polar_model.py new file mode 100644 index 0000000000..015888191c --- /dev/null +++ b/deepmd/tf2/model/polar_model.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from deepmd.dpmodel.model.polar_model import PolarModel as PolarModelDP +from deepmd.tf2.atomic_model.polar_atomic_model import ( + DPAtomicModelPolar, +) +from deepmd.tf2.model.base_model import ( + BaseModel, +) +from deepmd.tf2.model.dp_model import ( + make_tf2_dp_model_from_dpmodel, +) + + +@BaseModel.register("polar") +class PolarModel(make_tf2_dp_model_from_dpmodel(PolarModelDP, DPAtomicModelPolar)): + pass diff --git a/deepmd/tf2/model/property_model.py b/deepmd/tf2/model/property_model.py new file mode 100644 index 0000000000..87fbadaac4 --- /dev/null +++ b/deepmd/tf2/model/property_model.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later + +from deepmd.dpmodel.model.property_model import PropertyModel as PropertyModelDP +from deepmd.tf2.atomic_model.property_atomic_model import ( + DPAtomicModelProperty, +) +from deepmd.tf2.model.base_model import ( + BaseModel, +) +from deepmd.tf2.model.dp_model import ( + make_tf2_dp_model_from_dpmodel, +) + + +@BaseModel.register("property") +class PropertyModel( + make_tf2_dp_model_from_dpmodel(PropertyModelDP, DPAtomicModelProperty) +): + pass diff --git a/deepmd/tf2/transform_output.py b/deepmd/tf2/transform_output.py new file mode 100644 index 0000000000..6b1d22caab --- /dev/null +++ b/deepmd/tf2/transform_output.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import tensorflow as tf + +from deepmd.dpmodel.output_def import ( + ModelOutputDef, + OutputVariableDef, + get_deriv_name, + get_reduce_name, +) + + +def get_leading_dims( + vv: tf.Tensor, + vdef: OutputVariableDef, +) -> tf.Tensor: + """Get the dimensions of nf x nloc. + + Parameters + ---------- + vv : np.ndarray + The input array from which to compute the leading dimensions. + vdef : OutputVariableDef + The output variable definition containing the shape to exclude from `vv`. + + Returns + ------- + list + A list of leading dimensions of `vv`, excluding the last `len(vdef.shape)` dimensions. + """ + vshape = tf.shape(vv) + return vshape[: (vv.shape.rank - len(vdef.shape))] + + +def communicate_extended_output( + model_ret: dict[str, tf.Tensor], + model_output_def: ModelOutputDef, + mapping: tf.Tensor, # nf x nloc + do_atomic_virial: bool = False, +) -> dict[str, tf.Tensor]: + """Transform the output of the model network defined on + local and ghost (extended) atoms to local atoms. + + """ + new_ret = {} + for kk in model_output_def.keys_outp(): + vv = model_ret[kk] + vdef = model_output_def[kk] + new_ret[kk] = vv + if vdef.reducible: + kk_redu = get_reduce_name(kk) + new_ret[kk_redu] = model_ret[kk_redu] + kk_derv_r, kk_derv_c = get_deriv_name(kk) + mldims = tf.shape(mapping) + vldims = get_leading_dims(vv, vdef) + if vdef.r_differentiable: + if model_ret[kk_derv_r] is not None: + derv_r_ext_dims = list(vdef.shape) + [3] # noqa:RUF005 + indices = tf.reshape(mapping, [tf.shape(mapping)[0], -1, 1]) + # concat frame idx + indices = tf.concat( + [ + tf.reshape( + tf.repeat( + tf.range(tf.shape(indices)[0], dtype=indices.dtype), + tf.shape(mapping)[1], + ), + tf.shape(indices), + ), + indices, + ], + axis=-1, + ) + force = tf.scatter_nd( + indices, + model_ret[kk_derv_r], + tf.cast(tf.concat([vldims, derv_r_ext_dims], axis=0), tf.int64), + ) + new_ret[kk_derv_r] = tf.reshape( + force, + tf.concat([tf.shape(force)[:2], list(vdef.shape), [3]], axis=0), + ) + else: + # name holders + new_ret[kk_derv_r] = None + if vdef.c_differentiable: + assert vdef.r_differentiable + if model_ret[kk_derv_c] is not None: + derv_c_ext_dims = list(vdef.shape) + [9] # noqa:RUF005 + indices = tf.reshape(mapping, [tf.shape(mapping)[0], -1, 1]) + # concat frame idx + indices = tf.concat( + [ + tf.reshape( + tf.repeat( + tf.range(tf.shape(indices)[0], dtype=indices.dtype), + tf.shape(mapping)[1], + ), + tf.shape(indices), + ), + indices, + ], + axis=-1, + ) + virial = tf.scatter_nd( + indices, + model_ret[kk_derv_c], + tf.cast(tf.concat([vldims, derv_c_ext_dims], axis=0), tf.int64), + ) + new_ret[kk_derv_c] = tf.reshape( + virial, + tf.concat( + [tf.shape(virial)[:2], list(vdef.shape), [9]], axis=0 + ), + ) + new_ret[kk_derv_c + "_redu"] = tf.reduce_sum( + new_ret[kk_derv_c], axis=1 + ) + else: + new_ret[kk_derv_c] = None + new_ret[kk_derv_c + "_redu"] = None + if not do_atomic_virial: + # pop atomic virial, because it is not correctly calculated. + new_ret.pop(kk_derv_c) + return new_ret diff --git a/deepmd/tf2/utils/__init__.py b/deepmd/tf2/utils/__init__.py new file mode 100644 index 0000000000..6ceb116d85 --- /dev/null +++ b/deepmd/tf2/utils/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later diff --git a/deepmd/tf2/utils/_dpmodel.py b/deepmd/tf2/utils/_dpmodel.py new file mode 100644 index 0000000000..d09c1b9412 --- /dev/null +++ b/deepmd/tf2/utils/_dpmodel.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +from deepmd.dpmodel.utils.nlist import ( + build_neighbor_list as dpmodel_build_neighbor_list, +) +from deepmd.dpmodel.utils.nlist import ( + extend_coord_with_ghosts as dpmodel_extend_coord_with_ghosts, +) +from deepmd.dpmodel.utils.nlist import format_nlist as dpmodel_format_nlist +from deepmd.dpmodel.utils.region import inter2phys as dpmodel_inter2phys +from deepmd.dpmodel.utils.region import normalize_coord as dpmodel_normalize_coord +from deepmd.dpmodel.utils.region import to_face_distance as dpmodel_to_face_distance +from deepmd.tf2.common import ( + to_tensorflow_array, + to_tf_tensor, +) +from deepmd.tf2.env import ( + Array, +) + + +def build_neighbor_list( + coord: Any, + atype: Any, + nloc: int, + rcut: float, + sel: int | list[int], + distinguish_types: bool = True, +) -> Array: + ret = to_tf_tensor( + dpmodel_build_neighbor_list( + to_tensorflow_array(coord), + to_tensorflow_array(atype), + nloc, + rcut, + sel, + distinguish_types=distinguish_types, + ) + ) + ret.set_shape([None, None, sel if isinstance(sel, int) else sum(sel)]) + return to_tensorflow_array(ret) + + +def extend_coord_with_ghosts( + coord: Any, + atype: Any, + cell: Any | None, + rcut: float, +) -> tuple[Array, Array, Array]: + extended_coord, extended_atype, mapping = dpmodel_extend_coord_with_ghosts( + to_tensorflow_array(coord), + to_tensorflow_array(atype), + None if cell is None else to_tensorflow_array(cell), + rcut, + ) + return ( + to_tensorflow_array(extended_coord), + to_tensorflow_array(extended_atype), + to_tensorflow_array(mapping), + ) + + +def format_nlist( + extended_coord: Any, + nlist: Any, + nsel: int, + rcut: float, +) -> Array: + return to_tensorflow_array( + dpmodel_format_nlist( + to_tensorflow_array(extended_coord), + to_tensorflow_array(nlist), + nsel, + rcut, + ) + ) + + +def inter2phys(coord: Any, cell: Any) -> Array: + return to_tensorflow_array( + dpmodel_inter2phys(to_tensorflow_array(coord), to_tensorflow_array(cell)) + ) + + +def normalize_coord(coord: Any, cell: Any) -> Array: + return to_tensorflow_array( + dpmodel_normalize_coord(to_tensorflow_array(coord), to_tensorflow_array(cell)) + ) + + +def to_face_distance(cell: Any) -> Array: + return to_tensorflow_array(dpmodel_to_face_distance(to_tensorflow_array(cell))) + + +__all__ = [ + "build_neighbor_list", + "extend_coord_with_ghosts", + "format_nlist", + "inter2phys", + "normalize_coord", + "to_face_distance", +] diff --git a/deepmd/tf2/utils/exclude_mask.py b/deepmd/tf2/utils/exclude_mask.py new file mode 100644 index 0000000000..c76220fdea --- /dev/null +++ b/deepmd/tf2/utils/exclude_mask.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.utils.exclude_mask import AtomExcludeMask as AtomExcludeMaskDP +from deepmd.dpmodel.utils.exclude_mask import PairExcludeMask as PairExcludeMaskDP + +from ..common import ( + tf2_module, +) + + +@tf2_module +class AtomExcludeMask(AtomExcludeMaskDP): + pass + + +@tf2_module +class PairExcludeMask(PairExcludeMaskDP): + pass diff --git a/deepmd/tf2/utils/network.py b/deepmd/tf2/utils/network.py new file mode 100644 index 0000000000..966be78813 --- /dev/null +++ b/deepmd/tf2/utils/network.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, + ClassVar, +) + +import tensorflow as tf + +from deepmd.dpmodel.common import ( + PRECISION_DICT, + NativeOP, +) +from deepmd.dpmodel.utils.network import EmbeddingNet as EmbeddingNetDP +from deepmd.dpmodel.utils.network import FittingNet as FittingNetDP +from deepmd.dpmodel.utils.network import Identity as IdentityDP +from deepmd.dpmodel.utils.network import LayerNorm as LayerNormDP +from deepmd.dpmodel.utils.network import NativeLayer as NativeLayerDP +from deepmd.dpmodel.utils.network import NativeNet as NativeNetDP +from deepmd.dpmodel.utils.network import NetworkCollection as NetworkCollectionDP +from deepmd.dpmodel.utils.network import ( + make_embedding_network, + make_fitting_network, + make_multilayer_network, +) + +from ..common import ( + register_dpmodel_mapping, + tf2_module, + to_tensorflow_array, + to_tf_tensor, +) + + +class NativeLayer(NativeLayerDP, tf.Module): + _tf2_variable_attrs: ClassVar[set[str]] = {"w", "b", "idt"} + + def __init__(self, *args: Any, **kwargs: Any) -> None: + tf.Module.__init__(self) + NativeLayerDP.__init__(self, *args, **kwargs) + + @staticmethod + def _tf2_variable_storage_name(name: str) -> str: + return f"_tf2_{name}_variable" + + def _get_tf2_variable(self, name: str) -> tf.Variable | None: + return getattr(self, self._tf2_variable_storage_name(name), None) + + def _get_tf2_variable_array(self, name: str) -> Any | None: + variable = self._get_tf2_variable(name) + return None if variable is None else to_tensorflow_array(variable) + + def _set_tf2_variable(self, name: str, value: Any) -> None: + storage_name = self._tf2_variable_storage_name(name) + if value is None: + tf.Module.__setattr__(self, storage_name, None) + return + tensor = to_tf_tensor(value) + variable = tf.Variable( + tensor, + trainable=bool(getattr(self, "trainable", True)), + name=name, + ) + tf.Module.__setattr__(self, storage_name, variable) + + def _refresh_tf2_variable_trainability(self) -> None: + for name in self._tf2_variable_attrs: + variable = self._get_tf2_variable(name) + if variable is not None and variable.trainable != self.trainable: + self._set_tf2_variable(name, variable.read_value()) + + def __getattribute__(self, name: str) -> Any: + if name in {"w", "b", "idt"}: + return object.__getattribute__(self, "_get_tf2_variable_array")(name) + return super().__getattribute__(name) + + def __setattr__(self, name: str, value: Any) -> None: + if name in self._tf2_variable_attrs: + self._set_tf2_variable(name, value) + return + tf.Module.__setattr__(self, name, value) + if name == "trainable": + self._refresh_tf2_variable_trainability() + + def check_type_consistency(self) -> None: + precision = self.precision + + def check_var(var: Any | None) -> None: + if var is not None: + dtype_name = getattr(var.dtype, "name", str(var.dtype).split(".")[-1]) + assert PRECISION_DICT[dtype_name] is PRECISION_DICT[precision] + + check_var(self.w) + check_var(self.b) + check_var(self.idt) + + def serialize(self) -> dict: + data = super().serialize() + + def to_numpy(var: Any | None) -> Any | None: + tensor = to_tf_tensor(var) + return None if tensor is None else tensor.numpy() + + data["@variables"] = { + "w": to_numpy(self.w), + "b": to_numpy(self.b), + "idt": to_numpy(self.idt), + } + return data + + +@tf2_module +class NativeNet(make_multilayer_network(NativeLayer, NativeOP)): + pass + + +class EmbeddingNet(make_embedding_network(NativeNet, NativeLayer)): + pass + + +class FittingNet(make_fitting_network(EmbeddingNet, NativeNet, NativeLayer)): + pass + + +@tf2_module +class NetworkCollection(NetworkCollectionDP): + NETWORK_TYPE_MAP: ClassVar[dict[str, type]] = { + "network": NativeNet, + "embedding_network": EmbeddingNet, + "fitting_network": FittingNet, + } + + +class LayerNorm(LayerNormDP, NativeLayer): + pass + + +@tf2_module +class Identity(IdentityDP): + pass + + +register_dpmodel_mapping( + NativeNetDP, + lambda v: NativeNet.deserialize(v.serialize()), +) + +register_dpmodel_mapping( + EmbeddingNetDP, + lambda v: EmbeddingNet.deserialize(v.serialize()), +) + +register_dpmodel_mapping( + FittingNetDP, + lambda v: FittingNet.deserialize(v.serialize()), +) + +register_dpmodel_mapping( + NativeLayerDP, + lambda v: NativeLayer.deserialize(v.serialize()), +) + +register_dpmodel_mapping( + LayerNormDP, + lambda v: LayerNorm.deserialize(v.serialize()), +) + +register_dpmodel_mapping( + NetworkCollectionDP, + lambda v: NetworkCollection.deserialize(v.serialize()), +) + +register_dpmodel_mapping( + IdentityDP, + lambda v: Identity(), +) diff --git a/deepmd/tf2/utils/serialization.py b/deepmd/tf2/utils/serialization.py new file mode 100644 index 0000000000..f4aac52fd7 --- /dev/null +++ b/deepmd/tf2/utils/serialization.py @@ -0,0 +1,536 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import json +import os +from collections.abc import ( + Callable, +) +from typing import ( + Any, +) + +import numpy as np +import tensorflow as tf + +from deepmd._vendors import ndtensorflow as xp +from deepmd.tf2.common import ( + unwrap_value, +) +from deepmd.tf2.make_model import ( + model_call_from_call_lower, +) +from deepmd.tf2.model.base_model import ( + BaseModel, +) +from deepmd.tf2.utils._dpmodel import ( + format_nlist, +) + + +def _env_flag(name: str) -> bool: + return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"} + + +def _default_jit_compile() -> bool: + return _env_flag("DP_JIT") + + +class _ExportConstantArray: + """Array-like export constant that traces as ``tf.constant``.""" + + __array_priority__ = 2 + + def __init__(self, value: Any) -> None: + self._value = np.asarray(value) + + @property + def dtype(self) -> tf.DType: + return tf.as_dtype(self._value.dtype) + + @property + def shape(self) -> tuple[int, ...]: + return self._value.shape + + @property + def ndim(self) -> int: + return self._value.ndim + + @property + def size(self) -> int: + return self._value.size + + @property + def T(self) -> Any: + return self._array().T + + @property + def mT(self) -> Any: + return self._array().mT + + def _array(self) -> Any: + return xp.asarray(tf.constant(self._value)) + + def __array_namespace__(self, /, *, api_version: str | None = None) -> Any: + del api_version + return xp + + def __array__(self, dtype: Any | None = None) -> np.ndarray: + return self._value.astype(dtype) if dtype is not None else self._value + + def __tf_tensor__( + self, + dtype: tf.DType | None = None, + name: str | None = None, + ) -> tf.Tensor: + return tf.constant(self._value, dtype=dtype, name=name) + + def __getitem__(self, key: Any, /) -> Any: + return self._array()[key] + + def __len__(self) -> int: + return len(self._value) + + def __iter__(self) -> Any: + return iter(self._array()) + + def __bool__(self, /) -> bool: + return bool(self._value.item()) + + def __complex__(self, /) -> complex: + return complex(self._value.item()) + + def __float__(self, /) -> float: + return float(self._value.item()) + + def __index__(self, /) -> int: + return int(self._value.item()) + + def __int__(self, /) -> int: + return int(self._value.item()) + + def __repr__(self) -> str: + return f"_ExportConstantArray({self._value!r})" + + def astype( + self, + dtype: tf.DType, + /, + *, + copy: bool = True, + device: str | None = None, + ) -> Any: + return self._array().astype(dtype, copy=copy, device=device) + + def reshape(self, *shape: Any, copy: bool | None = None) -> Any: + return self._array().reshape(*shape, copy=copy) + + def ravel(self) -> Any: + return self._array().ravel() + + def squeeze(self, axis: int | tuple[int, ...] | None = None) -> Any: + return self._array().squeeze(axis=axis) + + def to_device( + self, + device: str, + /, + *, + stream: int | Any | None = None, + ) -> Any: + return self._array().to_device(device, stream=stream) + + def unwrap(self) -> tf.Tensor: + return tf.constant(self._value) + + +def _export_binary_forward(name: str) -> Callable[[Any, Any], Any]: + def method(self: _ExportConstantArray, other: Any, /) -> Any: + return getattr(xp, name)(self._array(), other) + + return method + + +def _export_binary_reflected(name: str) -> Callable[[Any, Any], Any]: + def method(self: _ExportConstantArray, other: Any, /) -> Any: + return getattr(xp, name)(other, self._array()) + + return method + + +def _export_unary(name: str) -> Callable[[Any], Any]: + def method(self: _ExportConstantArray) -> Any: + return getattr(xp, name)(self._array()) + + return method + + +_ExportConstantArray.__add__ = _export_binary_forward("add") # type: ignore[attr-defined] +_ExportConstantArray.__radd__ = _export_binary_reflected("add") # type: ignore[attr-defined] +_ExportConstantArray.__and__ = _export_binary_forward("bitwise_and") # type: ignore[attr-defined] +_ExportConstantArray.__rand__ = _export_binary_reflected("bitwise_and") # type: ignore[attr-defined] +_ExportConstantArray.__floordiv__ = _export_binary_forward("floor_divide") # type: ignore[attr-defined] +_ExportConstantArray.__rfloordiv__ = _export_binary_reflected("floor_divide") # type: ignore[attr-defined] +_ExportConstantArray.__ge__ = _export_binary_forward("greater_equal") # type: ignore[attr-defined] +_ExportConstantArray.__le__ = _export_binary_reflected("greater_equal") # type: ignore[attr-defined] +_ExportConstantArray.__gt__ = _export_binary_forward("greater") # type: ignore[attr-defined] +_ExportConstantArray.__lt__ = _export_binary_reflected("greater") # type: ignore[attr-defined] +_ExportConstantArray.__lshift__ = _export_binary_forward("bitwise_left_shift") # type: ignore[attr-defined] +_ExportConstantArray.__rlshift__ = _export_binary_reflected("bitwise_left_shift") # type: ignore[attr-defined] +_ExportConstantArray.__matmul__ = _export_binary_forward("matmul") # type: ignore[attr-defined] +_ExportConstantArray.__rmatmul__ = _export_binary_reflected("matmul") # type: ignore[attr-defined] +_ExportConstantArray.__mod__ = _export_binary_forward("remainder") # type: ignore[attr-defined] +_ExportConstantArray.__rmod__ = _export_binary_reflected("remainder") # type: ignore[attr-defined] +_ExportConstantArray.__mul__ = _export_binary_forward("multiply") # type: ignore[attr-defined] +_ExportConstantArray.__rmul__ = _export_binary_reflected("multiply") # type: ignore[attr-defined] +_ExportConstantArray.__or__ = _export_binary_forward("bitwise_or") # type: ignore[attr-defined] +_ExportConstantArray.__ror__ = _export_binary_reflected("bitwise_or") # type: ignore[attr-defined] +_ExportConstantArray.__pow__ = _export_binary_forward("pow") # type: ignore[attr-defined] +_ExportConstantArray.__rpow__ = _export_binary_reflected("pow") # type: ignore[attr-defined] +_ExportConstantArray.__rshift__ = _export_binary_forward("bitwise_right_shift") # type: ignore[attr-defined] +_ExportConstantArray.__rrshift__ = _export_binary_reflected("bitwise_right_shift") # type: ignore[attr-defined] +_ExportConstantArray.__sub__ = _export_binary_forward("subtract") # type: ignore[attr-defined] +_ExportConstantArray.__rsub__ = _export_binary_reflected("subtract") # type: ignore[attr-defined] +_ExportConstantArray.__truediv__ = _export_binary_forward("divide") # type: ignore[attr-defined] +_ExportConstantArray.__rtruediv__ = _export_binary_reflected("divide") # type: ignore[attr-defined] +_ExportConstantArray.__xor__ = _export_binary_forward("bitwise_xor") # type: ignore[attr-defined] +_ExportConstantArray.__rxor__ = _export_binary_reflected("bitwise_xor") # type: ignore[attr-defined] +_ExportConstantArray.__eq__ = _export_binary_forward("equal") # type: ignore[attr-defined] +_ExportConstantArray.__ne__ = _export_binary_forward("not_equal") # type: ignore[attr-defined] +_ExportConstantArray.__abs__ = _export_unary("abs") # type: ignore[attr-defined] +_ExportConstantArray.__invert__ = _export_unary("bitwise_invert") # type: ignore[attr-defined] +_ExportConstantArray.__neg__ = _export_unary("negative") # type: ignore[attr-defined] +_ExportConstantArray.__pos__ = _export_unary("positive") # type: ignore[attr-defined] + + +def _as_export_constant(value: Any) -> _ExportConstantArray: + if isinstance(value, xp.Array): + value = value.unwrap() + if isinstance(value, tf.Variable): + value = value.numpy() + if isinstance(value, tf.Tensor): + value = value.numpy() + return _ExportConstantArray(value) + + +def _freeze_tf2_constants_for_export(value: Any, seen: set[int]) -> Any: + """Freeze TensorFlow state inside deepmd objects to export constants.""" + if isinstance(value, (xp.Array, tf.Variable, tf.Tensor)): + return _as_export_constant(value) + if isinstance(value, list): + for ii, item in enumerate(value): + value[ii] = _freeze_tf2_constants_for_export(item, seen) + return value + if isinstance(value, dict): + for kk, item in list(value.items()): + value[kk] = _freeze_tf2_constants_for_export(item, seen) + return value + if isinstance(value, tuple): + return tuple(_freeze_tf2_constants_for_export(item, seen) for item in value) + if not hasattr(value, "__dict__"): + return value + if not type(value).__module__.startswith("deepmd."): + return value + + oid = id(value) + if oid in seen: + return value + seen.add(oid) + for name, item in list(value.__dict__.items()): + frozen = _freeze_tf2_constants_for_export(item, seen) + if frozen is not item: + object.__setattr__(value, name, frozen) + return value + + +def _freeze_tf2_variables_for_export(model: tf.Module) -> None: + """Freeze TF2 trainable state to constants for C++ SavedModel inference.""" + _freeze_tf2_constants_for_export(model, set()) + + +def deserialize_to_file( + model_file: str, data: dict, *, jit_compile: bool | None = None +) -> None: + """Deserialize the dictionary to a TensorFlow SavedModel.""" + if not model_file.endswith(".savedmodeltf"): + raise ValueError("TF2 backend only supports the .savedmodeltf extension") + return deserialize_to_savedmodel(model_file, data, jit_compile=jit_compile) + + +def deserialize_to_savedmodel( + model_file: str, data: dict, *, jit_compile: bool | None = None +) -> None: + """Deserialize the dictionary to a TensorFlow SavedModel directory.""" + if jit_compile is None: + jit_compile = _default_jit_compile() + + # Import model registrations before deserializing the dpmodel payload. + import deepmd.tf2.model.model # noqa: F401 + + model = BaseModel.deserialize(data["model"]) + _freeze_tf2_variables_for_export(model) + model_def_script = data["model_def_script"] + + tf_model = tf.Module() + + def call_lower_with_fixed_do_atomic_virial( + do_atomic_virial: bool, + ) -> Callable: + def call_lower( + coord: tf.Tensor, + atype: tf.Tensor, + nlist: tf.Tensor, + mapping: tf.Tensor, + fparam: tf.Tensor, + aparam: tf.Tensor, + ) -> dict[str, tf.Tensor]: + return unwrap_value( + model.call_common_lower( + coord, + atype, + nlist, + mapping, + fparam, + aparam, + do_atomic_virial=do_atomic_virial, + ) + ) + + return call_lower + + @tf.function( + autograph=True, + jit_compile=jit_compile, + input_signature=[ + tf.TensorSpec([None, None, 3], tf.float64), + tf.TensorSpec([None, None], tf.int32), + tf.TensorSpec([None, None, None], tf.int64), + tf.TensorSpec([None, None], tf.int64), + tf.TensorSpec([None, model.get_dim_fparam()], tf.float64), + tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), + ], + ) + def call_lower_without_atomic_virial( + coord: tf.Tensor, + atype: tf.Tensor, + nlist: tf.Tensor, + mapping: tf.Tensor, + fparam: tf.Tensor, + aparam: tf.Tensor, + ) -> dict[str, tf.Tensor]: + nlist = format_nlist(coord, nlist, model.get_nnei(), model.get_rcut()) + return call_lower_with_fixed_do_atomic_virial(False)( + coord, atype, nlist, mapping, fparam, aparam + ) + + tf_model.call_lower = call_lower_without_atomic_virial + + @tf.function( + autograph=True, + jit_compile=jit_compile, + input_signature=[ + tf.TensorSpec([None, None, 3], tf.float64), + tf.TensorSpec([None, None], tf.int32), + tf.TensorSpec([None, None, None], tf.int64), + tf.TensorSpec([None, None], tf.int64), + tf.TensorSpec([None, model.get_dim_fparam()], tf.float64), + tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), + ], + ) + def call_lower_with_atomic_virial( + coord: tf.Tensor, + atype: tf.Tensor, + nlist: tf.Tensor, + mapping: tf.Tensor, + fparam: tf.Tensor, + aparam: tf.Tensor, + ) -> dict[str, tf.Tensor]: + nlist = format_nlist(coord, nlist, model.get_nnei(), model.get_rcut()) + return call_lower_with_fixed_do_atomic_virial(True)( + coord, atype, nlist, mapping, fparam, aparam + ) + + tf_model.call_lower_atomic_virial = call_lower_with_atomic_virial + + def make_call_whether_do_atomic_virial(do_atomic_virial: bool) -> Callable: + call_lower = ( + call_lower_with_atomic_virial + if do_atomic_virial + else call_lower_without_atomic_virial + ) + + def call( + coord: tf.Tensor, + atype: tf.Tensor, + box: tf.Tensor, + fparam: tf.Tensor, + aparam: tf.Tensor, + ) -> dict[str, tf.Tensor]: + return unwrap_value( + model_call_from_call_lower( + call_lower=call_lower, + rcut=model.get_rcut(), + sel=model.get_sel(), + mixed_types=model.mixed_types(), + model_output_def=model.model_output_def(), + coord=coord, + atype=atype, + box=box, + fparam=fparam, + aparam=aparam, + do_atomic_virial=do_atomic_virial, + ) + ) + + return call + + @tf.function( + autograph=True, + input_signature=[ + tf.TensorSpec([None, None, 3], tf.float64), + tf.TensorSpec([None, None], tf.int32), + tf.TensorSpec([None, None, None], tf.float64), + tf.TensorSpec([None, model.get_dim_fparam()], tf.float64), + tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), + ], + ) + def call_with_atomic_virial( + coord: tf.Tensor, + atype: tf.Tensor, + box: tf.Tensor, + fparam: tf.Tensor, + aparam: tf.Tensor, + ) -> dict[str, tf.Tensor]: + return make_call_whether_do_atomic_virial(True)( + coord, atype, box, fparam, aparam + ) + + tf_model.call_atomic_virial = call_with_atomic_virial + + @tf.function( + autograph=True, + input_signature=[ + tf.TensorSpec([None, None, 3], tf.float64), + tf.TensorSpec([None, None], tf.int32), + tf.TensorSpec([None, None, None], tf.float64), + tf.TensorSpec([None, model.get_dim_fparam()], tf.float64), + tf.TensorSpec([None, None, model.get_dim_aparam()], tf.float64), + ], + ) + def call_without_atomic_virial( + coord: tf.Tensor, + atype: tf.Tensor, + box: tf.Tensor, + fparam: tf.Tensor, + aparam: tf.Tensor, + ) -> dict[str, tf.Tensor]: + return make_call_whether_do_atomic_virial(False)( + coord, atype, box, fparam, aparam + ) + + tf_model.call = call_without_atomic_virial + + @tf.function + def get_type_map() -> tf.Tensor: + return tf.constant(model.get_type_map(), dtype=tf.string) + + tf_model.get_type_map = get_type_map + + @tf.function + def get_rcut() -> tf.Tensor: + return tf.constant(model.get_rcut(), dtype=tf.double) + + tf_model.get_rcut = get_rcut + + @tf.function + def get_dim_fparam() -> tf.Tensor: + return tf.constant(model.get_dim_fparam(), dtype=tf.int64) + + tf_model.get_dim_fparam = get_dim_fparam + + @tf.function + def get_dim_aparam() -> tf.Tensor: + return tf.constant(model.get_dim_aparam(), dtype=tf.int64) + + tf_model.get_dim_aparam = get_dim_aparam + + @tf.function + def get_sel_type() -> tf.Tensor: + return tf.constant(model.get_sel_type(), dtype=tf.int64) + + tf_model.get_sel_type = get_sel_type + + @tf.function + def is_aparam_nall() -> tf.Tensor: + return tf.constant(model.is_aparam_nall(), dtype=tf.bool) + + tf_model.is_aparam_nall = is_aparam_nall + + @tf.function + def model_output_type() -> tf.Tensor: + return tf.constant(model.model_output_type(), dtype=tf.string) + + tf_model.model_output_type = model_output_type + + @tf.function + def mixed_types() -> tf.Tensor: + return tf.constant(model.mixed_types(), dtype=tf.bool) + + tf_model.mixed_types = mixed_types + + if model.get_min_nbor_dist() is not None: + + @tf.function + def get_min_nbor_dist() -> tf.Tensor: + return tf.constant(model.get_min_nbor_dist(), dtype=tf.double) + + tf_model.get_min_nbor_dist = get_min_nbor_dist + + @tf.function + def get_sel() -> tf.Tensor: + return tf.constant(model.get_sel(), dtype=tf.int64) + + tf_model.get_sel = get_sel + + @tf.function + def get_model_def_script() -> tf.Tensor: + return tf.constant( + json.dumps(model_def_script, separators=(",", ":")), dtype=tf.string + ) + + tf_model.get_model_def_script = get_model_def_script + + @tf.function + def has_message_passing() -> tf.Tensor: + return tf.constant(model.has_message_passing(), dtype=tf.bool) + + tf_model.has_message_passing = has_message_passing + tf_model.do_message_passing = has_message_passing + + @tf.function + def has_default_fparam() -> tf.Tensor: + return tf.constant(model.has_default_fparam(), dtype=tf.bool) + + tf_model.has_default_fparam = has_default_fparam + + @tf.function + def get_default_fparam() -> tf.Tensor: + default_fparam = model.get_default_fparam() + if default_fparam is None: + return tf.constant([], dtype=tf.double) + return tf.constant(default_fparam, dtype=tf.double) + + tf_model.get_default_fparam = get_default_fparam + + tf.saved_model.save( + tf_model, + model_file, + options=tf.saved_model.SaveOptions(experimental_custom_gradients=True), + ) + + +def serialize_from_file(model_file: str) -> dict: + """Serialize a TF2 SavedModel to a dictionary. + + SavedModel does not currently carry enough structured variable metadata to + round-trip back to the DeePMD dictionary format. + """ + raise ValueError(f"TF2 backend cannot serialize {model_file!r} to a model dict") diff --git a/deepmd/tf2/utils/type_embed.py b/deepmd/tf2/utils/type_embed.py new file mode 100644 index 0000000000..8bc6cd6eb1 --- /dev/null +++ b/deepmd/tf2/utils/type_embed.py @@ -0,0 +1,12 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +from deepmd.dpmodel.utils.type_embed import TypeEmbedNet as TypeEmbedNetDP + +from ..common import ( + tf2_module, +) +from . import network as _tf2_network # noqa: F401 + + +@tf2_module +class TypeEmbedNet(TypeEmbedNetDP): + pass diff --git a/pyproject.toml b/pyproject.toml index 35fc0fdb18..8541ccf413 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -444,7 +444,9 @@ runtime-evaluated-base-classes = ["torch.nn.Module"] "source/3rdparty/**" = ["ALL"] "backend/**" = ["ANN"] "data/**" = ["ANN"] +"deepmd/_vendors/**" = ["ALL"] "deepmd/tf/**" = ["TID253"] +"deepmd/tf2/**" = ["TID253"] "deepmd/pt/**" = ["TID253", "B905"] "deepmd/pt_expt/**" = ["TID253", "B905"] "deepmd/jax/**" = ["TID253"] diff --git a/source/api_cc/src/DeepPotJAX.cc b/source/api_cc/src/DeepPotJAX.cc index 8bc76e9edd..1549d5216e 100644 --- a/source/api_cc/src/DeepPotJAX.cc +++ b/source/api_cc/src/DeepPotJAX.cc @@ -321,11 +321,16 @@ void deepmd::DeepPotJAX::init(const std::string& model, sel = get_vector(ctx, "get_sel", func_vector, device, status); nnei = std::accumulate(sel.begin(), sel.end(), decltype(sel)::value_type(0)); try { - do_message_passing = get_scalar(ctx, "do_message_passing", + do_message_passing = get_scalar(ctx, "has_message_passing", func_vector, device, status); } catch (tf_function_not_found& e) { - // compatibile with models generated by v3.0.0rc0 - do_message_passing = false; + try { + do_message_passing = get_scalar(ctx, "do_message_passing", + func_vector, device, status); + } catch (tf_function_not_found& e) { + // compatible with models generated by v3.0.0rc0 + do_message_passing = false; + } } try { has_default_fparam_ = get_scalar(ctx, "has_default_fparam", diff --git a/source/api_cc/src/common.cc b/source/api_cc/src/common.cc index 0f59bb0e04..ded21e0c66 100644 --- a/source/api_cc/src/common.cc +++ b/source/api_cc/src/common.cc @@ -1465,6 +1465,10 @@ void deepmd::print_summary(const std::string& pre) { } deepmd::DPBackend deepmd::get_backend(const std::string& model) { + auto has_suffix = [](const std::string& value, const std::string& suffix) { + return value.length() >= suffix.length() && + value.substr(value.length() - suffix.length()) == suffix; + }; if (model.length() >= 4 && model.substr(model.length() - 4) == ".pth") { return deepmd::DPBackend::PyTorch; } else if (model.length() >= 4 && @@ -1472,8 +1476,8 @@ deepmd::DPBackend deepmd::get_backend(const std::string& model) { return deepmd::DPBackend::PyTorchExportable; } else if (model.length() >= 3 && model.substr(model.length() - 3) == ".pb") { return deepmd::DPBackend::TensorFlow; - } else if (model.length() >= 11 && - model.substr(model.length() - 11) == ".savedmodel") { + } else if (has_suffix(model, ".savedmodel") || + has_suffix(model, ".savedmodeltf")) { return deepmd::DPBackend::JAX; } else if ((model.length() >= 5 && model.substr(model.length() - 5) == ".json") || diff --git a/source/api_cc/tests/deeppot_universal_test_common.h b/source/api_cc/tests/deeppot_universal_test_common.h index 565a151760..29be32acd1 100644 --- a/source/api_cc/tests/deeppot_universal_test_common.h +++ b/source/api_cc/tests/deeppot_universal_test_common.h @@ -124,6 +124,12 @@ inline std::vector model_cases() { /*supports_float=*/true, /*supports_nframes=*/false, /*supports_lmp_nlist_mapping=*/true}, + {"tf2_savedmodeltf", Backend::JAX, + "../../tests/infer/deeppot_sea.savedmodeltf", + /*convert_pbtxt=*/false, &sea_deeppot_ref(), nullptr, 1e-10, 1e-4, + /*supports_float=*/true, + /*supports_nframes=*/false, + /*supports_lmp_nlist_mapping=*/true}, {"paddle_json", Backend::Paddle, "../../tests/infer/deeppot_sea.json", /*convert_pbtxt=*/false, &sea_deeppot_ref(), nullptr, 1e-7, 1e-4, /*supports_float=*/false, diff --git a/source/api_cc/tests/test_deeppot_universal.cc b/source/api_cc/tests/test_deeppot_universal.cc index e0ee6fc8f4..2a5ef05be8 100644 --- a/source/api_cc/tests/test_deeppot_universal.cc +++ b/source/api_cc/tests/test_deeppot_universal.cc @@ -99,6 +99,28 @@ std::vector variant_deeppot_cases() { /*supports_no_pbc_atomic=*/true, /*supports_no_pbc_lmp_nlist=*/true, /*supports_no_pbc_lmp_nlist_atomic=*/true}, + {"dpa_tf2_savedmodeltf", + Backend::JAX, + "../../tests/infer/deeppot_dpa.savedmodeltf", + /*convert_pbtxt=*/false, + &deepmd_test::jax_dpa_deeppot_ref(), + &deepmd_test::jax_dpa_deeppot_no_pbc_ref(), + "", + "", + "", + 1e-7, + 1e-1, + /*supports_float=*/true, + /*supports_finite_difference=*/false, + /*supports_lmp_nlist=*/false, + /*supports_lmp_nlist_atomic=*/false, + /*supports_lmp_nlist_cutoff_twice=*/false, + /*supports_lmp_nlist_type_sel=*/false, + /*supports_print_summary=*/false, + /*supports_no_pbc_simple=*/true, + /*supports_no_pbc_atomic=*/true, + /*supports_no_pbc_lmp_nlist=*/true, + /*supports_no_pbc_lmp_nlist_atomic=*/true}, {"dpa1_pytorch_pth", Backend::PyTorch, "../../tests/infer/deeppot_dpa1.pth", diff --git a/source/jax2tf_tests/test_format_nlist.py b/source/jax2tf_tests/test_format_nlist.py index cd157ac974..20201147ad 100644 --- a/source/jax2tf_tests/test_format_nlist.py +++ b/source/jax2tf_tests/test_format_nlist.py @@ -89,3 +89,31 @@ def test_format_nlist_larger_rcut(self) -> None: nlist = format_nlist(self.ecoord, nlist, sum(self.nsel), self.rcut) # we only need to ensure the result is correct, no need to check the order self.assertAllEqual(tnp.sort(nlist, axis=-1), tnp.sort(self.nlist, axis=-1)) + + def test_format_nlist_dynamic_nnei_graph(self) -> None: + @tf.function( + input_signature=[ + tf.TensorSpec([None, None, 3], tf.float64), + tf.TensorSpec([None, None, None], tf.int64), + ] + ) + def graph_format_nlist( + extended_coord: tf.Tensor, nlist: tf.Tensor + ) -> tf.Tensor: + return format_nlist(extended_coord, nlist, 3, 1.01) + + extended_coord = tf.constant( + [[[0.0, 0.0, 0.0], [0.5, 0.0, 0.0], [1.5, 0.0, 0.0], [3.0, 0.0, 0.0]]], + dtype=tf.float64, + ) + expected = tf.constant([[[1, -1, -1], [0, 2, -1]]], dtype=tf.int64) + + for nlist in [ + [[[1, 2], [0, 2]]], + [[[1, 2, -1], [0, 2, 3]]], + [[[2, 1, 3, -1], [3, 2, 0, -1]]], + ]: + self.assertAllEqual( + graph_format_nlist(extended_coord, tf.constant(nlist, tf.int64)), + expected, + ) diff --git a/source/jax2tf_tests/test_ndtensorflow_array.py b/source/jax2tf_tests/test_ndtensorflow_array.py new file mode 100644 index 0000000000..56a2d06556 --- /dev/null +++ b/source/jax2tf_tests/test_ndtensorflow_array.py @@ -0,0 +1,42 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +import numpy as np +import tensorflow as tf + +from deepmd._vendors import ndtensorflow as xp + + +def test_dynamic_boolean_mask_indexing() -> None: + @tf.function( + input_signature=[ + tf.TensorSpec([None, None, 3], tf.float64), + tf.TensorSpec([None, None], tf.bool), + ] + ) + def mask_values(array: tf.Tensor, mask: tf.Tensor) -> tf.Tensor: + return xp.asarray(array)[mask].unwrap() + + array = tf.reshape(tf.range(18, dtype=tf.float64), (2, 3, 3)) + mask = tf.constant([[True, False, True], [False, True, False]]) + + np.testing.assert_equal( + mask_values(array, mask).numpy(), + np.array( + [ + [0.0, 1.0, 2.0], + [6.0, 7.0, 8.0], + [12.0, 13.0, 14.0], + ] + ), + ) + + +def test_empty_boolean_mask_indexing() -> None: + out = xp.asarray(tf.constant([False]))[xp.asarray(tf.constant([], dtype=tf.bool))] + + np.testing.assert_equal(out.unwrap().numpy(), np.array([], dtype=bool)) + + +def test_tile_uint16() -> None: + out = xp.tile(xp.asarray(tf.constant(0, dtype=tf.uint16)), (2,)) + + np.testing.assert_equal(out.unwrap().numpy(), np.array([0, 0], dtype=np.uint16)) diff --git a/source/jax2tf_tests/test_nlist.py b/source/jax2tf_tests/test_nlist.py index c5906fa2fc..8ac9b8daa5 100644 --- a/source/jax2tf_tests/test_nlist.py +++ b/source/jax2tf_tests/test_nlist.py @@ -27,8 +27,8 @@ def setUp(self) -> None: [self.cell, self.icoord, self.atype] = [ tnp.expand_dims(ii, 0) for ii in [self.cell, self.icoord, self.atype] ] - self.coord = inter2phys(self.icoord, self.cell).reshape([-1, self.nloc * 3]) - self.cell = self.cell.reshape([-1, 9]) + self.coord = tf.reshape(inter2phys(self.icoord, self.cell), [-1, self.nloc * 3]) + self.cell = tf.reshape(self.cell, [-1, 9]) [self.cell, self.coord, self.atype] = [ tnp.tile(ii, [self.nf, 1]) for ii in [self.cell, self.coord, self.atype] ] @@ -57,7 +57,7 @@ def test_build_notype(self) -> None: ) self.assertAllClose(nlist[0], nlist[1]) nlist_mask = nlist[0] == -1 - nlist_loc = mapping[0][nlist[0]] + nlist_loc = tf.gather(mapping[0], tf.where(nlist_mask, 0, nlist[0])) nlist_loc = tnp.where(nlist_mask, tnp.full_like(nlist_loc, -1), nlist_loc) self.assertAllClose( tnp.sort(nlist_loc, axis=-1), @@ -78,7 +78,7 @@ def test_build_type(self) -> None: ) self.assertAllClose(nlist[0], nlist[1]) nlist_mask = nlist[0] == -1 - nlist_loc = mapping[0][nlist[0]] + nlist_loc = tf.gather(mapping[0], tf.where(nlist_mask, 0, nlist[0])) nlist_loc = tnp.where(nlist_mask, tnp.full_like(nlist_loc, -1), nlist_loc) for ii in range(2): self.assertAllClose( @@ -100,16 +100,16 @@ def test_extend_coord(self) -> None: ) # check the shift vectors are aligned with grid shift_vec = ( - ecoord.reshape([-1, self.ns, self.nloc, 3]) - - self.coord.reshape([-1, self.nloc, 3])[:, None, :, :] + tf.reshape(ecoord, [-1, self.ns, self.nloc, 3]) + - tf.reshape(self.coord, [-1, self.nloc, 3])[:, None, :, :] ) - shift_vec = shift_vec.reshape([-1, self.nall, 3]) + shift_vec = tf.reshape(shift_vec, [-1, self.nall, 3]) # hack!!! assumes identical cell across frames shift_vec = tnp.matmul( - shift_vec, tf.linalg.inv(self.cell.reshape([self.nf, 3, 3])[0]) + shift_vec, tf.linalg.inv(tf.reshape(self.cell, [self.nf, 3, 3])[0]) ) # nf x nall x 3 - shift_vec = tnp.round(shift_vec) + shift_vec = tf.round(shift_vec) # check: identical shift vecs self.assertAllClose(shift_vec[0], shift_vec[1], rtol=self.prec, atol=self.prec) # check: shift idx aligned with grid diff --git a/source/tests/common/test_pretrained_backend.py b/source/tests/common/test_pretrained_backend.py index c3c58d17da..b55cd11b9c 100644 --- a/source/tests/common/test_pretrained_backend.py +++ b/source/tests/common/test_pretrained_backend.py @@ -35,6 +35,16 @@ def test_detect_backend_by_pretrained_suffix_not_supported(self) -> None: with self.assertRaises(ValueError): Backend.detect_backend_by_model("DPA-3.2-5M.pretrained") + def test_detect_savedmodel_suffix_split(self) -> None: + self.assertEqual( + Backend.detect_backend_by_model("model.savedmodel").name, + "JAX", + ) + self.assertEqual( + Backend.detect_backend_by_model("model.savedmodeltf").name, + "TensorFlow2", + ) + def test_parse_pretrained_alias_plain_name(self) -> None: self.assertEqual(parse_pretrained_alias("DPA-3.2-5M"), "DPA-3.2-5M") self.assertEqual(parse_pretrained_alias("dpa-3.2-5m"), "DPA-3.2-5M") diff --git a/source/tests/consistent/common.py b/source/tests/consistent/common.py index d8fb9e64e0..44632ca6eb 100644 --- a/source/tests/consistent/common.py +++ b/source/tests/consistent/common.py @@ -37,10 +37,22 @@ from ..utils import ( CI, + DP_TEST_TF2_ONLY, TEST_DEVICE, ) -INSTALLED_TF = Backend.get_backend("tensorflow")().is_available() +RUN_TF2_BACKEND_TESTS = os.environ.get("DEEPMD_TEST_TF2") == "1" or DP_TEST_TF2_ONLY + +INSTALLED_TF = ( + not RUN_TF2_BACKEND_TESTS and Backend.get_backend("tensorflow")().is_available() +) +try: + _TF2_BACKEND = Backend.get_backend("tf2") +except (KeyError, RuntimeError): + _TF2_BACKEND = None +INSTALLED_TF2 = ( + RUN_TF2_BACKEND_TESTS and _TF2_BACKEND is not None and _TF2_BACKEND().is_available() +) INSTALLED_PT = Backend.get_backend("pytorch")().is_available() try: _PT_EXPT_BACKEND = Backend.get_backend("pytorch-exportable") @@ -51,7 +63,11 @@ INSTALLED_PD = Backend.get_backend("paddle")().is_available() INSTALLED_ARRAY_API_STRICT = find_spec("array_api_strict") is not None -if os.environ.get("CI") and not (INSTALLED_TF and INSTALLED_PT and INSTALLED_PD): +if ( + os.environ.get("CI") + and not RUN_TF2_BACKEND_TESTS + and not (INSTALLED_TF and INSTALLED_PT and INSTALLED_PD) +): raise ImportError("TensorFlow, PyTorch or Paddle should be tested in the CI") @@ -75,6 +91,8 @@ "INSTALLED_PT", "INSTALLED_PT_EXPT", "INSTALLED_TF", + "INSTALLED_TF2", + "RUN_TF2_BACKEND_TESTS", "CommonTest", "parameterize_func", "parameterized", @@ -91,6 +109,8 @@ class CommonTest(ABC): """Additional data that will not be checked.""" tf_class: ClassVar[type | None] """TensorFlow model class.""" + tf2_class: ClassVar[type | None] = None + """TensorFlow 2 model class.""" dp_class: ClassVar[type | None] """Native DP model class.""" pt_class: ClassVar[type | None] @@ -108,6 +128,8 @@ class CommonTest(ABC): """Whether to skip the native DP model.""" skip_tf: ClassVar[bool] = not INSTALLED_TF """Whether to skip the TensorFlow model.""" + skip_tf2: ClassVar[bool] = not INSTALLED_TF2 + """Whether to skip the TensorFlow 2 model.""" skip_pt: ClassVar[bool] = not INSTALLED_PT """Whether to skip the PyTorch model.""" skip_pt_expt: ClassVar[bool] = not INSTALLED_PT_EXPT @@ -199,6 +221,16 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: """ raise NotImplementedError("Not implemented") + def eval_tf2(self, tf2_obj: Any) -> Any: + """Evaluate the return value of TensorFlow 2. + + Parameters + ---------- + tf2_obj : Any + The object of TensorFlow 2 + """ + raise NotImplementedError("Not implemented") + def eval_jax(self, jax_obj: Any) -> Any: """Evaluate the return value of JAX. @@ -239,6 +271,7 @@ class RefBackend(Enum): PD = 5 JAX = 6 ARRAY_API_STRICT = 7 + TF2 = 8 @abstractmethod def extract_ret( @@ -341,6 +374,11 @@ def get_pt_expt_ret_serialization_from_cls(self, obj): data = obj.serialize() return ret, data + def get_tf2_ret_serialization_from_cls(self, obj): + ret = self.eval_tf2(obj) + data = obj.serialize() + return ret, data + def get_jax_ret_serialization_from_cls(self, obj): ret = self.eval_jax(obj) data = obj.serialize() @@ -375,6 +413,8 @@ def get_reference_backend(self): return self.RefBackend.PD if not self.skip_array_api_strict: return self.RefBackend.ARRAY_API_STRICT + if not self.skip_tf2 and self.tf2_class is not None: + return self.RefBackend.TF2 raise ValueError("No available reference") def get_reference_ret_serialization(self, ref: RefBackend): @@ -402,6 +442,11 @@ def get_reference_ret_serialization(self, ref: RefBackend): if ref == self.RefBackend.ARRAY_API_STRICT: obj = self.init_backend_cls(self.array_api_strict_class) return self.get_array_api_strict_ret_serialization_from_cls(obj) + if ref == self.RefBackend.TF2: + if self.tf2_class is None: + raise ValueError("TF2 class is not set") + obj = self.init_backend_cls(self.tf2_class) + return self.get_tf2_ret_serialization_from_cls(obj) raise ValueError("No available reference") def test_tf_consistent_with_ref(self) -> None: @@ -558,6 +603,45 @@ def test_pt_expt_self_consistent(self) -> None: else: self.assertEqual(rr1, rr2) + @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") + def test_tf2_consistent_with_ref(self) -> None: + """Test whether TF2 and reference are consistent.""" + if self.skip_tf2 or self.tf2_class is None: + self.skipTest("Unsupported backend") + ref_backend = self.get_reference_backend() + if ref_backend == self.RefBackend.TF2: + self.skipTest("Reference is self") + ret1, data1 = self.get_reference_ret_serialization(ref_backend) + ret1 = self.extract_ret(ret1, ref_backend) + obj = self.tf2_class.deserialize(data1) + ret2 = self.eval_tf2(obj) + ret2 = self.extract_ret(ret2, self.RefBackend.TF2) + data2 = obj.serialize() + if obj.__class__.__name__.startswith(("Polar", "Dipole", "DOS")): + common_keys = set(data1.keys()) & set(data2.keys()) + data1 = {k: data1[k] for k in common_keys} + data2 = {k: data2[k] for k in common_keys} + # drop @variables since they are not equal across backends + data1.pop("@variables", None) + data2.pop("@variables", None) + np.testing.assert_equal(data1, data2) + self._compare_ret(ret1, ret2) + + @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.") + def test_tf2_self_consistent(self) -> None: + """Test whether TF2 is self consistent.""" + if self.skip_tf2 or self.tf2_class is None: + self.skipTest("Unsupported backend") + obj1 = self.init_backend_cls(self.tf2_class) + ret1, data1 = self.get_tf2_ret_serialization_from_cls(obj1) + obj2 = self.tf2_class.deserialize(data1) + ret2, data2 = self.get_tf2_ret_serialization_from_cls(obj2) + np.testing.assert_equal(data1, data2) + self._compare_ret( + self.extract_ret(ret1, self.RefBackend.TF2), + self.extract_ret(ret2, self.RefBackend.TF2), + ) + def test_jax_consistent_with_ref(self) -> None: """Test whether JAX and reference are consistent.""" if self.skip_jax: diff --git a/source/tests/consistent/descriptor/common.py b/source/tests/consistent/descriptor/common.py index 078db4829f..1fe394b768 100644 --- a/source/tests/consistent/descriptor/common.py +++ b/source/tests/consistent/descriptor/common.py @@ -26,6 +26,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, ) if INSTALLED_PT or INSTALLED_PT_EXPT: @@ -41,6 +42,10 @@ GLOBAL_TF_FLOAT_PRECISION, tf, ) +if INSTALLED_TF2: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) if INSTALLED_JAX: from deepmd.jax.env import ( jnp, @@ -209,6 +214,47 @@ def eval_pt_expt_descriptor( for x in pt_expt_obj(ext_coords, ext_atype, **kwargs) ] + def eval_tf2_descriptor( + self, + tf2_obj: Any, + natoms: np.ndarray, + coords: np.ndarray, + atype: np.ndarray, + box: np.ndarray, + mixed_types: bool = False, + fparam: np.ndarray | None = None, + charge_spin: np.ndarray | None = None, + ) -> Any: + ext_coords, ext_atype, mapping = extend_coord_with_ghosts( + coords.reshape(1, -1, 3), + atype.reshape(1, -1), + box.reshape(1, 3, 3), + tf2_obj.get_rcut(), + ) + nlist = build_neighbor_list( + ext_coords, + ext_atype, + natoms[0], + tf2_obj.get_rcut(), + tf2_obj.get_sel(), + distinguish_types=(not mixed_types), + ) + kwargs = { + "nlist": to_tensorflow_array(nlist), + "mapping": to_tensorflow_array(mapping), + "fparam": to_tensorflow_array(fparam), + } + if hasattr(tf2_obj, "get_dim_chg_spin") and tf2_obj.get_dim_chg_spin() > 0: + kwargs["charge_spin"] = to_tensorflow_array(charge_spin) + return [ + to_numpy_array(x) + for x in tf2_obj( + to_tensorflow_array(ext_coords), + to_tensorflow_array(ext_atype), + **kwargs, + ) + ] + def eval_jax_descriptor( self, jax_obj: Any, diff --git a/source/tests/consistent/descriptor/test_dpa1.py b/source/tests/consistent/descriptor/test_dpa1.py index d1e8a97497..a3f7a5820c 100644 --- a/source/tests/consistent/descriptor/test_dpa1.py +++ b/source/tests/consistent/descriptor/test_dpa1.py @@ -21,6 +21,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -37,6 +38,10 @@ from deepmd.tf.descriptor.se_atten import DescrptDPA1Compat as DescrptDPA1TF else: DescrptDPA1TF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.dpa1 import DescrptDPA1 as DescrptDPA1TF2 +else: + DescrptDPA1TF2 = None if INSTALLED_JAX: from deepmd.jax.descriptor.dpa1 import DescrptDPA1 as DescriptorDPA1JAX else: @@ -403,7 +408,36 @@ def skip_tf(self) -> bool: ) ) + @property + def skip_tf2(self) -> bool: + ( + tebd_dim, + tebd_input_mode, + resnet_dt, + type_one_side, + attn, + attn_layer, + attn_dotr, + excluded_types, + env_protection, + set_davg_zero, + scaling_factor, + normalize, + temperature, + ln_eps, + smooth_type_embedding, + concat_output_tebd, + precision, + use_econf_tebd, + use_tebd_bias, + ) = self.param + return not INSTALLED_TF2 or self.is_meaningless_zero_attention_layer_tests( + attn_layer, + temperature, + ) + tf_class = DescrptDPA1TF + tf2_class = DescrptDPA1TF2 dp_class = DescrptDPA1DP pt_class = DescrptDPA1PT pt_expt_class = DescrptDPA1PTExpt @@ -528,6 +562,16 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: mixed_types=True, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + ) + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: return self.eval_array_api_strict_descriptor( array_api_strict_obj, diff --git a/source/tests/consistent/descriptor/test_dpa2.py b/source/tests/consistent/descriptor/test_dpa2.py index fe4540e240..c94b47b879 100644 --- a/source/tests/consistent/descriptor/test_dpa2.py +++ b/source/tests/consistent/descriptor/test_dpa2.py @@ -20,6 +20,7 @@ INSTALLED_PD, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -54,6 +55,10 @@ # not implemented DescrptDPA2TF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.dpa2 import DescrptDPA2 as DescrptDPA2TF2 +else: + DescrptDPA2TF2 = None from deepmd.dpmodel.descriptor.dpa2 import ( RepformerArgs, @@ -399,8 +404,10 @@ def skip_tf(self) -> bool: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pt_expt = not INSTALLED_PT_EXPT + skip_tf2 = not INSTALLED_TF2 tf_class = DescrptDPA2TF + tf2_class = DescrptDPA2TF2 dp_class = DescrptDPA2DP pt_class = DescrptDPA2PT pt_expt_class = DescrptDPA2PTExpt @@ -532,6 +539,16 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: mixed_types=True, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + ) + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: return self.eval_array_api_strict_descriptor( array_api_strict_obj, diff --git a/source/tests/consistent/descriptor/test_dpa3.py b/source/tests/consistent/descriptor/test_dpa3.py index 3f30d59435..f463a39530 100644 --- a/source/tests/consistent/descriptor/test_dpa3.py +++ b/source/tests/consistent/descriptor/test_dpa3.py @@ -23,6 +23,7 @@ INSTALLED_PD, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -57,6 +58,10 @@ # not implemented DescrptDPA3TF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.dpa3 import DescrptDPA3 as DescrptDPA3TF2 +else: + DescrptDPA3TF2 = None from deepmd.dpmodel.descriptor.dpa3 import ( RepFlowArgs, @@ -357,8 +362,10 @@ def skip_tf(self) -> bool: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pt_expt = not INSTALLED_PT_EXPT + skip_tf2 = not INSTALLED_TF2 tf_class = DescrptDPA3TF + tf2_class = DescrptDPA3TF2 dp_class = DescrptDPA3DP pt_class = DescrptDPA3PT pt_expt_class = DescrptDPA3PTExpt @@ -490,6 +497,17 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: charge_spin=self.charge_spin, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + charge_spin=self.charge_spin, + ) + def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: return self.eval_array_api_strict_descriptor( array_api_strict_obj, diff --git a/source/tests/consistent/descriptor/test_hybrid.py b/source/tests/consistent/descriptor/test_hybrid.py index 6557deb9a4..afb430db83 100644 --- a/source/tests/consistent/descriptor/test_hybrid.py +++ b/source/tests/consistent/descriptor/test_hybrid.py @@ -17,6 +17,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, ) from .common import ( @@ -32,6 +33,10 @@ from deepmd.tf.descriptor.hybrid import DescrptHybrid as DescrptHybridTF else: DescrptHybridTF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.hybrid import DescrptHybrid as DescrptHybridTF2 +else: + DescrptHybridTF2 = None if INSTALLED_JAX: from deepmd.jax.descriptor.hybrid import DescrptHybrid as DescrptHybridJAX else: @@ -86,6 +91,7 @@ def data(self) -> dict: } tf_class = DescrptHybridTF + tf2_class = DescrptHybridTF2 dp_class = DescrptHybridDP pt_class = DescrptHybridPT pt_expt_class = DescrptHybridPTExpt @@ -96,6 +102,7 @@ def data(self) -> dict: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pt_expt = not INSTALLED_PT_EXPT + skip_tf2 = not INSTALLED_TF2 def setUp(self) -> None: CommonTest.setUp(self) @@ -177,6 +184,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_descriptor( jax_obj, diff --git a/source/tests/consistent/descriptor/test_se_atten_v2.py b/source/tests/consistent/descriptor/test_se_atten_v2.py index 1cc644c73c..a7d0fa8e52 100644 --- a/source/tests/consistent/descriptor/test_se_atten_v2.py +++ b/source/tests/consistent/descriptor/test_se_atten_v2.py @@ -20,6 +20,7 @@ INSTALLED_PD, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -59,6 +60,12 @@ else: DescrptSeAttenV2PD = None DescrptSeAttenV2TF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.se_atten_v2 import ( + DescrptSeAttenV2 as DescrptSeAttenV2TF2, + ) +else: + DescrptSeAttenV2TF2 = None from deepmd.utils.argcheck import ( descrpt_se_atten_args, ) @@ -322,7 +329,36 @@ def skip_pd(self) -> bool: temperature, ) + @property + def skip_tf2(self) -> bool: + ( + tebd_dim, + resnet_dt, + type_one_side, + attn, + attn_layer, + attn_dotr, + excluded_types, + env_protection, + set_davg_zero, + scaling_factor, + normalize, + temperature, + ln_eps, + concat_output_tebd, + precision, + use_econf_tebd, + use_tebd_bias, + ) = self.param + return not INSTALLED_TF2 or self.is_meaningless_zero_attention_layer_tests( + attn_layer, + attn_dotr, + normalize, + temperature, + ) + tf_class = DescrptSeAttenV2TF + tf2_class = DescrptSeAttenV2TF2 dp_class = DescrptSeAttenV2DP pt_class = DescrptSeAttenV2PT pt_expt_class = DescrptSeAttenV2PTExpt @@ -425,6 +461,16 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: mixed_types=True, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + ) + def eval_pd(self, pd_obj: Any) -> Any: return self.eval_pd_descriptor( pd_obj, diff --git a/source/tests/consistent/descriptor/test_se_e2_a.py b/source/tests/consistent/descriptor/test_se_e2_a.py index 2174c5edda..0e0f4bd20d 100644 --- a/source/tests/consistent/descriptor/test_se_e2_a.py +++ b/source/tests/consistent/descriptor/test_se_e2_a.py @@ -18,6 +18,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized_cases, ) @@ -43,6 +44,10 @@ from deepmd.tf.descriptor.se_a import DescrptSeA as DescrptSeATF else: DescrptSeATF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.se_e2_a import DescrptSeA as DescrptSeATF2 +else: + DescrptSeATF2 = None if INSTALLED_PD: import paddle @@ -189,6 +194,10 @@ def skip_tf(self) -> bool: ) = self.param return env_protection != 0.0 or CommonTest.skip_tf + @property + def skip_tf2(self) -> bool: + return not INSTALLED_TF2 + @property def skip_jax(self) -> bool: ( @@ -223,6 +232,7 @@ def skip_array_api_strict(self) -> bool: return not type_one_side or not INSTALLED_ARRAY_API_STRICT tf_class = DescrptSeATF + tf2_class = DescrptSeATF2 dp_class = DescrptSeADP pt_class = DescrptSeAPT pt_expt_class = DescrptSeAPTExpt @@ -314,6 +324,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_descriptor( jax_obj, @@ -451,6 +470,10 @@ def skip_dp(self) -> bool: def skip_tf(self) -> bool: return True + @property + def skip_tf2(self) -> bool: + return not INSTALLED_TF2 + @property def skip_jax(self) -> bool: ( @@ -485,6 +508,7 @@ def skip_array_api_strict(self) -> bool: return not type_one_side or not INSTALLED_ARRAY_API_STRICT tf_class = DescrptSeATF + tf2_class = DescrptSeATF2 dp_class = DescrptSeADP pt_class = DescrptSeAPT pt_expt_class = DescrptSeAPTExpt @@ -613,6 +637,26 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + tf2_obj.compute_input_stats( + [ + { + "r0": None, + "coord": self.coords.reshape(-1, self.natoms[0], 3), + "atype": self.atype.reshape(1, -1), + "box": self.box.reshape(1, 3, 3), + "natoms": self.natoms[0], + } + ] + ) + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: jax_obj.compute_input_stats( [ diff --git a/source/tests/consistent/descriptor/test_se_r.py b/source/tests/consistent/descriptor/test_se_r.py index 826eaf2145..4ceaec57d1 100644 --- a/source/tests/consistent/descriptor/test_se_r.py +++ b/source/tests/consistent/descriptor/test_se_r.py @@ -17,6 +17,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -37,6 +38,10 @@ from deepmd.tf.descriptor.se_r import DescrptSeR as DescrptSeRTF else: DescrptSeRTF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.se_e2_r import DescrptSeR as DescrptSeRTF2 +else: + DescrptSeRTF2 = None from deepmd.utils.argcheck import ( descrpt_se_r_args, ) @@ -131,7 +136,18 @@ def skip_array_api_strict(self) -> bool: ) = self.param return not type_one_side or not INSTALLED_ARRAY_API_STRICT + @property + def skip_tf2(self) -> bool: + ( + resnet_dt, + type_one_side, + excluded_types, + precision, + ) = self.param + return not type_one_side or not INSTALLED_TF2 + tf_class = DescrptSeRTF + tf2_class = DescrptSeRTF2 dp_class = DescrptSeRDP pt_class = DescrptSeRPT pt_expt_class = DescrptSeRPTExpt @@ -210,6 +226,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_descriptor( jax_obj, diff --git a/source/tests/consistent/descriptor/test_se_t.py b/source/tests/consistent/descriptor/test_se_t.py index 7d2a33aba9..5980c04e87 100644 --- a/source/tests/consistent/descriptor/test_se_t.py +++ b/source/tests/consistent/descriptor/test_se_t.py @@ -17,6 +17,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -37,6 +38,10 @@ from deepmd.tf.descriptor.se_t import DescrptSeT as DescrptSeTTF else: DescrptSeTTF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.se_t import DescrptSeT as DescrptSeTTF2 +else: + DescrptSeTTF2 = None if INSTALLED_JAX: from deepmd.jax.descriptor.se_t import DescrptSeT as DescrptSeTJAX else: @@ -120,8 +125,10 @@ def skip_tf(self) -> bool: skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_jax = not INSTALLED_JAX + skip_tf2 = not INSTALLED_TF2 tf_class = DescrptSeTTF + tf2_class = DescrptSeTTF2 dp_class = DescrptSeTDP pt_class = DescrptSeTPT pt_expt_class = DescrptSeTPTExpt @@ -210,6 +217,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_descriptor( jax_obj, diff --git a/source/tests/consistent/descriptor/test_se_t_tebd.py b/source/tests/consistent/descriptor/test_se_t_tebd.py index 4017e059f5..b9d10d53a9 100644 --- a/source/tests/consistent/descriptor/test_se_t_tebd.py +++ b/source/tests/consistent/descriptor/test_se_t_tebd.py @@ -20,6 +20,7 @@ INSTALLED_PD, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -39,6 +40,10 @@ else: DescrptSeTTebdPTExpt = None DescrptSeTTebdTF = None +if INSTALLED_TF2: + from deepmd.tf2.descriptor.se_t_tebd import DescrptSeTTebd as DescrptSeTTebdTF2 +else: + DescrptSeTTebdTF2 = None if INSTALLED_JAX: from deepmd.jax.descriptor.se_t_tebd import DescrptSeTTebd as DescrptSeTTebdJAX else: @@ -180,8 +185,10 @@ def skip_tf(self) -> bool: skip_pd = not INSTALLED_PD skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_tf2 = not INSTALLED_TF2 tf_class = DescrptSeTTebdTF + tf2_class = DescrptSeTTebdTF2 dp_class = DescrptSeTTebdDP pt_class = DescrptSeTTebdPT pt_expt_class = DescrptSeTTebdPTExpt @@ -277,6 +284,16 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: mixed_types=True, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_descriptor( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + mixed_types=True, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_descriptor( jax_obj, diff --git a/source/tests/consistent/fitting/test_dipole.py b/source/tests/consistent/fitting/test_dipole.py index 245744a93e..4d052cac73 100644 --- a/source/tests/consistent/fitting/test_dipole.py +++ b/source/tests/consistent/fitting/test_dipole.py @@ -20,6 +20,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -45,6 +46,13 @@ from deepmd.tf.fit.dipole import DipoleFittingSeA as DipoleFittingTF else: DipoleFittingTF = object +if INSTALLED_TF2: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + from deepmd.tf2.fitting.fitting import DipoleFittingNet as DipoleFittingTF2 +else: + DipoleFittingTF2 = object if INSTALLED_JAX: from deepmd.jax.env import ( jnp, @@ -122,6 +130,7 @@ def skip_pt(self) -> bool: return CommonTest.skip_pt tf_class = DipoleFittingTF + tf2_class = DipoleFittingTF2 dp_class = DipoleFittingDP pt_class = DipoleFittingPT pt_expt_class = DipoleFittingPTExpt @@ -130,6 +139,7 @@ def skip_pt(self) -> bool: args = fitting_dipole() skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_tf2 = not INSTALLED_TF2 @property def skip_pt_expt(self) -> bool: @@ -223,6 +233,16 @@ def eval_dp(self, dp_obj: Any) -> Any: None, )["dipole"] + def eval_tf2(self, tf2_obj: Any) -> Any: + return to_numpy_array( + tf2_obj( + to_tensorflow_array(self.inputs), + to_tensorflow_array(self.atype.reshape(1, -1)), + to_tensorflow_array(self.gr), + None, + )["dipole"] + ) + def eval_jax(self, jax_obj: Any) -> Any: return np.asarray( jax_obj( diff --git a/source/tests/consistent/fitting/test_dos.py b/source/tests/consistent/fitting/test_dos.py index f758c9d317..a1f8530c5d 100644 --- a/source/tests/consistent/fitting/test_dos.py +++ b/source/tests/consistent/fitting/test_dos.py @@ -20,6 +20,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -43,6 +44,13 @@ from deepmd.tf.fit.dos import DOSFitting as DOSFittingTF else: DOSFittingTF = object +if INSTALLED_TF2: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + from deepmd.tf2.fitting.fitting import DOSFittingNet as DOSFittingTF2 +else: + DOSFittingTF2 = object from deepmd.utils.argcheck import ( fitting_dos, ) @@ -116,7 +124,10 @@ def skip_array_api_strict(self) -> bool: def skip_pt_expt(self) -> bool: return CommonTest.skip_pt_expt + skip_tf2 = not INSTALLED_TF2 + tf_class = DOSFittingTF + tf2_class = DOSFittingTF2 dp_class = DOSFittingDP pt_class = DOSFittingPT pt_expt_class = DOSFittingPTExpt @@ -239,6 +250,24 @@ def eval_dp(self, dp_obj: Any) -> Any: aparam=self.aparam if numb_aparam else None, )["dos"] + def eval_tf2(self, tf2_obj: Any) -> Any: + ( + resnet_dt, + precision, + mixed_types, + numb_fparam, + numb_aparam, + numb_dos, + ) = self.param + return to_numpy_array( + tf2_obj( + to_tensorflow_array(self.inputs), + to_tensorflow_array(self.atype.reshape(1, -1)), + fparam=to_tensorflow_array(self.fparam) if numb_fparam else None, + aparam=to_tensorflow_array(self.aparam) if numb_aparam else None, + )["dos"] + ) + def eval_jax(self, jax_obj: Any) -> Any: ( resnet_dt, diff --git a/source/tests/consistent/fitting/test_ener.py b/source/tests/consistent/fitting/test_ener.py index ba0f68c163..1c2caa7027 100644 --- a/source/tests/consistent/fitting/test_ener.py +++ b/source/tests/consistent/fitting/test_ener.py @@ -21,6 +21,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -46,6 +47,13 @@ from deepmd.tf.fit.ener import EnerFitting as EnerFittingTF else: EnerFittingTF = object +if INSTALLED_TF2: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + from deepmd.tf2.fitting.fitting import EnergyFittingNet as EnerFittingTF2 +else: + EnerFittingTF2 = None if INSTALLED_PD: import paddle @@ -159,6 +167,24 @@ def skip_tf(self) -> bool: ) = self.param return not INSTALLED_TF or default_fparam is not None + @property + def skip_tf2(self) -> bool: + ( + resnet_dt, + precision, + mixed_types, + (numb_fparam, default_fparam), + (numb_aparam, use_aparam_as_mask), + atom_ener, + ) = self.param + return ( + not INSTALLED_TF2 + or precision == "bfloat16" + or default_fparam is not None + or use_aparam_as_mask + or atom_ener != [] + ) + @property def skip_pt_expt(self) -> bool: ( @@ -173,6 +199,7 @@ def skip_pt_expt(self) -> bool: return CommonTest.skip_pt_expt or precision == "bfloat16" tf_class = EnerFittingTF + tf2_class = EnerFittingTF2 dp_class = EnerFittingDP pt_class = EnerFittingPT pt_expt_class = EnerFittingPTExpt @@ -304,6 +331,26 @@ def eval_dp(self, dp_obj: Any) -> Any: aparam=self.aparam if numb_aparam else None, )["energy"] + def eval_tf2(self, tf2_obj: Any) -> Any: + ( + resnet_dt, + precision, + mixed_types, + (numb_fparam, default_fparam), + (numb_aparam, use_aparam_as_mask), + atom_ener, + ) = self.param + return to_numpy_array( + tf2_obj( + to_tensorflow_array(self.inputs), + to_tensorflow_array(self.atype.reshape(1, -1)), + fparam=to_tensorflow_array(self.fparam) + if (numb_fparam and default_fparam is None) + else None, + aparam=to_tensorflow_array(self.aparam) if numb_aparam else None, + )["energy"] + ) + def eval_jax(self, jax_obj: Any) -> Any: ( resnet_dt, @@ -473,7 +520,26 @@ def skip_array_api_strict(self) -> bool: def skip_pd(self) -> bool: return not INSTALLED_PD + @property + def skip_tf2(self) -> bool: + ( + resnet_dt, + precision, + mixed_types, + (numb_fparam, default_fparam), + (numb_aparam, use_aparam_as_mask), + atom_ener, + ) = self.param + return ( + not INSTALLED_TF2 + or precision == "bfloat16" + or default_fparam is not None + or use_aparam_as_mask + or atom_ener != [] + ) + tf_class = EnerFittingTF + tf2_class = EnerFittingTF2 dp_class = EnerFittingDP pt_class = EnerFittingPT pt_expt_class = EnerFittingPTExpt @@ -661,6 +727,38 @@ def eval_dp(self, dp_obj: Any) -> Any: aparam=self.aparam, )["energy"] + def eval_tf2(self, tf2_obj: Any) -> Any: + ( + resnet_dt, + precision, + mixed_types, + (numb_fparam, default_fparam), + (numb_aparam, use_aparam_as_mask), + atom_ener, + ) = self.param + tf2_stat_data = [ + { + "fparam": to_tensorflow_array(d["fparam"]), + "aparam": to_tensorflow_array(d["aparam"]), + "find_fparam": d["find_fparam"], + "find_aparam": d["find_aparam"], + } + for d in self.stat_data + ] + tf2_obj.compute_input_stats(tf2_stat_data, protection=1e-2) + return to_numpy_array( + tf2_obj( + to_tensorflow_array(self.inputs), + to_tensorflow_array(self.atype.reshape(1, -1)), + fparam=to_tensorflow_array(self.fparam) + if self.fparam is not None + else None, + aparam=to_tensorflow_array(self.aparam) + if self.aparam is not None + else None, + )["energy"] + ) + def eval_jax(self, jax_obj: Any) -> Any: ( resnet_dt, diff --git a/source/tests/consistent/fitting/test_polar.py b/source/tests/consistent/fitting/test_polar.py index 142cbefdc8..80bf369848 100644 --- a/source/tests/consistent/fitting/test_polar.py +++ b/source/tests/consistent/fitting/test_polar.py @@ -20,6 +20,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -45,6 +46,13 @@ from deepmd.tf.fit.polar import PolarFittingSeA as PolarFittingTF else: PolarFittingTF = object +if INSTALLED_TF2: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + from deepmd.tf2.fitting.fitting import PolarFittingNet as PolarFittingTF2 +else: + PolarFittingTF2 = object if INSTALLED_JAX: from deepmd.jax.env import ( jnp, @@ -96,6 +104,7 @@ def skip_pt(self) -> bool: return CommonTest.skip_pt tf_class = PolarFittingTF + tf2_class = PolarFittingTF2 dp_class = PolarFittingDP pt_class = PolarFittingPT pt_expt_class = PolarFittingPTExpt @@ -104,6 +113,7 @@ def skip_pt(self) -> bool: args = fitting_polar() skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_tf2 = not INSTALLED_TF2 @property def skip_pt_expt(self) -> bool: @@ -193,6 +203,16 @@ def eval_dp(self, dp_obj: Any) -> Any: None, )["polarizability"] + def eval_tf2(self, tf2_obj: Any) -> Any: + return to_numpy_array( + tf2_obj( + to_tensorflow_array(self.inputs), + to_tensorflow_array(self.atype.reshape(1, -1)), + to_tensorflow_array(self.gr), + None, + )["polarizability"] + ) + def eval_jax(self, jax_obj: Any) -> Any: return np.asarray( jax_obj( diff --git a/source/tests/consistent/fitting/test_property.py b/source/tests/consistent/fitting/test_property.py index a9da348410..02e49ede4b 100644 --- a/source/tests/consistent/fitting/test_property.py +++ b/source/tests/consistent/fitting/test_property.py @@ -24,6 +24,7 @@ INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -62,6 +63,13 @@ PropertyFittingStrict = object PropertyFittingTF = object +if INSTALLED_TF2: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + from deepmd.tf2.fitting.fitting import PropertyFittingNet as PropertyFittingTF2 +else: + PropertyFittingTF2 = object @parameterized( @@ -117,12 +125,14 @@ def skip_tf(self) -> bool: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_tf2 = not INSTALLED_TF2 @property def skip_pt_expt(self) -> bool: return CommonTest.skip_pt_expt tf_class = PropertyFittingTF + tf2_class = PropertyFittingTF2 dp_class = PropertyFittingDP pt_class = PropertyFittingPT pt_expt_class = PropertyFittingPTExpt @@ -250,6 +260,25 @@ def eval_dp(self, dp_obj: Any) -> Any: aparam=self.aparam if numb_aparam else None, )[dp_obj.var_name] + def eval_tf2(self, tf2_obj: Any) -> Any: + ( + resnet_dt, + precision, + mixed_types, + numb_fparam, + numb_aparam, + task_dim, + intensive, + ) = self.param + return to_numpy_array( + tf2_obj( + to_tensorflow_array(self.inputs), + to_tensorflow_array(self.atype.reshape(1, -1)), + fparam=to_tensorflow_array(self.fparam) if numb_fparam else None, + aparam=to_tensorflow_array(self.aparam) if numb_aparam else None, + )[tf2_obj.var_name] + ) + def eval_jax(self, jax_obj: Any) -> Any: ( resnet_dt, diff --git a/source/tests/consistent/io/test_io.py b/source/tests/consistent/io/test_io.py index 982d56d8fa..ff51883375 100644 --- a/source/tests/consistent/io/test_io.py +++ b/source/tests/consistent/io/test_io.py @@ -146,6 +146,7 @@ def test_deep_eval(self) -> None: for backend_name, suffix_idx in ( # unfortunately, jax2tf cannot work with tf v1 behaviors ("jax", 2) if DP_TEST_TF2_ONLY else ("tensorflow", 0), + ("tf2", 0) if DP_TEST_TF2_ONLY else (None, None), ("pytorch", 0), ("dpmodel", 0), ("jax", 0) if DP_TEST_TF2_ONLY else (None, None), diff --git a/source/tests/consistent/loss/common.py b/source/tests/consistent/loss/common.py index efe4a33968..abd69b6938 100644 --- a/source/tests/consistent/loss/common.py +++ b/source/tests/consistent/loss/common.py @@ -1,5 +1,48 @@ # SPDX-License-Identifier: LGPL-3.0-or-later +from typing import ( + Any, +) + +import numpy as np + +from deepmd.dpmodel.common import ( + to_numpy_array, +) + + +def _to_tf2_loss_data(value: Any) -> Any: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + + if isinstance(value, dict): + return {kk: _to_tf2_loss_data(vv) for kk, vv in value.items()} + if isinstance(value, tuple): + return tuple(_to_tf2_loss_data(vv) for vv in value) + if isinstance(value, list): + return [_to_tf2_loss_data(vv) for vv in value] + if isinstance(value, np.ndarray): + return to_tensorflow_array(value) + return value class LossTest: """Useful utilities for loss tests.""" + + def eval_tf2_loss( + self, + tf2_obj: Any, + predict: dict[str, Any], + label: dict[str, Any], + **kwargs: Any, + ) -> tuple[Any, dict[str, Any]]: + loss, more_loss = tf2_obj( + self.learning_rate, + self.natoms, + _to_tf2_loss_data(predict), + _to_tf2_loss_data(label), + **kwargs, + ) + loss = to_numpy_array(loss) + more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} + return loss, more_loss diff --git a/source/tests/consistent/loss/test_dos.py b/source/tests/consistent/loss/test_dos.py index 8ed91873ec..a920f07728 100644 --- a/source/tests/consistent/loss/test_dos.py +++ b/source/tests/consistent/loss/test_dos.py @@ -19,6 +19,7 @@ INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -69,8 +70,10 @@ def data(self) -> dict: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pd = True + skip_tf2 = not INSTALLED_TF2 dp_class = DOSLossDP + tf2_class = DOSLossDP pt_class = DOSLossPT pt_expt_class = DOSLossPTExpt jax_class = DOSLossDP @@ -169,6 +172,9 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} return loss, more_loss + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_loss(tf2_obj, self.predict, self.label) + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: loss = ret[0] result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} diff --git a/source/tests/consistent/loss/test_ener.py b/source/tests/consistent/loss/test_ener.py index c87e7409cb..c167a9f43d 100644 --- a/source/tests/consistent/loss/test_ener.py +++ b/source/tests/consistent/loss/test_ener.py @@ -21,6 +21,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -128,8 +129,10 @@ def skip_pd(self) -> bool: skip_pt_expt = not INSTALLED_PT_EXPT skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_tf2 = not INSTALLED_TF2 tf_class = EnerLossTF + tf2_class = EnerLossDP dp_class = EnerLossDP pt_class = EnerLossPT pt_expt_class = EnerLossPTExpt @@ -333,6 +336,14 @@ def eval_pd(self, pd_obj: Any) -> Any: more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} return loss, more_loss + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_loss( + tf2_obj, + self.predict_dpmodel_style, + self.label, + mae=self.mae, + ) + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: loss = ret[0] result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} @@ -386,8 +397,10 @@ def data(self) -> dict: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pd = not INSTALLED_PD + skip_tf2 = not INSTALLED_TF2 tf_class = EnerLossTF + tf2_class = EnerLossDP dp_class = EnerLossDP pt_class = EnerLossPT pt_expt_class = EnerLossPTExpt @@ -568,6 +581,14 @@ def eval_pd(self, pd_obj: Any) -> Any: more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} return loss, more_loss + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_loss( + tf2_obj, + self.predict_dpmodel_style, + self.label, + mae=True, + ) + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: loss = ret[0] result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} diff --git a/source/tests/consistent/loss/test_ener_spin.py b/source/tests/consistent/loss/test_ener_spin.py index bd6561bb78..dacb354e68 100644 --- a/source/tests/consistent/loss/test_ener_spin.py +++ b/source/tests/consistent/loss/test_ener_spin.py @@ -19,6 +19,7 @@ INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -74,8 +75,10 @@ def data(self) -> dict: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pd = True + skip_tf2 = not INSTALLED_TF2 dp_class = EnerSpinLossDP + tf2_class = EnerSpinLossDP pt_class = EnerSpinLossPT pt_expt_class = EnerSpinLossPTExpt jax_class = EnerSpinLossDP @@ -192,6 +195,14 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} return loss, more_loss + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_loss( + tf2_obj, + self.predict, + self.label, + mae=self.mae, + ) + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: loss = ret[0] result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} diff --git a/source/tests/consistent/loss/test_property.py b/source/tests/consistent/loss/test_property.py index 7750eb6dae..1bb838cc5a 100644 --- a/source/tests/consistent/loss/test_property.py +++ b/source/tests/consistent/loss/test_property.py @@ -19,6 +19,7 @@ INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -61,8 +62,10 @@ def data(self) -> dict: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pd = True + skip_tf2 = not INSTALLED_TF2 dp_class = PropertyLossDP + tf2_class = PropertyLossDP pt_class = PropertyLossPT pt_expt_class = PropertyLossPTExpt jax_class = PropertyLossDP @@ -157,6 +160,9 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} return loss, more_loss + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_loss(tf2_obj, self.predict, self.label) + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: loss = ret[0] result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} diff --git a/source/tests/consistent/loss/test_tensor.py b/source/tests/consistent/loss/test_tensor.py index 06feb908fa..1c41fe425c 100644 --- a/source/tests/consistent/loss/test_tensor.py +++ b/source/tests/consistent/loss/test_tensor.py @@ -19,6 +19,7 @@ INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -63,8 +64,10 @@ def data(self) -> dict: skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT skip_pd = True + skip_tf2 = not INSTALLED_TF2 dp_class = TensorLossDP + tf2_class = TensorLossDP pt_class = TensorLossPT pt_expt_class = TensorLossPTExpt jax_class = TensorLossDP @@ -168,6 +171,9 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: more_loss = {kk: to_numpy_array(vv) for kk, vv in more_loss.items()} return loss, more_loss + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_loss(tf2_obj, self.predict, self.label) + def extract_ret(self, ret: Any, backend) -> dict[str, np.ndarray]: loss = ret[0] result = {"loss": np.atleast_1d(np.asarray(loss, dtype=np.float64))} diff --git a/source/tests/consistent/model/common.py b/source/tests/consistent/model/common.py index 3dff24dcba..0c2c4f98fe 100644 --- a/source/tests/consistent/model/common.py +++ b/source/tests/consistent/model/common.py @@ -18,6 +18,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, ) if INSTALLED_PT: @@ -28,6 +29,8 @@ GLOBAL_TF_FLOAT_PRECISION, tf, ) +if INSTALLED_TF2: + import tensorflow as tf if INSTALLED_JAX: from deepmd.jax.common import to_jax_array as numpy_to_jax from deepmd.jax.env import ( @@ -122,6 +125,18 @@ def eval_pt_expt_model(self, pt_expt_obj: Any, natoms, coords, atype, box) -> An ).items() } + def eval_tf2_model(self, tf2_obj: Any, natoms, coords, atype, box) -> Any: + del natoms + return { + kk: vv.numpy() if isinstance(vv, tf.Tensor) else to_numpy_array(vv) + for kk, vv in tf2_obj( + coords, + atype, + box=box, + do_atomic_virial=True, + ).items() + } + def eval_jax_model(self, jax_obj: Any, natoms, coords, atype, box) -> Any: def assert_jax_array(arr): assert isinstance(arr, jnp.ndarray) or arr is None diff --git a/source/tests/consistent/model/test_dipole.py b/source/tests/consistent/model/test_dipole.py index 8b9c24cd57..bfa4bb2411 100644 --- a/source/tests/consistent/model/test_dipole.py +++ b/source/tests/consistent/model/test_dipole.py @@ -27,6 +27,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -46,6 +47,11 @@ from deepmd.tf.model.tensor import DipoleModel as DipoleModelTF else: DipoleModelTF = None +if INSTALLED_TF2: + from deepmd.tf2.model.dipole_model import DipoleModel as DipoleModelTF2 + from deepmd.tf2.model.model import get_model as get_model_tf2 +else: + DipoleModelTF2 = None if INSTALLED_JAX: from deepmd.jax.model.dipole_model import DipoleModel as DipoleModelJAX from deepmd.jax.model.model import get_model as get_model_jax @@ -89,6 +95,7 @@ def data(self) -> dict: } tf_class = DipoleModelTF + tf2_class = DipoleModelTF2 dp_class = DipoleModelDP pt_class = DipoleModelPT pt_expt_class = DipoleModelPTExpt @@ -115,6 +122,8 @@ def get_reference_backend(self): def skip_tf(self): return not INSTALLED_TF + skip_tf2 = not INSTALLED_TF2 + @property def skip_jax(self) -> bool: return not INSTALLED_JAX @@ -131,6 +140,8 @@ def pass_data_to_cls(self, cls, data) -> Any: elif cls is DipoleModelPTExpt: dp_model = get_model_dp(data) return DipoleModelPTExpt.deserialize(dp_model.serialize()) + elif cls is DipoleModelTF2: + return get_model_tf2(data) elif cls is DipoleModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -212,6 +223,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_model( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -232,6 +252,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: self.RefBackend.DP, self.RefBackend.PT, self.RefBackend.PT_EXPT, + self.RefBackend.TF2, self.RefBackend.JAX, }: return ( diff --git a/source/tests/consistent/model/test_dos.py b/source/tests/consistent/model/test_dos.py index 016b4ffc04..dfbceeaf4c 100644 --- a/source/tests/consistent/model/test_dos.py +++ b/source/tests/consistent/model/test_dos.py @@ -27,6 +27,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -46,6 +47,11 @@ from deepmd.tf.model.dos import DOSModel as DOSModelTF else: DOSModelTF = None +if INSTALLED_TF2: + from deepmd.tf2.model.dos_model import DOSModel as DOSModelTF2 + from deepmd.tf2.model.model import get_model as get_model_tf2 +else: + DOSModelTF2 = None if INSTALLED_JAX: from deepmd.jax.model.dos_model import DOSModel as DOSModelJAX from deepmd.jax.model.model import get_model as get_model_jax @@ -90,6 +96,7 @@ def data(self) -> dict: } tf_class = DOSModelTF + tf2_class = DOSModelTF2 dp_class = DOSModelDP pt_class = DOSModelPT pt_expt_class = DOSModelPTExpt @@ -115,6 +122,8 @@ def get_reference_backend(self): def skip_tf(self): return not INSTALLED_TF + skip_tf2 = not INSTALLED_TF2 + @property def skip_jax(self) -> bool: return not INSTALLED_JAX @@ -131,6 +140,8 @@ def pass_data_to_cls(self, cls, data) -> Any: elif cls is DOSModelPTExpt: dp_model = get_model_dp(data) return DOSModelPTExpt.deserialize(dp_model.serialize()) + elif cls is DOSModelTF2: + return get_model_tf2(data) elif cls is DOSModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -206,6 +217,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_model( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -226,6 +246,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: self.RefBackend.DP, self.RefBackend.PT, self.RefBackend.PT_EXPT, + self.RefBackend.TF2, self.RefBackend.JAX, }: return ( diff --git a/source/tests/consistent/model/test_ener.py b/source/tests/consistent/model/test_ener.py index d62f84bea8..33e2c0c27b 100644 --- a/source/tests/consistent/model/test_ener.py +++ b/source/tests/consistent/model/test_ener.py @@ -28,6 +28,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, SKIP_FLAG, CommonTest, parameterized, @@ -48,6 +49,11 @@ from deepmd.tf.model.ener import EnerModel as EnergyModelTF else: EnergyModelTF = None +if INSTALLED_TF2: + from deepmd.tf2.model.ener_model import EnergyModel as EnergyModelTF2 + from deepmd.tf2.model.model import get_model as get_model_tf2 +else: + EnergyModelTF2 = None if INSTALLED_PD: from deepmd.pd.model.model import get_model as get_model_pd from deepmd.pd.model.model.ener_model import EnergyModel as EnergyModelPD @@ -119,6 +125,7 @@ def data(self) -> dict: } tf_class = EnergyModelTF + tf2_class = EnergyModelTF2 dp_class = EnergyModelDP pt_class = EnergyModelPT pd_class = EnergyModelPD @@ -152,6 +159,13 @@ def skip_tf(self): or self.data["atom_exclude_types"] != [] ) + @property + def skip_tf2(self) -> bool: + return not INSTALLED_TF2 or ( + self.data["pair_exclude_types"] != [] + or self.data["atom_exclude_types"] != [] + ) + @property def skip_jax(self) -> bool: return not INSTALLED_JAX @@ -168,6 +182,8 @@ def pass_data_to_cls(self, cls, data) -> Any: elif cls is EnergyModelPTExpt: dp_model = get_model_dp(data) return EnergyModelPTExpt.deserialize(dp_model.serialize()) + elif cls is EnergyModelTF2: + return get_model_tf2(data) elif cls is EnergyModelJAX: return get_model_jax(data) elif cls is EnergyModelPD: @@ -250,6 +266,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_model( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -290,6 +315,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: self.RefBackend.PT, self.RefBackend.PT_EXPT, self.RefBackend.JAX, + self.RefBackend.TF2, self.RefBackend.PD, }: return ( @@ -347,6 +373,7 @@ def data(self) -> dict: } tf_class = EnergyModelTF + tf2_class = EnergyModelTF2 dp_class = EnergyModelDP pt_class = EnergyModelPT pt_expt_class = EnergyModelPTExpt @@ -376,6 +403,13 @@ def skip_tf(self) -> bool: # TF does not have lower interface return True + @property + def skip_tf2(self) -> bool: + return not INSTALLED_TF2 or ( + self.data["pair_exclude_types"] != [] + or self.data["atom_exclude_types"] != [] + ) + @property def skip_jax(self) -> bool: return not INSTALLED_JAX @@ -390,6 +424,8 @@ def pass_data_to_cls(self, cls, data) -> Any: elif cls is EnergyModelPTExpt: dp_model = get_model_dp(data) return EnergyModelPTExpt.deserialize(dp_model.serialize()) + elif cls is EnergyModelTF2: + return get_model_tf2(data) elif cls is EnergyModelJAX: return get_model_jax(data) elif cls is EnergyModelPD: @@ -490,6 +526,18 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: ).items() } + def eval_tf2(self, tf2_obj: Any) -> Any: + return { + kk: to_numpy_array(vv) + for kk, vv in tf2_obj.call_lower( + self.extended_coord, + self.extended_atype, + self.nlist, + self.mapping, + do_atomic_virial=True, + ).items() + } + def eval_jax(self, jax_obj: Any) -> Any: return { kk: to_numpy_array(vv) @@ -535,6 +583,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: elif backend in { self.RefBackend.PT, self.RefBackend.JAX, + self.RefBackend.TF2, self.RefBackend.PD, }: return ( diff --git a/source/tests/consistent/model/test_polar.py b/source/tests/consistent/model/test_polar.py index 4fe3a2c6df..b4eac70735 100644 --- a/source/tests/consistent/model/test_polar.py +++ b/source/tests/consistent/model/test_polar.py @@ -27,6 +27,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -46,6 +47,11 @@ from deepmd.tf.model.tensor import PolarModel as PolarModelTF else: PolarModelTF = None +if INSTALLED_TF2: + from deepmd.tf2.model.model import get_model as get_model_tf2 + from deepmd.tf2.model.polar_model import PolarModel as PolarModelTF2 +else: + PolarModelTF2 = None if INSTALLED_JAX: from deepmd.jax.model.model import get_model as get_model_jax from deepmd.jax.model.polar_model import PolarModel as PolarModelJAX @@ -89,6 +95,7 @@ def data(self) -> dict: } tf_class = PolarModelTF + tf2_class = PolarModelTF2 dp_class = PolarModelDP pt_class = PolarModelPT pt_expt_class = PolarModelPTExpt @@ -115,6 +122,8 @@ def get_reference_backend(self): def skip_tf(self): return not INSTALLED_TF + skip_tf2 = not INSTALLED_TF2 + @property def skip_jax(self) -> bool: return not INSTALLED_JAX @@ -131,6 +140,8 @@ def pass_data_to_cls(self, cls, data) -> Any: elif cls is PolarModelPTExpt: dp_model = get_model_dp(data) return PolarModelPTExpt.deserialize(dp_model.serialize()) + elif cls is PolarModelTF2: + return get_model_tf2(data) elif cls is PolarModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -206,6 +217,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_model( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -226,6 +246,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: self.RefBackend.DP, self.RefBackend.PT, self.RefBackend.PT_EXPT, + self.RefBackend.TF2, self.RefBackend.JAX, }: return ( diff --git a/source/tests/consistent/model/test_property.py b/source/tests/consistent/model/test_property.py index ea68c1e838..28affedc11 100644 --- a/source/tests/consistent/model/test_property.py +++ b/source/tests/consistent/model/test_property.py @@ -26,6 +26,7 @@ INSTALLED_JAX, INSTALLED_PT, INSTALLED_PT_EXPT, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -46,6 +47,11 @@ from deepmd.jax.model.property_model import PropertyModel as PropertyModelJAX else: PropertyModelJAX = None +if INSTALLED_TF2: + from deepmd.tf2.model.model import get_model as get_model_tf2 + from deepmd.tf2.model.property_model import PropertyModel as PropertyModelTF2 +else: + PropertyModelTF2 = None if INSTALLED_PT_EXPT: from deepmd.pt_expt.common import to_torch_array as pt_expt_numpy_to_torch from deepmd.pt_expt.model import PropertyModel as PropertyModelPTExpt @@ -85,6 +91,7 @@ def data(self) -> dict: } tf_class = None + tf2_class = PropertyModelTF2 dp_class = PropertyModelDP pt_class = PropertyModelPT pt_expt_class = PropertyModelPTExpt @@ -110,6 +117,8 @@ def get_reference_backend(self): def skip_tf(self): return True # need to fix tf consistency + skip_tf2 = not INSTALLED_TF2 + @property def skip_jax(self) -> bool: return not INSTALLED_JAX @@ -126,6 +135,8 @@ def pass_data_to_cls(self, cls, data) -> Any: elif cls is PropertyModelPTExpt: dp_model = get_model_dp(data) return PropertyModelPTExpt.deserialize(dp_model.serialize()) + elif cls is PropertyModelTF2: + return get_model_tf2(data) elif cls is PropertyModelJAX: return get_model_jax(data) return cls(**data, **self.additional_data) @@ -207,6 +218,15 @@ def eval_pt_expt(self, pt_expt_obj: Any) -> Any: self.box, ) + def eval_tf2(self, tf2_obj: Any) -> Any: + return self.eval_tf2_model( + tf2_obj, + self.natoms, + self.coords, + self.atype, + self.box, + ) + def eval_jax(self, jax_obj: Any) -> Any: return self.eval_jax_model( jax_obj, @@ -223,6 +243,7 @@ def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: self.RefBackend.DP, self.RefBackend.PT, self.RefBackend.PT_EXPT, + self.RefBackend.TF2, self.RefBackend.JAX, }: return ( diff --git a/source/tests/consistent/test_activation.py b/source/tests/consistent/test_activation.py index 803336c15c..21194ce66b 100644 --- a/source/tests/consistent/test_activation.py +++ b/source/tests/consistent/test_activation.py @@ -21,6 +21,7 @@ INSTALLED_PT, INSTALLED_PT_EXPT, INSTALLED_TF, + INSTALLED_TF2, parameterized, ) @@ -99,6 +100,16 @@ def test_jax_consistent_with_ref(self) -> None: self.assertTrue(isinstance(test, jnp.ndarray)) np.testing.assert_allclose(self.ref, np.from_dlpack(test), atol=1e-10) + @unittest.skipUnless(INSTALLED_TF2, "TensorFlow 2 is not installed") + def test_tf2_consistent_with_ref(self) -> None: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + + input = to_tensorflow_array(self.random_input) + test = get_activation_fn_dp(self.activation)(input) + np.testing.assert_allclose(self.ref, to_numpy_array(test), atol=1e-7) + @unittest.skipUnless(INSTALLED_PD, "Paddle is not installed") def test_pd_consistent_with_ref(self): if INSTALLED_PD: diff --git a/source/tests/consistent/test_array_api.py b/source/tests/consistent/test_array_api.py index bec48b198e..ad40d321d5 100644 --- a/source/tests/consistent/test_array_api.py +++ b/source/tests/consistent/test_array_api.py @@ -19,6 +19,7 @@ INSTALLED_ARRAY_API_STRICT, INSTALLED_JAX, INSTALLED_PT, + INSTALLED_TF2, ) if INSTALLED_PT: @@ -36,6 +37,9 @@ if INSTALLED_ARRAY_API_STRICT: import array_api_strict as xp +if INSTALLED_TF2: + from deepmd._vendors import ndtensorflow as tnp + class TestArrayConversion(unittest.TestCase): @unittest.skipUnless(INSTALLED_PT, "PyTorch is not installed") @@ -342,3 +346,27 @@ def test_array_api_strict_consistent_with_ref(self) -> None: values_xp = xp.asarray(self.values_np) result = xp_setitem_at(x_xp, mask_xp, values_xp) np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10) + + @unittest.skipUnless(INSTALLED_TF2, "TensorFlow 2 is not installed") + def test_tf2_consistent_with_ref(self) -> None: + x_tf2 = tnp.asarray(self.x_np) + mask_tf2 = tnp.asarray(self.mask_np) + values_tf2 = tnp.asarray(self.values_np) + result = xp_setitem_at(x_tf2, mask_tf2, values_tf2) + np.testing.assert_allclose(self.ref, to_numpy_array(result), atol=1e-10) + + @unittest.skipUnless(INSTALLED_TF2, "TensorFlow 2 is not installed") + def test_tf2_full_rank_mask_consistent_with_ref(self) -> None: + x_np = np.zeros((1, 6, 10), dtype=np.int64) + mask_np = np.zeros((1, 6, 10), dtype=bool) + mask_np[:, :, :5] = True + values_np = np.arange(np.count_nonzero(mask_np), dtype=np.int64) + ref = x_np.copy() + ref[mask_np] = values_np + + result = xp_setitem_at( + tnp.asarray(x_np), + tnp.asarray(mask_np), + tnp.asarray(values_np), + ) + np.testing.assert_allclose(ref, to_numpy_array(result), atol=1e-10) diff --git a/source/tests/consistent/test_learning_rate.py b/source/tests/consistent/test_learning_rate.py index cd862c24ef..56142a675d 100644 --- a/source/tests/consistent/test_learning_rate.py +++ b/source/tests/consistent/test_learning_rate.py @@ -18,6 +18,7 @@ INSTALLED_ARRAY_API_STRICT, INSTALLED_JAX, INSTALLED_PT, + INSTALLED_TF2, parameterized, ) @@ -135,3 +136,15 @@ def test_jax_consistent_with_ref(self) -> None: self.compare_test_with_ref(jnp.array(self.step)) if self.warmup_step is not None: self.compare_test_with_warmup_ref(jnp.array(self.warmup_step)) + + @unittest.skipUnless(INSTALLED_TF2, "TensorFlow 2 is not installed") + def test_tf2_consistent_with_ref(self) -> None: + from deepmd.tf2.common import ( + to_tensorflow_array, + ) + + self.compare_test_with_ref(to_tensorflow_array(np.asarray(self.step))) + if self.warmup_step is not None: + self.compare_test_with_warmup_ref( + to_tensorflow_array(np.asarray(self.warmup_step)) + ) diff --git a/source/tests/consistent/test_type_embedding.py b/source/tests/consistent/test_type_embedding.py index 9c1de0e8c5..20853c4a7c 100644 --- a/source/tests/consistent/test_type_embedding.py +++ b/source/tests/consistent/test_type_embedding.py @@ -20,6 +20,7 @@ INSTALLED_PD, INSTALLED_PT, INSTALLED_TF, + INSTALLED_TF2, CommonTest, parameterized, ) @@ -53,6 +54,10 @@ from deepmd.pd.utils.env import DEVICE as PD_DEVICE else: TypeEmbedNetPD = object +if INSTALLED_TF2: + from deepmd.tf2.utils.type_embed import TypeEmbedNet as TypeEmbedNetTF2 +else: + TypeEmbedNetTF2 = None @parameterized( @@ -84,6 +89,7 @@ def data(self) -> dict: } tf_class = TypeEmbedNetTF + tf2_class = TypeEmbedNetTF2 dp_class = TypeEmbedNetDP pt_class = TypeEmbedNetPT jax_class = TypeEmbedNetJAX @@ -92,6 +98,7 @@ def data(self) -> dict: args = type_embedding_args() skip_jax = not INSTALLED_JAX skip_array_api_strict = not INSTALLED_ARRAY_API_STRICT + skip_tf2 = not INSTALLED_TF2 @property def additional_data(self) -> dict: @@ -152,6 +159,13 @@ def eval_array_api_strict(self, array_api_strict_obj: Any) -> Any: for x in (out,) ] + def eval_tf2(self, tf2_obj: Any) -> Any: + out = tf2_obj() + return [ + to_numpy_array(x) if hasattr(x, "__array_namespace__") else x + for x in (out,) + ] + def extract_ret(self, ret: Any, backend) -> tuple[np.ndarray, ...]: return (ret[0],) diff --git a/source/tests/infer/convert-models.sh b/source/tests/infer/convert-models.sh index d74023b9fd..f99f8d6478 100755 --- a/source/tests/infer/convert-models.sh +++ b/source/tests/infer/convert-models.sh @@ -4,5 +4,10 @@ set -ev SCRIPT_PATH=$(dirname $(realpath -s $0)) +# .savedmodel is the JAX/JAX2TF output suffix. .savedmodeltf is the TF2 output +# suffix. The C++ API loads both SavedModel artifacts through the TensorFlow C +# API loader historically named DeepPotJAX. dp convert-backend ${SCRIPT_PATH}/deeppot_sea.yaml ${SCRIPT_PATH}/deeppot_sea.savedmodel dp convert-backend ${SCRIPT_PATH}/deeppot_dpa.yaml ${SCRIPT_PATH}/deeppot_dpa.savedmodel +dp convert-backend ${SCRIPT_PATH}/deeppot_sea.yaml ${SCRIPT_PATH}/deeppot_sea.savedmodeltf +dp convert-backend ${SCRIPT_PATH}/deeppot_dpa.yaml ${SCRIPT_PATH}/deeppot_dpa.savedmodeltf