Make JAX-Warp FFI more robust. Prepare for warp-lang 1.10. Fix #2894

PiperOrigin-RevId: 819381906
Change-Id: I1f16abcbe02884445725c3087aa97af941131533

mjx/mujoco/mjx/warp/ffi.py

@@ -110,9 +110,16 @@ def jax_callable_variadic_tuple(
       return func(*unflat_args, **kwargs)
 
     # Provide a flattened signature for the Warp callable machinery.
-    func_wrapper.__signature__ = flatten_signature(
-        inspect.signature(func), args
-    )
+    new_signature = flatten_signature(inspect.signature(func), args)
+    func_wrapper.__signature__ = new_signature
+    func_wrapper.__annotations__ = {
+        p.name: p.annotation
+        for p in new_signature.parameters.values()
+        if p.annotation is not inspect.Parameter.empty
+    }
+    if new_signature.return_annotation is not inspect.Signature.empty:
+      func_wrapper.__annotations__['return'] = new_signature.return_annotation
+
     my_callable = ffi.jax_callable(
         func_wrapper,
         num_outputs=num_outputs,