deepmd.backend.jax

Classes

JAXBackend

JAX backend.

Module Contents

class deepmd.backend.jax.JAXBackend

Bases: deepmd.backend.backend.Backend

JAX backend.

name = 'JAX'

The formal name of the backend.

features: ClassVar[deepmd.backend.backend.Backend.Feature]

The features of the backend.

suffixes: ClassVar[list[str]] = ['.hlo', '.jax', '.savedmodel']

The suffixes of the backend.

is_available() → bool

Check if the backend is available.

Returns:

bool

Whether the backend is available.

property entry_point_hook: Callable[[argparse.Namespace], None]

Abstractmethod:

The entry point hook of the backend.

Returns:

Callable[[Namespace], None]

The entry point hook of the backend.

property deep_eval: type[deepmd.infer.deep_eval.DeepEvalBackend]

The Deep Eval backend of the backend.

Returns:

type[DeepEvalBackend]

The Deep Eval backend of the backend.

property neighbor_stat: type[deepmd.jax.utils.neighbor_stat.NeighborStat]

The neighbor statistics of the backend.

Returns:

type[NeighborStat]

The neighbor statistics of the backend.

property serialize_hook: Callable[[str], dict]

The serialize hook to convert the model file to a dictionary.

Returns:

Callable[[str], dict]

The serialize hook of the backend.

property deserialize_hook: Callable[[str, dict], None]

The deserialize hook to convert the dictionary to a model file.

Returns:

Callable[[str, dict], None]

The deserialize hook of the backend.