From 0d4f6231be43b0222b9cf2b543e6bbd6cd028453 Mon Sep 17 00:00:00 2001 From: Baruch Tabanpour Date: Tue, 14 Oct 2025 13:48:00 -0700 Subject: [PATCH] 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 | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/mjx/mujoco/mjx/warp/ffi.py b/mjx/mujoco/mjx/warp/ffi.py index 0931575a..a04a232b 100644 --- a/mjx/mujoco/mjx/warp/ffi.py +++ b/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,