diff --git a/.github/workflows/test_python.yml b/.github/workflows/test_python.yml
index b7c43dbcc5..ba8858d6b9 100644
--- a/.github/workflows/test_python.yml
+++ b/.github/workflows/test_python.yml
@@ -58,7 +58,7 @@ jobs:
       env:
         NUM_WORKERS: 0
     - name: Test TF2 eager mode
-      run: pytest --cov=deepmd --cov-append source/tests/consistent/io/test_io.py source/tests/jax/jax2tf --durations=0
+      run: pytest --cov=deepmd --cov-append source/tests/consistent/io/test_io.py source/jax2tf_tests --durations=0
       env:
         NUM_WORKERS: 0
         DP_TEST_TF2_ONLY: 1
diff --git a/pyproject.toml b/pyproject.toml
index b238ed42cf..cf2c8c3b93 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -424,6 +424,7 @@ banned-module-level-imports = [
 "source/tests/pt/**" = ["TID253"]
 "source/tests/jax/**" = ["TID253"]
 "source/tests/universal/pt/**" = ["TID253"]
+"source/jax2tf_tests/**" = ["TID253"]
 "source/ipi/tests/**" = ["TID253"]
 "source/lmp/tests/**" = ["TID253"]
 "**/*.ipynb" = ["T20"]  # printing in a nb file is expected
diff --git a/source/tests/jax/jax2tf/__init__.py b/source/jax2tf_tests/__init__.py
similarity index 100%
rename from source/tests/jax/jax2tf/__init__.py
rename to source/jax2tf_tests/__init__.py
diff --git a/source/tests/jax/jax2tf/test_nlist.py b/source/jax2tf_tests/test_nlist.py
similarity index 100%
rename from source/tests/jax/jax2tf/test_nlist.py
rename to source/jax2tf_tests/test_nlist.py
diff --git a/source/tests/jax/jax2tf/test_region.py b/source/jax2tf_tests/test_region.py
similarity index 100%
rename from source/tests/jax/jax2tf/test_region.py
rename to source/jax2tf_tests/test_region.py