style: extend no-explicit-dtype check to xp and jnp (#4247)

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

## Summary by CodeRabbit

- **New Features**
- Expanded the `DPChecker` to recognize additional libraries ("xp" and
"jnp") for enhanced validation of function calls.
  
- **Bug Fixes**
- Improved compatibility of the `offset` calculation in the
`xp_take_along_axis` function to ensure it matches the data type of the
`indices` array.

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

---------

Signed-off-by: Jinzhe Zeng <[email protected]>

Author: njzjz

deepmd/dpmodel/array_api.py

@@ -61,7 +61,7 @@ def xp_take_along_axis(arr, indices, axis):
     else:
         indices = xp.reshape(indices, (0, 0))
 
-    offset = (xp.arange(indices.shape[0]) * m)[:, xp.newaxis]
+    offset = (xp.arange(indices.shape[0], dtype=indices.dtype) * m)[:, xp.newaxis]
     indices = xp.reshape(offset + indices, (-1,))
 
     out = xp.take(arr, indices)

source/checker/deepmd_checker.py

@@ -37,7 +37,7 @@ def visit_call(self, node):
         if (
             isinstance(node.func, Attribute)
             and isinstance(node.func.expr, Name)
-            and node.func.expr.name in {"np", "tf", "torch"}
+            and node.func.expr.name in {"np", "tf", "torch", "xp", "jnp"}
             and node.func.attrname
             in {
                 # https://pytorch.org/docs/stable/torch.html#creation-ops