fix(jax): make display_if_exist jit-able (#4766)

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

- **Bug Fixes**
- Enhanced loss display to support element-wise handling of properties,
improving accuracy and consistency with array inputs.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Signed-off-by: Jinzhe Zeng <jinzhe.zeng@ustc.edu.cn>
Signed-off-by: Jinzhe Zeng <njzjz@qq.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: njzjz

deepmd/dpmodel/loss/loss.py

@@ -51,7 +51,9 @@ def display_if_exist(loss: np.ndarray, find_property: float) -> np.ndarray:
             the loss scalar or NaN
         """
         xp = array_api_compat.array_namespace(loss)
-        return loss if bool(find_property) else xp.nan
+        return xp.where(
+            xp.asarray(find_property, dtype=xp.bool), loss, xp.asarray(xp.nan)
+        )
 
     @classmethod
     def get_loss(cls, loss_params: dict) -> "Loss":