warp.jax_experimental.custom_call

Deprecated JAX integration using custom calls.

Deprecated since version 1.10.0: This module is deprecated. Use warp.jax_experimental.ffi instead, which provides the current implementation of JAX integration and supports JAX 0.4.25 and later, including JAX 0.8.0+.

This module contains the legacy custom call-based implementation that is only compatible with JAX 0.4.25 - 0.7.x and is retained for backward compatibility.

API

jax_kernel

Create a Jax primitive from a Warp kernel.