From 9c767adfba8a56c92501591199db60af11d53f94 Mon Sep 17 00:00:00 2001
From: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
Date: Thu, 31 Oct 2024 09:19:12 -0400
Subject: [PATCH] feat(dpmodel/jax): add entry point for dpmodel and jax
 backend (#4284)

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

## Summary by CodeRabbit

- **New Features**
- Introduced entry point loading functionality for enhanced module
initialization in both `dpmodel` and `jax` components of the DeepMD
framework.

These changes improve the framework's functionality and streamline
backend configuration.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@rutgers.edu>
---
 deepmd/dpmodel/__init__.py | 7 +++++++
 deepmd/jax/__init__.py     | 6 ++++++
 2 files changed, 13 insertions(+)

diff --git a/deepmd/dpmodel/__init__.py b/deepmd/dpmodel/__init__.py
index 6f83f849a3..111c2d6ced 100644
--- a/deepmd/dpmodel/__init__.py
+++ b/deepmd/dpmodel/__init__.py
@@ -1,4 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
+from deepmd.utils.entry_point import (
+    load_entry_point,
+)
+
 from .common import (
     DEFAULT_PRECISION,
     PRECISION_DICT,
@@ -32,3 +36,6 @@
     "get_deriv_name",
     "get_hessian_name",
 ]
+
+
+load_entry_point("deepmd.dpmodel")
diff --git a/deepmd/jax/__init__.py b/deepmd/jax/__init__.py
index 2ff078e797..bb5c0a5206 100644
--- a/deepmd/jax/__init__.py
+++ b/deepmd/jax/__init__.py
@@ -1,2 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 """JAX backend."""
+
+from deepmd.utils.entry_point import (
+    load_entry_point,
+)
+
+load_entry_point("deepmd.jax")