warp.JaxCallableGraphMode#

class warp.JaxCallableGraphMode(*values)[source]#

CUDA graph capture modes for warp.jax_callable().

These modes control whether JAX or Warp captures a CUDA graph, and whether staging buffers are used when capturing with Warp.

Attributes

`NONE`	Disable graph capture.
`JAX`	Let JAX capture the graph so the callable can be used as a subgraph within a larger JAX capture.
`WARP`	Let Warp capture the graph and replay it for matching buffer addresses.
`WARP_STAGED`	Capture a Warp graph using staging buffers and insert memcpy nodes inside the graph.
`WARP_STAGED_EX`	Capture a Warp graph using staging buffers and perform memcpy outside the graph.

NONE = 0#: Disable graph capture. Use when operations are not CUDA-graph compatible (for example, host synchronization).

JAX = 1#: Let JAX capture the graph so the callable can be used as a subgraph within a larger JAX capture.

WARP = 2#: Let Warp capture the graph and replay it for matching buffer addresses.

WARP_STAGED = 3#: Capture a Warp graph using staging buffers and insert memcpy nodes inside the graph.

WARP_STAGED_EX = 4#: Capture a Warp graph using staging buffers and perform memcpy outside the graph.