|
| 1 | +"""Copyright (c) 2026 Nathaniel Starkman. All rights reserved.""" |
| 2 | + |
| 3 | +__all__ = ("error_if",) |
| 4 | + |
| 5 | +from typing import Any, Literal, TypeVar |
| 6 | + |
| 7 | +import jax |
| 8 | +import jax.numpy as jnp |
| 9 | + |
| 10 | +T = TypeVar("T") |
| 11 | + |
| 12 | + |
| 13 | +def error_if( |
| 14 | + value: T, |
| 15 | + condition: Any, |
| 16 | + msg: str, |
| 17 | + /, |
| 18 | + *, |
| 19 | + on_error: Literal["raise", "off"] = "raise", |
| 20 | +) -> T: |
| 21 | + """Raise an error if a condition is true (JAX-compatible). |
| 22 | +
|
| 23 | + This function is compatible with JAX transformations (jit, vmap, grad) and |
| 24 | + raises an error at runtime if the condition evaluates to True. The raised |
| 25 | + exception is always `jax.errors.JaxRuntimeError`. |
| 26 | +
|
| 27 | + Uses `jax.debug.callback` internally to execute the error-raising logic |
| 28 | + outside of JAX's tracing system. This has important consequences: |
| 29 | +
|
| 30 | + - **Runtime execution**: The callback executes at runtime, not during |
| 31 | + tracing, so it can inspect actual runtime values (not tracer objects). |
| 32 | + - **CPU execution**: The callback always runs on CPU, even if the rest of |
| 33 | + the JAX computation is on GPU/TPU. This is a mild performance cost but |
| 34 | + necessary for error handling. |
| 35 | + - **Not dead code eliminated**: Despite using a callback, the error check is |
| 36 | + NOT eliminated by JAX's optimizations (unlike pure computations that may |
| 37 | + be DCE'd). The callback is considered a side effect and is preserved. |
| 38 | + - **With jit/vmap**: Works seamlessly inside `jax.jit` and `jax.vmap` |
| 39 | + contexts. When vmapped, the callback is called for each mapped instance. |
| 40 | + - **Exception wrapping**: Under `jax.jit`, exceptions raised from |
| 41 | + `jax.debug.callback` are typically wrapped by JAX as |
| 42 | + `jax.errors.JaxRuntimeError` (backend-dependent). Outside jit, the |
| 43 | + `JaxRuntimeError` is raised directly. |
| 44 | +
|
| 45 | + Parameters |
| 46 | + ---------- |
| 47 | + value |
| 48 | + The value to return if no error is raised. |
| 49 | + condition |
| 50 | + A boolean or JAX array (scalar or multi-element) indicating whether to |
| 51 | + raise an error. If an array, the error is raised if any element is True. |
| 52 | + msg |
| 53 | + The error message to raise if condition is True. |
| 54 | + on_error |
| 55 | + Controls error handling behavior: |
| 56 | +
|
| 57 | + - ``"raise"`` (default): Raise an error when condition is True. |
| 58 | + - ``"off"``: |
| 59 | + Disable error checking entirely. This is a complete no-op that skips |
| 60 | + even the callback, providing maximum performance when error checks |
| 61 | + are not needed. |
| 62 | +
|
| 63 | + Returns |
| 64 | + ------- |
| 65 | + T |
| 66 | + The input value if condition is False, otherwise an error is raised. |
| 67 | +
|
| 68 | + Examples |
| 69 | + -------- |
| 70 | + Basic usage (raises error if condition is true): |
| 71 | +
|
| 72 | + .. code-block:: python |
| 73 | +
|
| 74 | + import jax.numpy as jnp |
| 75 | + from jaxmore._src import error_if |
| 76 | +
|
| 77 | + x = jnp.array(5) |
| 78 | + result = error_if(x, x > 10, "x is too large") |
| 79 | + # result = Array(5, dtype=int32) |
| 80 | +
|
| 81 | + Disable error checking for performance: |
| 82 | +
|
| 83 | + .. code-block:: python |
| 84 | +
|
| 85 | + import jax.numpy as jnp |
| 86 | + from jaxmore._src import error_if |
| 87 | +
|
| 88 | + x = jnp.array(15) |
| 89 | + result = error_if(x, x > 10, "ignored", on_error="off") |
| 90 | + # No callback is executed; this is a complete no-op |
| 91 | +
|
| 92 | + Notes |
| 93 | + ----- |
| 94 | + For most production code requiring flexible error handling, |
| 95 | + `equinox.error_if` is recommended. The `equinox.error_if` implementation is |
| 96 | + generally more feature-rich and better integrated with the broader |
| 97 | + JAX/Equinox ecosystem: |
| 98 | +
|
| 99 | + - **Multiple error handling modes**: Equinox supports ``"raise"``, |
| 100 | + ``"breakpoint"`` (with configurable debugger frames), ``"nan"`` (replace |
| 101 | + values and continue), ``"warn"`` (emit warning), and ``"off"``. |
| 102 | + - **Better error messages**: Integration with `equinox.filter_jit` provides |
| 103 | + cleaner stack traces and more informative error reporting. |
| 104 | + - **Platform-specific handling**: Special logic for TPU runtime, which |
| 105 | + normally squelches errors. |
| 106 | + - **Automatic differentiation**: Custom JVP rules ensure correct behavior |
| 107 | + under `jax.grad` and other AD transforms. |
| 108 | + - **Dead code elimination awareness**: Uses `lax.cond` to conditionally |
| 109 | + branch, allowing JAX's compiler to optimize away unused checks (though |
| 110 | + this requires the return value to be used in the computation). |
| 111 | +
|
| 112 | + """ |
| 113 | + if on_error == "off": |
| 114 | + # Complete no-op: skip the callback entirely for maximum performance. |
| 115 | + return value |
| 116 | + |
| 117 | + # Reduce array conditions to scalar boolean using 'any' semantics. |
| 118 | + # This ensures the callback can use 'if cond_val:' without ambiguity. |
| 119 | + scalar_condition = jnp.asarray(condition).astype(bool) |
| 120 | + if scalar_condition.ndim > 0: |
| 121 | + scalar_condition = jnp.any(scalar_condition) |
| 122 | + |
| 123 | + def callback(cond_val: Any) -> None: |
| 124 | + # Callback to check condition and raise error. |
| 125 | + if cond_val: |
| 126 | + raise jax.errors.JaxRuntimeError(str(msg)) |
| 127 | + |
| 128 | + jax.debug.callback(callback, scalar_condition) |
| 129 | + return value |
0 commit comments