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,13 @@ def lower_wrapper(*args):
     nonlocal jax_lower_static_kwargs
 
     jaxfn_args = []
-    jaxfn_kwargs = jax_lower_static_kwargs.copy()
+    # TODO(junjiang): revert to jax_lower_static_kwargs.copy() once NumPy 2.0 is
+    # the minimum supported version.
+    jaxfn_kwargs = {
+        k: jax.numpy.array(v) if isinstance(v, float) 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)