-
Notifications
You must be signed in to change notification settings - Fork 504
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'devel' into add_paddle_backend
- Loading branch information
Showing
34 changed files
with
524 additions
and
22 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,110 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later | ||
from importlib.util import ( | ||
find_spec, | ||
) | ||
from typing import ( | ||
TYPE_CHECKING, | ||
Callable, | ||
ClassVar, | ||
List, | ||
Type, | ||
) | ||
|
||
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("jax") | ||
class JAXBackend(Backend): | ||
"""JAX backend.""" | ||
|
||
name = "JAX" | ||
"""The formal name of the backend.""" | ||
features: ClassVar[Backend.Feature] = ( | ||
Backend.Feature(0) | ||
# Backend.Feature.ENTRY_POINT | ||
# | Backend.Feature.DEEP_EVAL | ||
# | Backend.Feature.NEIGHBOR_STAT | ||
# | Backend.Feature.IO | ||
) | ||
"""The features of the backend.""" | ||
suffixes: ClassVar[List[str]] = [] | ||
"""The suffixes of the backend.""" | ||
|
||
def is_available(self) -> bool: | ||
"""Check if the backend is available. | ||
Returns | ||
------- | ||
bool | ||
Whether the backend is available. | ||
""" | ||
return find_spec("jax") is not None | ||
|
||
@property | ||
def entry_point_hook(self) -> Callable[["Namespace"], None]: | ||
"""The entry point hook of the backend. | ||
Returns | ||
------- | ||
Callable[[Namespace], None] | ||
The entry point hook of the backend. | ||
""" | ||
raise NotImplementedError | ||
|
||
@property | ||
def deep_eval(self) -> Type["DeepEvalBackend"]: | ||
"""The Deep Eval backend of the backend. | ||
Returns | ||
------- | ||
type[DeepEvalBackend] | ||
The Deep Eval backend of the backend. | ||
""" | ||
raise NotImplementedError | ||
|
||
@property | ||
def neighbor_stat(self) -> Type["NeighborStat"]: | ||
"""The neighbor statistics of the backend. | ||
Returns | ||
------- | ||
type[NeighborStat] | ||
The neighbor statistics of the backend. | ||
""" | ||
raise NotImplementedError | ||
|
||
@property | ||
def serialize_hook(self) -> Callable[[str], dict]: | ||
"""The serialize hook to convert the model file to a dictionary. | ||
Returns | ||
------- | ||
Callable[[str], dict] | ||
The serialize hook of the backend. | ||
""" | ||
raise NotImplementedError | ||
|
||
@property | ||
def deserialize_hook(self) -> 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. | ||
""" | ||
raise NotImplementedError |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later | ||
"""JAX backend.""" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,37 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later | ||
from typing import ( | ||
Union, | ||
overload, | ||
) | ||
|
||
import numpy as np | ||
|
||
from deepmd.jax.env import ( | ||
jnp, | ||
) | ||
|
||
|
||
@overload | ||
def to_jax_array(array: np.ndarray) -> jnp.ndarray: ... | ||
|
||
|
||
@overload | ||
def to_jax_array(array: None) -> None: ... | ||
|
||
|
||
def to_jax_array(array: Union[np.ndarray]) -> Union[jnp.ndarray]: | ||
"""Convert a numpy array to a JAX array. | ||
Parameters | ||
---------- | ||
array : np.ndarray | ||
The numpy array to convert. | ||
Returns | ||
------- | ||
jnp.ndarray | ||
The JAX tensor. | ||
""" | ||
if array is None: | ||
return None | ||
return jnp.array(array) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,14 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later | ||
import os | ||
|
||
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" | ||
|
||
import jax | ||
import jax.numpy as jnp | ||
|
||
jax.config.update("jax_enable_x64", True) | ||
|
||
__all__ = [ | ||
"jax", | ||
"jnp", | ||
] |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1 @@ | ||
# SPDX-License-Identifier: LGPL-3.0-or-later |
Oops, something went wrong.