deepmd._vendors.ndtensorflow._array

deepmd._vendors.ndtensorflow._array#

Classes#

`DLDeviceType`	Enum where members are also (and must be) ints
`Array`	User-facing TensorFlow-backed array object.

Module Contents#

class deepmd._vendors.ndtensorflow._array.DLDeviceType[source]#

Bases: enum.IntEnum

Enum where members are also (and must be) ints

CPU = 1[source]#

CUDA = 2[source]#

class deepmd._vendors.ndtensorflow._array.Array(tensor: Any | None = None, /)[source]#

Bases: 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[source]#

__composite_gradient__[source]#

_tensor: tf.Tensor[source]#

classmethod _from_tensor(tensor: Any, /) → Array[source]#

_replace_tensor(tensor: tf.Tensor) → None[source]#

unwrap() → tf.Tensor[source]#: Return the wrapped TensorFlow tensor.

__tf_tensor__(dtype: tensorflow.DType | None = None, name: str | None = None) → tf.Tensor[source]#

property device: str[source]#

property dtype: tensorflow.DType[source]#

property mT: Array[source]#

property ndim: int[source]#

property shape: tuple[int | tf.Tensor, Ellipsis][source]#

property size: int | tf.Tensor[source]#

property T: Array[source]#

astype(dtype: tensorflow.DType, /, *, copy: bool = True, device: str | None = None) → Array[source]#

reshape(*shape: Any, copy: bool | None = None) → Array[source]#

ravel() → Array[source]#

squeeze(axis: int | tuple[int, Ellipsis] | None = None) → Array[source]#

to_device(device: str, /, *, stream: int | Any | None = None) → Array[source]#

__dlpack__(*, stream: int | Any | None = None, max_version: tuple[int, int] | None = None, dl_device: tuple[int, int] | None = None, copy: bool | None = None) → Any[source]#

__dlpack_device__() → tuple[DLDeviceType, int][source]#

__array_namespace__(/, *, api_version: str | None = None) → Any[source]#

__array__(dtype: Any | None = None) → Any[source]#

__len__() → int[source]#

__iter__() → collections.abc.Iterator[Array][source]#

__getitem__(key: Any, /) → Array[source]#

__setitem__(key: Any, value: Any, /) → None[source]#

_scalar_value() → Any[source]#

__bool__(/) → bool[source]#

__complex__(/) → complex[source]#

__float__(/) → float[source]#

__index__(/) → int[source]#

__int__(/) → int[source]#

__repr__() → str[source]#

__eq__(other: Any) → Array[source]#

__ne__(other: Any) → Array[source]#

__abs__() → Array[source]#

__invert__() → Array[source]#

__neg__() → Array[source]#

__pos__() → Array[source]#