fix(jax): setattr case_embd (#5104)

`case_embd` was supported but the JAX backend was not touched.

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

* **Bug Fixes**
* Ensured the case_embd parameter is consistently converted and handled
during fitting, improving compatibility across array backends and
preventing mis-coercion.

* **Tests**
* Adjusted test setup to reset the default computation graph before
enabling eager execution, stabilizing related test runs.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>

Author: njzjz

deepmd/jax/fitting/fitting.py

@@ -43,6 +43,7 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any:
         "fparam_inv_std",
         "aparam_avg",
         "aparam_inv_std",
+        "case_embd",
         "default_fparam_tensor",
     }:
         value = to_jax_array(value)

source/tests/array_api_strict/fitting/fitting.py

@@ -31,6 +31,7 @@ def setattr_for_general_fitting(name: str, value: Any) -> Any:
         "fparam_inv_std",
         "aparam_avg",
         "aparam_inv_std",
+        "case_embd",
         "default_fparam_tensor",
     }:
         value = to_array_api_strict_array(value)

source/tests/pt/test_tabulate.py

@@ -41,6 +41,7 @@ def get_activation_function(functype: int):
 
 
 def setUpModule() -> None:
+    tf.reset_default_graph()
     tf.compat.v1.enable_eager_execution()