Fix JAX bridge incompatibility with NumPy < 2.0

PiperOrigin-RevId: 823564334

Author: junjiang-lab

ai_edge_torch/odml_torch/jax_bridge/_wrap.py

@@ -79,7 +79,15 @@ def lower_wrapper(*args):
     nonlocal jax_lower_static_kwargs
 
     jaxfn_args = []
-    jaxfn_kwargs = jax_lower_static_kwargs.copy()
+    jaxfn_kwargs = {
+        k: (
+            jax.numpy.array(v)
+            if isinstance(v, (int, float)) and not isinstance(v, bool)
+            else v
+        )
+        for k, v in jax_lower_static_kwargs.items()
+    }
+
     for name, arg in zip(jax_lower_argnames, args):
       if name is None:
         jaxfn_args.append(arg)