feat(jax): export call_lower to SavedModel via jax2tf

.github/workflows/test_python.yml

@@ -53,6 +53,12 @@ jobs:
     - run: pytest --cov=deepmd source/tests --durations=0 --splits 6 --group ${{ matrix.group }} --store-durations --durations-path=.test_durations --splitting-algorithm least_duration
       env:
         NUM_WORKERS: 0
+    - name: Test TF2 eager mode
+      run: pytest --cov=deepmd source/tests/consistent/io/test_io.py --durations=0
+      env:
+        NUM_WORKERS: 0
+        DP_TEST_TF2_ONLY: 1
+      if: matrix.group == 1
     - run: mv .test_durations .test_durations_${{ matrix.group }}
     - name: Upload partial durations
       uses: actions/upload-artifact@v4

source/tests/consistent/io/test_io.py

@@ -23,6 +23,7 @@
 
 from ...utils import (
     CI,
+    DP_TEST_TF2_ONLY,
     TEST_DEVICE,
 )
 
@@ -72,6 +73,7 @@ def tearDown(self):
                 shutil.rmtree(ii)
 
     @unittest.skipIf(TEST_DEVICE != "cpu" and CI, "Only test on CPU.")
+    @unittest.skipIf(DP_TEST_TF2_ONLY, "Conflict with TF2 eager mode.")
     def test_data_equal(self):
         prefix = "test_consistent_io_" + self.__class__.__name__.lower()
         for backend_name, suffix_idx in (
@@ -141,12 +143,11 @@ def test_deep_eval(self):
         prefix = "test_consistent_io_" + self.__class__.__name__.lower()
         rets = []
         for backend_name, suffix_idx in (
-            ("tensorflow", 0),
+            # unfortunately, jax2tf cannot work with tf v1 behaviors
+            ("jax", 2) if DP_TEST_TF2_ONLY else ("tensorflow", 0),
             ("pytorch", 0),
             ("dpmodel", 0),
             ("jax", 0),
-            # unfortunately, jax2tf cannot work with tf v1 behaviors
-            # ("jax", 2),
         ):
             backend = Backend.get_backend(backend_name)()
             if not backend.is_available():

source/tests/utils.py

@@ -8,3 +8,4 @@
 
 # see https://docs.github.com/en/actions/writing-workflows/choosing-what-your-workflow-does/store-information-in-variables#default-environment-variables
 CI = os.environ.get("CI") == "true"
+DP_TEST_TF2_ONLY = os.environ.get("DP_TEST_TF2_ONLY") == "1"