Checkify with nested JITs

[thomaspinder]

Consider the following MWE

import jax 
import jax.numpy as jnp
from jax.experimental import checkify


def _is_positive(value):
    checkify.check(
        jnp.all(value > 0), "value needs to be positive, got {value}", value=value
    )


def _check_is_positive(value) -> None:
    """Check if a value is positive.

    Args:
        value: The value to check.

    Raises:
        ValueError: If any element of value is negative.
    """
    checked_f = checkify.checkify(_is_positive)
    error, _ = jax.jit(checked_f)(value)
    checkify.check_error(error)

This is behaving as expected, and the following unit test would pass

import pytest
from jax.experimental import checkify


def test_check_is_positive():
    _check_is_positive(jnp.array(3.0))
    with pytest.raises(ValueError):
        _check_is_positive(jnp.array(-3.0))

However, should an end user try to JIT a function that wraps _check_is_positive, then an error is throw unless the new jitted fn. has checkify applied. You may see this below, were an end-user to implement _dummy_fn and then JIT it, then an error would be thrown. This is undesirable, and I would prefer if there was an alternative way that I can apply checkify, whilst safeguarding against end-users jitting code.

# Does not work
def test_jittable_checkify_1():

    def _dummy_fn(value):
        return _check_is_positive(value)

    jitted_fn = jax.jit(_dummy_fn)
    jitted_fn(jnp.array(3.0))


# Works
def test_jittable_checkify_2():

    def _dummy_fn(value):
        return _check_is_positive(value)

    jitted_fn = jax.jit(checkify.checkify(_dummy_fn))
    jitted_fn(jnp.array(3.0))

Stacktrace

    def test_jittable_checkify_1():

        def _dummy_fn(value):
            return _check_is_positive(value)

        jitted_fn = jit(_dummy_fn)
>       jitted_fn(jnp.array(3.0))
E       ValueError: Cannot abstractly evaluate a checkify.check which was not functionalized. This probably means you tried to stage (jit/scan/pmap/...) a `check` without functionalizing it through `checkify.checkify`.

tests/test_parameters.py:78: ValueError