partial eval silently skips effects about jax HOT 3 OPEN

import jax
from jax.experimental import checkify
from jax._src.interpreters import partial_eval as pe


def fn(x):
    checkify.check(x < 10, "checkify")
    jax.debug.print("callback: {}", x)
    return x


jaxpr = jax.make_jaxpr(fn)(1)
jaxpr_known, jaxpr_unkown, out_unknowns, *_ = (
    pe.partial_eval_jaxpr_stateful(jaxpr.jaxpr, in_unknowns=[True],
                                   in_inst=[True], ensure_out_unknowns=False,
                                   ensure_out_inst=False, saveable=lambda *_,
                                   **__: True)
)
print(jaxpr_known, jaxpr_unkown, out_unknowns)

import jax
from jax import Array
from jax import numpy as jnp
from jax._src import checkify, effects

# effects.custom_derivatives_allowed_effects.add_type(checkify.ErrorEffect)


@jax.custom_jvp
def fn(x: Array):
    return x


@fn.defjvp
def testfn_jvp(primals: tuple[Array, ...], tangents: tuple[Array, ...]):

    (x,) = primals
    (tg,) = tangents

    jax.debug.print("callback: {}", tg)
    # checkify.check(tg < 0, "invalid tangent {}", tg)

    return x, tg


# jvp works
# jax.jvp(fn, [jnp.array(1.0)], [jnp.array(2.0)])

val, tang = jax.linearize(fn, jnp.array(1.0))
tang(jnp.array(2.0))