-
Notifications
You must be signed in to change notification settings - Fork 522
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(jax/array-api): se_t_tebd (#4288)
<!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit ## Release Notes - **New Features** - Introduced support for JAX as a backend for the "se_e3_tebd" descriptor, enhancing flexibility in computational options. - Added serialization and deserialization methods to the descriptor classes for better state management. - **Bug Fixes** - Improved handling of attributes in the descriptor classes to ensure correct data types and transformations. - **Tests** - Enhanced the test suite to support multiple backends, including JAX and Array API Strict, improving the robustness of testing. <!-- end of auto-generated comment: release notes by coderabbit.ai --> Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
- Loading branch information
Showing
5 changed files
with
253 additions
and
41 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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,56 @@ | ||
| # SPDX-License-Identifier: LGPL-3.0-or-later | ||
| from typing import ( | ||
| Any, | ||
| ) | ||
|
|
||
| from deepmd.dpmodel.descriptor.se_t_tebd import ( | ||
| DescrptBlockSeTTebd as DescrptBlockSeTTebdDP, | ||
| ) | ||
| from deepmd.dpmodel.descriptor.se_t_tebd import DescrptSeTTebd as DescrptSeTTebdDP | ||
| from deepmd.jax.common import ( | ||
| ArrayAPIVariable, | ||
| flax_module, | ||
| to_jax_array, | ||
| ) | ||
| from deepmd.jax.descriptor.base_descriptor import ( | ||
| BaseDescriptor, | ||
| ) | ||
| from deepmd.jax.utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) | ||
| from deepmd.jax.utils.network import ( | ||
| NetworkCollection, | ||
| ) | ||
| from deepmd.jax.utils.type_embed import ( | ||
| TypeEmbedNet, | ||
| ) | ||
|
|
||
|
|
||
| @flax_module | ||
| class DescrptBlockSeTTebd(DescrptBlockSeTTebdDP): | ||
| def __setattr__(self, name: str, value: Any) -> None: | ||
| if name in {"mean", "stddev"}: | ||
| value = to_jax_array(value) | ||
| if value is not None: | ||
| value = ArrayAPIVariable(value) | ||
| elif name in {"embeddings", "embeddings_strip"}: | ||
| if value is not None: | ||
| value = NetworkCollection.deserialize(value.serialize()) | ||
| elif name == "env_mat": | ||
| # env_mat doesn't store any value | ||
| pass | ||
| elif name == "emask": | ||
| value = PairExcludeMask(value.ntypes, value.exclude_types) | ||
|
|
||
| return super().__setattr__(name, value) | ||
|
|
||
|
|
||
| @BaseDescriptor.register("se_e3_tebd") | ||
| @flax_module | ||
| class DescrptSeTTebd(DescrptSeTTebdDP): | ||
| def __setattr__(self, name: str, value: Any) -> None: | ||
| if name == "se_ttebd": | ||
| value = DescrptBlockSeTTebd.deserialize(value.serialize()) | ||
| elif name == "type_embedding": | ||
| value = TypeEmbedNet.deserialize(value.serialize()) | ||
| return super().__setattr__(name, value) |
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,47 @@ | ||
| # SPDX-License-Identifier: LGPL-3.0-or-later | ||
| from typing import ( | ||
| Any, | ||
| ) | ||
|
|
||
| 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 ( | ||
| to_array_api_strict_array, | ||
| ) | ||
| from ..utils.exclude_mask import ( | ||
| PairExcludeMask, | ||
| ) | ||
| from ..utils.network import ( | ||
| NetworkCollection, | ||
| ) | ||
| from ..utils.type_embed import ( | ||
| TypeEmbedNet, | ||
| ) | ||
|
|
||
|
|
||
| class DescrptBlockSeTTebd(DescrptBlockSeTTebdDP): | ||
| def __setattr__(self, name: str, value: Any) -> None: | ||
| if name in {"mean", "stddev"}: | ||
| value = to_array_api_strict_array(value) | ||
| elif name in {"embeddings", "embeddings_strip"}: | ||
| if value is not None: | ||
| value = NetworkCollection.deserialize(value.serialize()) | ||
| elif name == "env_mat": | ||
| # env_mat doesn't store any value | ||
| pass | ||
| elif name == "emask": | ||
| value = PairExcludeMask(value.ntypes, value.exclude_types) | ||
|
|
||
| return super().__setattr__(name, value) | ||
|
|
||
|
|
||
| class DescrptSeTTebd(DescrptSeTTebdDP): | ||
| def __setattr__(self, name: str, value: Any) -> None: | ||
| if name == "se_ttebd": | ||
| value = DescrptBlockSeTTebd.deserialize(value.serialize()) | ||
| elif name == "type_embedding": | ||
| value = TypeEmbedNet.deserialize(value.serialize()) | ||
| return super().__setattr__(name, value) |
Oops, something went wrong.