liblaf.jarp ¶

Create a data field whose default is normalized to a JAX array.

When default is a concrete array-like value, array rewrites it into a factory so each instance receives its own array object.

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: ... ) –

67
68
69
70
71
72
73
74
75
76
77
78
79

def array(**kwargs: Unpack[FieldOptions[Any]]) -> Array:
    """Create a data field whose default is normalized to a JAX array.

    When `default` is a concrete array-like value, `array` rewrites it into
    a factory so each instance receives its own array object.
    """
    if "default" in kwargs and "factory" not in kwargs:
        default: Any = 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)  # ty:ignore[no-matching-overload]

Create a field whose PyTree role is chosen from the runtime value.

82
83
84
85
86

@_wraps(attrs.field)
def auto(**kwargs) -> Any:
    """Create a field whose PyTree role is chosen from the runtime value."""
    kwargs.setdefault("static", FieldType.AUTO)
    return field(**kwargs)

Choose between two branches, then retry eagerly if JAX rejects them.

The wrapper first calls jax.lax.cond. If that raises jax.errors.JAXTypeError or jax.errors.JAXIndexError, it logs the exception and reruns the selected branch in plain Python.

Parameters:

• pred (ScalarLike) –

  Scalar predicate. Python truthiness decides which branch runs on the fallback path.

• true_fun (Callable[[*Ts], T]) –

  Branch evaluated when pred is true.

• false_fun (Callable[[*Ts], T]) –

  Branch evaluated when pred is false.

• *operands (*Ts, default: () ) –

  Positional operands forwarded to the selected branch.

Returns:

• T –

  The value returned by the selected branch.

12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38

@lax_wrapper(jax.lax.cond)
def cond[*Ts, T](
    pred: ScalarLike,
    true_fun: Callable[[*Ts], T],
    false_fun: Callable[[*Ts], T],
    *operands: *Ts,
) -> T:
    """Choose between two branches, then retry eagerly if JAX rejects them.

    The wrapper first calls [`jax.lax.cond`][jax.lax.cond]. If that raises
    [`jax.errors.JAXTypeError`][jax.errors.JAXTypeError] or
    [`jax.errors.JAXIndexError`][jax.errors.JAXIndexError], it logs the
    exception and reruns the selected branch in plain Python.

    Args:
        pred: Scalar predicate. Python truthiness decides which branch runs on
            the fallback path.
        true_fun: Branch evaluated when `pred` is true.
        false_fun: Branch evaluated when `pred` is false.
        *operands: Positional operands forwarded to the selected branch.

    Returns:
        The value returned by the selected branch.
    """
    if pred:
        return true_fun(*operands)
    return false_fun(*operands)

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 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.

67
68
69
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

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 [`attrs.define`][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

Wrap a callable and cache Python fallbacks for failing metadata shapes.

The wrapper first uses the same partitioned call path as filter_jit. If that path raises jax.errors.JAXTypeError or jax.errors.JAXIndexError, the exception is logged, the current static-metadata signature is marked as unsupported, and the original callable is invoked directly in Python. Later calls with the same static metadata skip the partitioned path and reuse the Python fallback immediately.

Examples:

>>> import jax.numpy as jnp
>>> from liblaf import jarp
>>> @jarp.fallback_jit
... def add_one(value):
...     return value + 1
>>> int(add_one(jnp.array(2)))
3

Parameters:

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

  Callable to wrap. When omitted, return a configured decorator.

• keep_unused (bool, default: ... ) –
• device (Any | None, default: ... ) –
• backend (str | None, default: ... ) –
• inline (bool, default: ... ) –

Returns:

• Callable –

  The wrapped callable, or a decorator that produces one.

 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
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115

def fallback_jit[**P, T](
    fun: Callable[P, T] | None = None, **kwargs: Unpack[FilterJitOptions]
) -> Callable:
    """Wrap a callable and cache Python fallbacks for failing metadata shapes.

    The wrapper first uses the same partitioned call path as
    [`filter_jit`][liblaf.jarp.filter_jit]. If that path raises
    [`jax.errors.JAXTypeError`][jax.errors.JAXTypeError] or
    [`jax.errors.JAXIndexError`][jax.errors.JAXIndexError], the exception is
    logged, the current static-metadata signature is marked as unsupported,
    and the original callable is invoked directly in Python. Later calls with
    the same static metadata skip the partitioned path and reuse the Python
    fallback immediately.

    Examples:
        >>> import jax.numpy as jnp
        >>> from liblaf import jarp
        >>> @jarp.fallback_jit
        ... def add_one(value):
        ...     return value + 1
        >>> int(add_one(jnp.array(2)))
        3

    Args:
        fun: Callable to wrap. When omitted, return a configured decorator.
        **kwargs: Options forwarded to [`jax.jit`][jax.jit] for the inner
            filtered callable.

    Returns:
        The wrapped callable, or a decorator that produces one.
    """
    if fun is None:
        return functools.partial(fallback_jit, **kwargs)
    fun_data, fun_meta = tree.partition(fun)
    inner: FilterInner = FilterInner(fun_meta=fun_meta)
    inner_jit: jax.stages.Wrapped = jax.jit(inner, **kwargs)
    outer: FallbackOuter = FallbackOuter(
        fun_data=fun_data, fun_meta=fun_meta, inner=inner_jit
    )
    functools.update_wrapper(outer, fun)
    return outer

Create an attrs field using jarp's static metadata convention.

89
90
91
92
93
94
95
96
97

@_wraps(attrs.field)
def field(**kwargs) -> Any:
    """Create an `attrs` field using jarp's `static` metadata convention."""
    if "static" in kwargs:
        kwargs["metadata"] = {
            "static": kwargs.pop("static"),
            **(kwargs.get("metadata") or {}),
        }
    return attrs.field(**kwargs)

Wrap a callable with liblaf.jarp data-versus-metadata partitioning.

The wrapper partitions the callable and each invocation's arguments with partition, rebuilds the original call shape, and partitions the return value again before handing it back. This keeps JAX arrays on the dynamic side of the partition while preserving ordinary Python metadata such as strings, bound methods, or configuration objects.

Examples:

>>> import jax.numpy as jnp
>>> from liblaf import jarp
>>> @jarp.filter_jit
... def scale(value, *, label):
...     assert label == "active"
...     return value * 2
>>> scale(jnp.array([1, 2]), label="active").tolist()
[2, 4]

Parameters:

• fun (F | None, default: None ) –

  Callable to wrap. When omitted, return a configured decorator.

• keep_unused (bool, default: ... ) –
• device (Any | None, default: ... ) –
• backend (str | None, default: ... ) –
• inline (bool, default: ... ) –

Returns:

• Callable[..., Any] –

  The wrapped callable, or a decorator that produces one.

 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
106
107
108
109
110
111
112
113
114
115
116
117

def filter_jit[F: Callable[..., Any]](
    fun: F | None = None, **kwargs: Unpack[FilterJitOptions]
) -> Callable[..., Any]:
    """Wrap a callable with `liblaf.jarp` data-versus-metadata partitioning.

    The wrapper partitions the callable and each invocation's arguments with
    [`partition`][liblaf.jarp.tree.partition], rebuilds the original call shape,
    and partitions the return value again before handing it back. This keeps
    JAX arrays on the dynamic side of the partition while preserving ordinary
    Python metadata such as strings, bound methods, or configuration objects.

    Examples:
        >>> import jax.numpy as jnp
        >>> from liblaf import jarp
        >>> @jarp.filter_jit
        ... def scale(value, *, label):
        ...     assert label == "active"
        ...     return value * 2
        >>> scale(jnp.array([1, 2]), label="active").tolist()
        [2, 4]

    Args:
        fun: Callable to wrap. When omitted, return a configured decorator.
        **kwargs: Options forwarded to [`jax.jit`][jax.jit] for the inner
            filtered callable.

    Returns:
        The wrapped callable, or a decorator that produces one.
    """
    if fun is None:
        return functools.partial(filter_jit, **kwargs)
    fun_data, fun_meta = tree.partition(fun)
    inner: FilterInner = FilterInner(fun_meta=fun_meta)
    inner_jit: InnerLike = jax.jit(inner, **kwargs)
    outer: FilterOuter = FilterOuter(fun_data=fun_data, inner=inner_jit)
    functools.update_wrapper(outer, fun)
    return outer

Run a counted loop, then retry in Python if JAX rejects the body.

The wrapper first calls jax.lax.fori_loop. If that raises jax.errors.JAXTypeError or jax.errors.JAXIndexError, it logs the exception and runs an ordinary Python for loop instead.

Parameters:

• lower (int) –

  Inclusive loop lower bound.

• upper (int) –

  Exclusive loop upper bound.

• body_fun (Callable[[int, T], T]) –

  Callback that receives the iteration index and current loop value, then returns the next loop value.

• init_val (T) –

  Initial loop value.

• **kwargs (Any, default: {} ) –

  Extra keyword arguments forwarded to jax.lax.fori_loop on the first attempt. They are ignored on the Python fallback path.

Returns:

• T –

  The final loop value.

41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73

@lax_wrapper(jax.lax.fori_loop)
def fori_loop[T](
    lower: int,
    upper: int,
    body_fun: Callable[[int, T], T],
    init_val: T,
    **kwargs: Any,
) -> T:
    """Run a counted loop, then retry in Python if JAX rejects the body.

    The wrapper first calls [`jax.lax.fori_loop`][jax.lax.fori_loop]. If that
    raises [`jax.errors.JAXTypeError`][jax.errors.JAXTypeError] or
    [`jax.errors.JAXIndexError`][jax.errors.JAXIndexError], it logs the
    exception and runs an ordinary Python `for` loop instead.

    Args:
        lower: Inclusive loop lower bound.
        upper: Exclusive loop upper bound.
        body_fun: Callback that receives the iteration index and current loop
            value, then returns the next loop value.
        init_val: Initial loop value.
        **kwargs: Extra keyword arguments forwarded to
            [`jax.lax.fori_loop`][jax.lax.fori_loop] on the first attempt.
            They are ignored on the Python fallback path.

    Returns:
        The final loop value.
    """
    del kwargs
    val: T = init_val
    for i in range(lower, upper):
        val: T = body_fun(i, val)
    return val

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

This is the common choice for immutable structures whose array fields should participate in JAX transformations.

112
113
114
115
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.

    This is the common choice for immutable structures whose array fields
    should participate in JAX transformations.
    """
    _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.

Use this for immutable helper objects that should be treated as static metadata instead of flattening into JAX leaves.

137
138
139
140
141
142
143
144
145
146
147
148

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.

    Use this for immutable helper objects that should be treated as static
    metadata instead of flattening into JAX leaves.
    """
    _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.

When generic=True, func is treated as a factory keyed by the Warp scalar dtypes inferred from the runtime JAX arguments. The factory output is cached, so repeated calls with the same dtype signature reuse the same Warp callable.

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 (VmapMethod | 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: ... ) –
• has_side_effect (bool, default: ... ) –

Returns:

• Any –

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

• Any –

  The callable returns the output arrays produced by Warp's FFI wrapper.

 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117

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

    When `generic=True`, `func` is treated as a factory keyed by the Warp
    scalar dtypes inferred from the runtime JAX arguments. The factory output is
    cached, so repeated calls with the same dtype signature reuse the same Warp
    callable.

    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.
        The callable returns the output arrays produced by Warp's FFI wrapper.
    """
    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)

Wrap warp.jax_experimental.jax_kernel with optional overload lookup.

When arg_types_factory is provided, the wrapper infers Warp scalar dtypes from the runtime JAX arguments, builds an overload signature, and resolves the corresponding Warp kernel before dispatch.

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 [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.

• Any –

  The callable returns the output arrays produced by Warp's FFI wrapper.

 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
106
107
108
109
110

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.

    When `arg_types_factory` is provided, the wrapper infers Warp scalar
    dtypes from the runtime JAX arguments, builds an overload signature, and
    resolves the corresponding Warp kernel before dispatch.

    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
            [warp.overload][].
        **kwargs: Options forwarded to Warp's JAX kernel adapter.

    Returns:
        A callable compatible with JAX tracing, or a decorator producing one.
        The callable returns the output arrays produced by Warp's FFI wrapper.
    """
    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,
        arg_types_factory=arg_types_factory,
    )

Partially apply a callable and keep bound values visible to JAX trees.

48
49
50

def partial[T](func: Callable[..., T], /, *args: Any, **kwargs: Any) -> Partial[..., T]:
    """Partially apply a callable and keep bound values visible to 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 Structure instead of being concatenated into the flat array.

Examples:

>>> import jax.numpy as jnp
>>> from liblaf import jarp
>>> flat, structure = jarp.ravel({"x": jnp.array([1.0, 2.0]), "tag": "train"})
>>> flat.tolist()
[1.0, 2.0]
>>> rebuilt = structure.unravel(jnp.array([3.0, 4.0]))
>>> rebuilt["x"].tolist(), rebuilt["tag"]
([3.0, 4.0], 'train')

Parameters:

• tree (T) –

  PyTree to flatten.

Returns:

• Array –

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

• Structure[T] –

  array and structure can rebuild compatible trees later.

 85
 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
111
112
113
114
115
116
117
118
119

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 [`Structure`][liblaf.jarp.tree.Structure] instead of being
    concatenated into the flat array.

    Examples:
        >>> import jax.numpy as jnp
        >>> from liblaf import jarp
        >>> flat, structure = jarp.ravel({"x": jnp.array([1.0, 2.0]), "tag": "train"})
        >>> flat.tolist()
        [1.0, 2.0]
        >>> rebuilt = structure.unravel(jnp.array([3.0, 4.0]))
        >>> rebuilt["x"].tolist(), rebuilt["tag"]
        ([3.0, 4.0], 'train')

    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

Create a field that is always treated as static metadata.

100
101
102
103
104
105

@_wraps(attrs.field)
def static(**kwargs) -> Any:
    """Create a field that is always treated as static metadata."""
    # for consistency with `jax.tree_util.register_dataclass`
    kwargs.setdefault("static", True)
    return field(**kwargs)

Decorate a class as a Warp struct.

Plain classes are forwarded to warp.struct. Classes that define __annotations_factory__(dtype) stay generic: MyStruct[wp.float64] builds and caches a specialized Warp struct from the factory annotations, while MyStruct() instantiates MyStruct[liblaf.jarp.warp.types.floating] so the default follows JAX's active precision mode.

Parameters:

• cls (T) –

  Class to decorate.

Returns:

• T –

  The Warp struct for plain classes, or the original generic class with

• T –

  dtype subscription and default construction hooks installed.

11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

def struct[T: type](cls: T) -> T:
    """Decorate a class as a Warp struct.

    Plain classes are forwarded to `warp.struct`. Classes that define
    `__annotations_factory__(dtype)` stay generic: `MyStruct[wp.float64]`
    builds and caches a specialized Warp struct from the factory annotations,
    while `MyStruct()` instantiates `MyStruct[liblaf.jarp.warp.types.floating]`
    so the default follows JAX's active precision mode.

    Args:
        cls: Class to decorate.

    Returns:
        The Warp struct for plain classes, or the original generic class with
        dtype subscription and default construction hooks installed.
    """
    if not hasattr(cls, "__annotations_factory__"):
        return cast("T", wp.struct(cls))

    @functools.cache
    def __class_getitem__(cls: T, key: Any) -> T:  # noqa: N807
        c: type = type(
            cls.__name__,
            (cls,),
            {
                "__module__": cls.__module__,
                "__qualname__": cls.__qualname__,
                "__annotations__": cls.__annotations_factory__(key),  # ty:ignore[unresolved-attribute]
            },
        )
        return cast("T", wp.struct(c, module="unique"))

    def __new__(owner: type) -> object:  # noqa: N807
        if owner is cls:
            return __class_getitem__(cls, wpt.floating)()
        return object.__new__(owner)

    cls.__class_getitem__ = classmethod(__class_getitem__)  # ty:ignore[invalid-assignment]
    cls.__new__ = staticmethod(__new__)  # ty:ignore[invalid-assignment]
    return cls

Choose one branch by index, then retry eagerly if JAX rejects it.

The wrapper first calls jax.lax.switch. If that raises jax.errors.JAXTypeError or jax.errors.JAXIndexError, it logs the exception, clamps index into the valid range, and dispatches in plain Python.

Parameters:

• index (ArrayLike) –

  Branch index. The fallback path clamps the value into the valid range before dispatch.

• branches (Sequence[Callable[[*Ts], T]]) –

  Candidate branch functions.

• *operands (*Ts, default: () ) –

  Positional operands forwarded to the selected branch.

Returns:

• T –

  The value returned by the selected branch.

76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98

@lax_wrapper(jax.lax.switch)
def switch[*Ts, T](
    index: ArrayLike, branches: Sequence[Callable[[*Ts], T]], *operands: *Ts
) -> T:
    """Choose one branch by index, then retry eagerly if JAX rejects it.

    The wrapper first calls [`jax.lax.switch`][jax.lax.switch]. If that raises
    [`jax.errors.JAXTypeError`][jax.errors.JAXTypeError] or
    [`jax.errors.JAXIndexError`][jax.errors.JAXIndexError], it logs the
    exception, clamps `index` into the valid range, and dispatches in plain
    Python.

    Args:
        index: Branch index. The fallback path clamps the value into the valid
            range before dispatch.
        branches: Candidate branch functions.
        *operands: Positional operands forwarded to the selected branch.

    Returns:
        The value returned by the selected branch.
    """
    index: Array = jax.lax.clamp(index, 0, len(branches) - 1)
    return branches[cast("int", index)](*operands)

Convert a supported array object into a [warp.array][liblaf.jarp.warp.array].

The dispatcher supports existing Warp arrays, 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. Use (-1, Any) for vector inference and (-1, -1, Any) for matrix inference when the element type should follow the source array.

Parameters:

• arr (array | ndarray | Array) –

  Array object to convert.

• *_args (Any, default: () ) –

  Reserved for singledispatch compatibility.

• **_kwargs (Any, default: {} ) –

  Reserved for singledispatch compatibility.

Returns:

• array –

  A Warp array view or converted array, depending on the source type.

Raises:

• TypeError –

  If arr uses an unsupported type.

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

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

    The dispatcher supports existing Warp arrays, 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`.
    Use `(-1, Any)` for vector inference and `(-1, -1, Any)` for matrix
    inference when the element type should follow the source array.

    Args:
        arr: Array object to convert.
        *_args: Reserved for singledispatch compatibility.
        **_kwargs: Reserved for singledispatch compatibility.

    Returns:
        A Warp array view or converted array, depending on the source type.

    Raises:
        TypeError: If `arr` uses an unsupported type.
    """
    raise TypeError(arr)

Run a loop, then retry in Python if JAX rejects the callbacks.

The wrapper first calls jax.lax.while_loop. If that raises jax.errors.JAXTypeError or jax.errors.JAXIndexError, it logs the exception and reruns the loop eagerly in Python.

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.

Returns:

• T –

  The final loop state.

101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123

@lax_wrapper(jax.lax.while_loop)
def while_loop[T](
    cond_fun: Callable[[T], BooleanNumeric], body_fun: Callable[[T], T], init_val: T
) -> T:
    """Run a loop, then retry in Python if JAX rejects the callbacks.

    The wrapper first calls [`jax.lax.while_loop`][jax.lax.while_loop]. If
    that raises [`jax.errors.JAXTypeError`][jax.errors.JAXTypeError] or
    [`jax.errors.JAXIndexError`][jax.errors.JAXIndexError], it logs the
    exception and reruns the loop eagerly in Python.

    Args:
        cond_fun: Predicate evaluated on the loop state.
        body_fun: Function that produces the next loop state.
        init_val: Initial loop state.

    Returns:
        The final loop state.
    """
    val: T = init_val
    while cond_fun(val):
        val: T = body_fun(val)
    return val