Import NVIDIA/warp 1.13.0 from GitHub.

PiperOrigin-RevId: 923878216
Change-Id: If9df5e0672e956eb62874fb6f4f74ba3a02977b3

mjx/cuda_requirements.txt

@@ -20,8 +20,8 @@ jax-cuda12-pjrt==0.8.3; python_version >= '3.13' \
 jax-cuda12-pjrt==0.5.3; python_version < '3.13' \
     --hash=sha256:04ee111eaf5fc2692978ad4a5c84d5925e42eb05c1701849ba3a53f6515400cc \
     --hash=sha256:c5378306568ba0c81b230a779dd3194c9dd10339ab6360ae80928108d37e7f75
-warp-lang==1.12.1 \
+warp-lang==1.13.0 \
-    --hash=sha256:98df3533a6c40a33cce961f8efa991006b30c9d286356e4cd77ea8ce86928f1d \
+    --hash=sha256:4375f572301991fe0fbf0af29fc84d76cd27d531432d6df8452b25088c21ea5a \
-    --hash=sha256:6bf01f10509488ba8eacaf4ec7fcf7cfbd503118b22e002ecba407b40a17424e \
+    --hash=sha256:ac2479c70ad410d58deb088c2a64168792be588efc1bec22dc51f92238e8c8c3 \
-    --hash=sha256:af6d680e79c1be6e46ddf80ecaa358f222804f882f4683260a7b4abd80a0981b \
+    --hash=sha256:476b54f0dcf6767f23305a328660a9a74d01e371b6b7c3d77e03e18b4a1bb1a5 \
-    --hash=sha256:826b2f93df8e47eac0c751a8eb5a0533e2fc5434158c8896a63be53bfbd728c7
+    --hash=sha256:47975ea07252d45a4d09d2d1a6cffc55002fa7fbde771c8dceb1baf4b32d4fb8

mjx/mujoco/mjx/_src/io.py

@@ -179,8 +179,8 @@ def _wp_to_np_type(wp_field: Any, name: str = '') -> Any:
     wp_field.shape = wp_field.shape[1:]
   # warp scalars
   wp_dtype = type(wp_field)
-  if wp_dtype in wp.types.warp_type_to_np_dtype:
+  if wp_dtype in wp._src.types.warp_type_to_np_dtype:
-    return wp.types.warp_type_to_np_dtype[wp_dtype](wp_field)
+    return wp.dtype_to_numpy(wp_dtype)(wp_field)
 
   # warp arrays
   if isinstance(wp_field, wp.array):

mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/custom_call.py

@@ -131,13 +131,13 @@ def _warp_custom_callback(stream, buffers, opaque, opaque_len):
 
     # Launch the kernel.
     wp._src.context.runtime.core.wp_cuda_launch_kernel(
-        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream
+        device.context, hooks.forward, bounds.size, 0, 256, hooks.forward_smem_bytes, kernel_params, stream, None
     )
 
 
 def _create_jax_warp_primitive():
     import jax  # noqa: PLC0415
-    from jax._src.interpreters import batching  # noqa: PLC0415
+    from jax.interpreters import batching  # noqa: PLC0415
     from jax.interpreters import mlir  # noqa: PLC0415
     from jax.interpreters.mlir import ir  # noqa: PLC0415
     from jaxlib.hlo_helpers import custom_call  # noqa: PLC0415
@@ -170,9 +170,9 @@ def _create_jax_warp_primitive():
         # Figure out the number of outputs.
         wp_kernel = _registered_kernels[params["kernel"]]
         output_count = len(wp_kernel.adj.args) - len(args)
-        shape, dim = next((a.shape, d) for a, d in zip(args, dims) if d is not None)
+        shape, dim = next((a.shape, d) for a, d in zip(args, dims, strict=True) if d is not None)
         size = shape[dim]
-        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims)]
+        args = [batching.bdim_at_front(a, d, size) if len(a.shape) else a for a, d in zip(args, dims, strict=True)]
         # Create the batched primitive.
         return _jax_warp_p.bind(*args, **params), [dims[0]] * output_count
 
@@ -211,6 +211,7 @@ def _create_jax_warp_primitive():
     def base_type_to_jax_ir(warp_dtype):
         warp_to_jax_dict = {
             wp.float16: ir.F16Type.get(),
+            wp.bfloat16: ir.BF16Type.get(),
             wp.float32: ir.F32Type.get(),
             wp.float64: ir.F64Type.get(),
             wp.int8: ir.IntegerType.get_signless(8),
@@ -232,6 +233,7 @@ def _create_jax_warp_primitive():
     def base_type_is_compatible(warp_type, jax_ir_type):
         jax_ir_to_warp = {
             "f16": wp.float16,
+            "bf16": wp.bfloat16,
             "f32": wp.float32,
             "f64": wp.float64,
             "i8": wp.int8,
@@ -324,7 +326,7 @@ def _create_jax_warp_primitive():
         # Figure out the types and shapes of the input arrays.
         arg_strings = []
         operand_layouts = []
-        for actual, warg in zip(args, wp_kernel.adj.args):
+        for actual, warg in zip(args, wp_kernel.adj.args, strict=False):
             wtype = warg.type
             rtt = ir.RankedTensorType(actual.type)
 

mjx/mujoco/mjx/third_party/warp/_src/jax_experimental/ffi.py

@@ -8,8 +8,8 @@ import ctypes
 import inspect
 import threading
 import traceback
+from collections.abc import Callable
 from enum import IntEnum
-from typing import Callable
 
 import jax
 
@@ -25,6 +25,7 @@ from warp._src.types import (
     type_size_in_bytes,
     type_to_warp,
 )
+from warp._src.utils import warn
 
 from .xla_ffi import *
 
@@ -249,19 +250,18 @@ class FfiKernel:
                 # check dtype
                 if input_value.dtype != input_arg.jax_scalar_type:
                     raise TypeError(
-                            f"Invalid data type for array argument '{input_arg.name}',"
+                        f"Invalid data type for array argument '{input_arg.name}', expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
-                            f" expected {input_arg.jax_scalar_type}, got {input_value.dtype}"
                     )
                 # check ndim
                 if input_value.ndim != input_arg.jax_ndim:
                     raise TypeError(
-                            f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
+                        f"Invalid dimensionality for array argument '{input_arg.name}', expected {input_arg.jax_ndim} dimensions, got {input_value.ndim}"
                     )
                 # check inner dims
                 for d in range(input_arg.dtype_ndim):
                     if input_value.shape[input_arg.type.ndim + d] != input_arg.dtype_shape[d]:
                         raise TypeError(
-                                f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
+                            f"Invalid inner dimensions for array argument '{input_arg.name}', expected {input_arg.dtype_shape}, got {input_value.shape[-input_arg.dtype_ndim :]}"
                         )
             else:
                 # make sure scalar is not a traced variable, should be static
@@ -455,6 +455,7 @@ class FfiKernel:
                         hooks.forward_smem_bytes,
                         kernel_params,
                         stream,
+                        None,  # apic_info
                     )
                 else:
                     wp._src.context.runtime.core.wp_cpu_launch_kernel(
@@ -467,9 +468,7 @@ class FfiKernel:
         except Exception as e:
             print(traceback.format_exc())
             return create_ffi_error(
-                call_frame.contents.api,
+                call_frame.contents.api, XLA_FFI_Error_Code.UNKNOWN, f"FFI callback error: {type(e).__name__}: {e}"
-                XLA_FFI_Error_Code.UNKNOWN,
-                f"FFI callback error: {type(e).__name__}: {e}",
             )
 
 
@@ -617,7 +616,9 @@ class FfiCallable:
         self.input_output_aliases = input_output_aliases
 
         # register the callback
-        FFI_CCALLFUNC = ctypes.CFUNCTYPE(ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame))
+        FFI_CCALLFUNC = ctypes.CFUNCTYPE(
+                ctypes.c_void_p, ctypes.POINTER(XLA_FFI_CallFrame)
+        )
 
         self.callback_func_cuda = FFI_CCALLFUNC(lambda call_frame: self.ffi_callback(call_frame, platform="CUDA"))
         ffi_ccall_address_cuda = ctypes.cast(self.callback_func_cuda, ctypes.c_void_p)
@@ -1071,7 +1072,7 @@ class FfiCallable:
                         call_desc.capture = capture
 
                     else:
-                        # not capturing or on CPU
+                        # not capturing
                         self.func(*arg_list)
 
         except Exception as e:
@@ -1359,7 +1360,7 @@ def jax_kernel(
                 try:
                     gi.zero_()
                 except Exception as e:
-                    wp.utils.warn(f"Failed to zero gradient array: {e}", stacklevel=2)
+                    warn(f"Failed to zero gradient array: {e}", stacklevel=2)
                     raise e
 
         # NOTE: We cannot use a passed launch_dims here, the backward rule doesn't receive it (and it could be wrong under pmap/vmap).
@@ -1445,7 +1446,7 @@ def jax_kernel(
         non_static_inputs, output_vals_tuple = residuals
 
         input_vals = list(non_static_inputs)
-        for i, v in zip(static_args, nondiff_vals):
+        for i, v in zip(static_args, nondiff_vals, strict=True):
             input_vals.insert(i, v)
 
         # Normalize grad outputs and handle nested containers (e.g., single tuple for multi-output)
@@ -1461,7 +1462,7 @@ def jax_kernel(
 
         out_dims_map = {}
         param_ann = {p.name: p.annotation for p in parameters[:num_inputs]}
-        for name, val in zip(differentiable_input_names, non_static_inputs):
+        for name, val in zip(differentiable_input_names, non_static_inputs, strict=True):
             ann = param_ann.get(name)
             if ann is None:
                 continue

mjx/pyproject.toml

@@ -37,7 +37,7 @@ dependencies = [
 
 [project.optional-dependencies]
 warp = [
-    "warp-lang==1.12.1",
+    "warp-lang==1.13.0",
 ]
 dev = [
     "isort",