API Reference¶

Use these pages when you need the exact module surface rather than the workflow guides.

• jarp covers the top-level re-exports.
• jarp.lax documents control-flow wrappers with automatic Python fallbacks.
• jarp.tree covers PyTree partitioning, flattening, and registration helpers.
• jarp.tree.attrs covers the attrs decorators and field specifiers.
• jarp.tree.codegen covers the generated callback helpers behind custom registrations.
• jarp.tree.prelude covers Partial, PyTreeProxy, and the built-in adapter registrations.
• jarp.utils covers the lower-level logging helpers behind eager fallbacks.
• jarp.warp covers array conversion and Warp/JAX adapter wrappers.
• jarp.warp.types covers precision-aware Warp dtype helpers.