jarp ¶

Parameters:

• default ¶

  (T, default: ... ) –

• validator ¶

  (_ValidatorArgType[T] | None, default: ... ) –

• repr ¶

  (_ReprArgType, default: ... ) –

• hash ¶

  (bool | None, default: ... ) –

• init ¶

  (bool, default: ... ) –

• metadata ¶

  (Mapping[Any, Any] | None, default: ... ) –

• converter ¶

  (_ConverterType | list[_ConverterType] | tuple[_ConverterType, ...] | None, default: ... ) –

• factory ¶

  (Callable[[], T] | None, default: ... ) –

• kw_only ¶

  (bool | None, default: ... ) –

• eq ¶

  (_EqOrderType | None, default: ... ) –

• order ¶

  (_EqOrderType | None, default: ... ) –

• on_setattr ¶

  (_OnSetAttrArgType | None, default: ... ) –

• alias ¶

  (str | None, default: ... ) –

• type ¶

  (type | None, default: ... ) –

• static ¶

  (FieldType | bool | None, default: ... ) –

63
64
65
66
67
68
69
70

def array(**kwargs: Unpack[FieldOptions[ArrayLike | None]]) -> Array:
    if "default" in kwargs and "factory" not in kwargs:
        default: ArrayLike | None = kwargs["default"]
        if not (default is None or isinstance(default, attrs.Factory)):  # ty:ignore[invalid-argument-type]
            default: Array = jnp.asarray(default)
            kwargs.pop("default")
            kwargs["factory"] = lambda: default
    return field(**kwargs)

Define an attrs class and optionally register it as a PyTree.

Parameters:

• maybe_cls ¶

  (T | None, default: None ) –

  Class being decorated. When omitted, return a configured decorator.

• **kwargs ¶

  (Any, default: {} ) –

  Options forwarded to :func:attrs.define, plus pytree to control JAX registration. pytree="data" registers fields with fieldz semantics, "static" registers the whole instance as a static value, and "none" leaves the class unregistered.

Returns:

• Any –

  The decorated class or a class decorator.

70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

def define[T: type](maybe_cls: T | None = None, **kwargs: Any) -> Any:
    """Define an ``attrs`` class and optionally register it as a PyTree.

    Args:
        maybe_cls: Class being decorated. When omitted, return a configured
            decorator.
        **kwargs: Options forwarded to :func:`attrs.define`, plus ``pytree`` to
            control JAX registration. ``pytree="data"`` registers fields with
            ``fieldz`` semantics, ``"static"`` registers the whole instance as a
            static value, and ``"none"`` leaves the class unregistered.

    Returns:
        The decorated class or a class decorator.
    """
    if maybe_cls is None:
        return functools.partial(define, **kwargs)
    pytree: PyTreeType = PyTreeType(kwargs.pop("pytree", None))
    frozen: bool = kwargs.get("frozen", False)
    if pytree is PyTreeType.STATIC and not frozen:
        warnings.warn(
            "Defining a static class that is not frozen may lead to unexpected behavior.",
            stacklevel=2,
        )
    cls: T = attrs.define(maybe_cls, **kwargs)  # ty:ignore[invalid-assignment]
    match pytree:
        case PyTreeType.DATA:
            register_fieldz(cls)
        case PyTreeType.STATIC:
            jtu.register_static(cls)
    return cls

Define a frozen attrs class and register it as a data PyTree.

116
117
118
119
120
121
122

def frozen[T: type](maybe_cls: T | None = None, **kwargs: Any) -> Any:
    """Define a frozen ``attrs`` class and register it as a data PyTree."""
    _warnings_hide = True
    if maybe_cls is None:
        return functools.partial(frozen, **kwargs)
    kwargs.setdefault("frozen", True)
    return define(maybe_cls, **kwargs)

Define a frozen attrs class and register it as a static PyTree.

139
140
141
142
143
144
145
146

def frozen_static[T: type](maybe_cls: T | None = None, **kwargs: Any) -> Any:
    """Define a frozen ``attrs`` class and register it as a static PyTree."""
    _warnings_hide = True
    if maybe_cls is None:
        return functools.partial(frozen_static, **kwargs)
    kwargs.setdefault("frozen", True)
    kwargs.setdefault("pytree", PyTreeType.STATIC)
    return define(maybe_cls, **kwargs)

Wrap warp.jax_experimental.jax_callable with optional dtype dispatch.

Parameters:

• func ¶

  (Callable | None, default: None ) –

  Warp callable function or factory. When omitted, return a decorator.

• generic ¶

  (bool, default: False ) –

  When true, func is treated as a factory that receives Warp scalar dtypes inferred from the runtime JAX arguments and returns a concrete Warp callable implementation.

• num_outputs ¶

  (int, default: ... ) –

• graph_mode ¶

  (GraphMode, default: ... ) –

• vmap_method ¶

  (str | None, default: ... ) –

• output_dims ¶

  (dict[str, ShapeLike] | None, default: ... ) –

• in_out_argnames ¶

  (Iterable[str], default: ... ) –

• stage_in_argnames ¶

  (Iterable[str], default: ... ) –

• stage_out_argnames ¶

  (Iterable[str], default: ... ) –

• graph_cache_max ¶

  (int | None, default: ... ) –

• module_preload_mode ¶

  (ModulePreloadMode, default: ... ) –

Returns:

• Any –

  A callable compatible with JAX tracing, or a decorator producing one.

 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110

def jax_callable(
    func: Callable | None = None,
    *,
    generic: bool = False,
    **kwargs: Unpack[JaxCallableOptions],
) -> Any:
    """Wrap ``warp.jax_experimental.jax_callable`` with optional dtype dispatch.

    Args:
        func: Warp callable function or factory. When omitted, return a
            decorator.
        generic: When true, ``func`` is treated as a factory that receives Warp
            scalar dtypes inferred from the runtime JAX arguments and returns a
            concrete Warp callable implementation.
        **kwargs: Options forwarded to Warp's JAX callable adapter.

    Returns:
        A callable compatible with JAX tracing, or a decorator producing one.
    """
    if func is None:
        return functools.partial(jax_callable, generic=generic, **kwargs)
    if not generic:
        return warp.jax_experimental.jax_callable(func, **kwargs)
    factory: _FfiCallableFactory = functools.lru_cache(func)
    return _FfiCallable(factory=factory, options=kwargs)  # ty:ignore[invalid-argument-type]

Wrap warp.jax_experimental.jax_kernel with optional overload lookup.

Parameters:

• kernel ¶

  (Callable | None, default: None ) –

  Warp kernel to expose to JAX. When omitted, return a decorator.

• arg_types_factory ¶

  (Callable[[WarpScalarDType], ArgTypes] | None, default: None ) –

  Optional callback that maps runtime Warp scalar dtypes to the overloaded kernel argument types expected by :func:warp.overload.

• num_outputs ¶

  (int, default: ... ) –

• vmap_method ¶

  (VmapMethod, default: ... ) –

• launch_dims ¶

  (ShapeLike | None, default: ... ) –

• output_dims ¶

  (ShapeLike | dict[str, ShapeLike] | None, default: ... ) –

• in_out_argnames ¶

  (Iterable[str], default: ... ) –

• module_preload_mode ¶

  (ModulePreloadMode, default: ... ) –

• enable_backward ¶

  (bool, default: ... ) –

Returns:

• Any –

  A callable compatible with JAX tracing, or a decorator producing one.

 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105

def jax_kernel(
    kernel: Callable | None = None,
    *,
    arg_types_factory: Callable[[WarpScalarDType], ArgTypes] | None = None,
    **kwargs: Unpack[JaxKernelOptions],
) -> Any:
    """Wrap ``warp.jax_experimental.jax_kernel`` with optional overload lookup.

    Args:
        kernel: Warp kernel to expose to JAX. When omitted, return a decorator.
        arg_types_factory: Optional callback that maps runtime Warp scalar dtypes
            to the overloaded kernel argument types expected by
            :func:`warp.overload`.
        **kwargs: Options forwarded to Warp's JAX kernel adapter.

    Returns:
        A callable compatible with JAX tracing, or a decorator producing one.
    """
    if kernel is None:
        return functools.partial(
            jax_kernel, arg_types_factory=arg_types_factory, **kwargs
        )
    if arg_types_factory is None:
        return warp.jax_experimental.jax_kernel(kernel, **kwargs)
    return _FfiKernel(
        kernel=cast("wp.Kernel", kernel),
        options=kwargs,  # ty:ignore[invalid-argument-type]
        arg_types_factory=arg_types_factory,
    )

Compile a callable with JAX, optionally preserving static PyTree leaves.

When filter=False this is a thin wrapper around :func:jax.jit. When filter=True the function and its inputs are partitioned into dynamic array leaves and static metadata so mixed PyTrees can cross the JIT boundary without requiring manual static_argnums wiring.

Parameters:

• fun ¶

  (Callable[P, T] | None, default: None ) –

  Callable to compile. When omitted, return a decorator.

• **kwargs ¶

  (Any, default: {} ) –

  Options forwarded to :func:jax.jit. With filter=True, only the subset in :class:FilterJitOptions is supported because static argument handling is managed internally.

Returns:

• Callable –

  A compiled callable or decorator with the same public call signature as

• Callable –

  fun.

56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87

def jit[**P, T](fun: Callable[P, T] | None = None, **kwargs: Any) -> Callable:
    """Compile a callable with JAX, optionally preserving static PyTree leaves.

    When ``filter=False`` this is a thin wrapper around :func:`jax.jit`.
    When ``filter=True`` the function and its inputs are partitioned into
    dynamic array leaves and static metadata so mixed PyTrees can cross the JIT
    boundary without requiring manual ``static_argnums`` wiring.

    Args:
        fun: Callable to compile. When omitted, return a decorator.
        **kwargs: Options forwarded to :func:`jax.jit`. With ``filter=True``,
            only the subset in :class:`FilterJitOptions` is supported because
            static argument handling is managed internally.

    Returns:
        A compiled callable or decorator with the same public call signature as
        ``fun``.
    """
    if fun is None:
        return functools.partial(jit, **kwargs)
    filter_: bool = kwargs.pop("filter", False)
    if not filter_:
        return jax.jit(fun, **kwargs)

    fun_data: _Data
    fun_meta: _Meta[Callable[P, T]]
    fun_data, fun_meta = tree.partition(fun)
    inner: _Inner[T] = _Inner(fun_meta)
    inner_jit: _InnerProtocol[T] = jax.jit(inner, static_argnums=(2,), **kwargs)
    outer: _Outer[P, T] = _Outer(fun_data=fun_data, inner=inner_jit)
    functools.update_wrapper(outer, fun)
    return outer

Partially apply a callable and keep the result compatible with JAX trees.

30
31
32

def partial[T](func: Callable[..., T], /, *args: Any, **kwargs: Any) -> Partial[..., T]:
    """Partially apply a callable and keep the result compatible with JAX trees."""
    return Partial(func, *args, **kwargs)

Flatten a PyTree's dynamic leaves into one vector.

Non-array leaves are treated as static metadata and preserved in the returned :class:Structure instead of being concatenated into the flat array.

Parameters:

• tree ¶

  (T) –

  PyTree to flatten.

Returns:

• Array –

  A tuple of (flat, structure) where flat is a one-dimensional

• Structure[T] –

  JAX array and structure can rebuild compatible trees later.

 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105

def ravel[T](tree: T) -> tuple[Array, Structure[T]]:
    """Flatten a PyTree's dynamic leaves into one vector.

    Non-array leaves are treated as static metadata and preserved in the
    returned :class:`Structure` instead of being concatenated into the flat
    array.

    Args:
        tree: PyTree to flatten.

    Returns:
        A tuple of ``(flat, structure)`` where ``flat`` is a one-dimensional
        JAX array and ``structure`` can rebuild compatible trees later.
    """
    leaves, treedef = jax.tree.flatten(tree)
    dynamic_leaves, static_leaves = partition_leaves(leaves)
    flat: Array = _ravel(dynamic_leaves)
    structure: Structure[T] = Structure(
        offsets=_offsets_from_leaves(dynamic_leaves),
        shapes=_shapes_from_leaves(dynamic_leaves),
        meta_leaves=tuple(static_leaves),
        treedef=treedef,
        dtype=flat.dtype,
    )
    return flat, structure

Convert a supported array object into a :class:warp.array.

The generic dispatcher currently supports NumPy arrays and JAX arrays. A dtype hint may be a concrete Warp dtype or a tuple that describes a vector or matrix dtype inferred from the trailing dimensions of arr.

12
13
14
15
16
17
18
19
20

@functools.singledispatch
def to_warp(arr: Any, *_args, **_kwargs) -> wp.array:
    """Convert a supported array object into a :class:`warp.array`.

    The generic dispatcher currently supports NumPy arrays and JAX arrays. A
    ``dtype`` hint may be a concrete Warp dtype or a tuple that describes a
    vector or matrix dtype inferred from the trailing dimensions of ``arr``.
    """
    raise TypeError(arr)

Run a loop with either jax.lax.while_loop or Python control flow.

Parameters:

• cond_fun ¶

  (Callable[[T], BooleanNumeric]) –

  Predicate evaluated on the loop state.

• body_fun ¶

  (Callable[[T], T]) –

  Function that produces the next loop state.

• init_val ¶

  (T) –

  Initial loop state.

• jit ¶

  (bool, default: True ) –

  When true, dispatch to :func:jax.lax.while_loop. When false, run an eager Python while loop with the same callbacks.

Returns:

• T –

  The final loop state.

 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

def while_loop[T](
    cond_fun: Callable[[T], BooleanNumeric],
    body_fun: Callable[[T], T],
    init_val: T,
    *,
    jit: bool = True,
) -> T:
    """Run a loop with either ``jax.lax.while_loop`` or Python control flow.

    Args:
        cond_fun: Predicate evaluated on the loop state.
        body_fun: Function that produces the next loop state.
        init_val: Initial loop state.
        jit: When true, dispatch to :func:`jax.lax.while_loop`. When false, run
            an eager Python ``while`` loop with the same callbacks.

    Returns:
        The final loop state.
    """
    if jit:
        return jax.lax.while_loop(cond_fun, body_fun, init_val)
    val: T = init_val
    while cond_fun(val):
        val: T = body_fun(val)
    return val