jax-ml/jax · error · TypeError

Value returned by a remat policy should be a bool or `ad_che

Error message

Value returned by a remat policy should be a bool or `ad_checkpoint.Recompute`, `ad_checkpoint.Saveable` or `ad_checkpoint.Offloadable(...)`. Got {case} of type {type(case)}.

What it means

A remat policy (the `saveable` callable passed to jax.checkpoint / jax.remat policies) must return booleans or members of jax.ad_checkpoint (Recompute, Saveable, Offloadable(...)). ensure_enum validates each returned value; anything else — including the class Offloadable itself instead of an instance — raises TypeError with a targeted hint.

Source

Thrown at jax/_src/interpreters/partial_eval.py:987

class Offloadable(NamedTuple):
  src: MemoryKind
  dst: MemoryKind

RematCases = RecomputeType | SaveableType | Offloadable
RematCases_ = RematCases | bool

def ensure_enum(case: bool | RematCases) -> RematCases:
  if isinstance(case, bool):
    return Saveable if case else Recompute
  if not isinstance(case, (RecomputeType, SaveableType, Offloadable)):
    msg = ("Value returned by a remat policy should be a bool or"
           " `ad_checkpoint.Recompute`, `ad_checkpoint.Saveable` or"
           " `ad_checkpoint.Offloadable(...)`."
           f" Got {case} of type {type(case)}.")
    if isinstance(case, Offloadable):
      msg += ("Did you return `Offloadable` instead of an instantiated"
              " `Offloadable(...)`?")
    raise TypeError(msg)
  return case

# A primitive rule for policy-driven partial evaluation returns a 5-tuple
# with the components representing, respectively:
#  * the JaxprEqn for the 'known' side (or None if there is no known component),
#  * the JaxprEqn for the 'unknown' side (or None),
#  * a list of booleans indicating which of the original outputs are unknown,
#  * a list of booleans indicating which of the original outputs are
#    instantiated (i.e. available) in the 'unknown' side,
#  * a list of Var instances representing residuals to be added (i.e. to be
#    plumbed as outputs of the 'known' side jaxpr and added as input binders to
#    the 'unknown' jaxpr).
PartialEvalCustomResult = tuple[JaxprEqn | None, JaxprEqn | None,
                                Sequence[bool], Sequence[bool], list[Var]]
PartialEvalCustomRule = Callable[
    [Callable[..., RematCases_], Sequence[bool], Sequence[bool], JaxprEqn],
    PartialEvalCustomResult]
partial_eval_jaxpr_custom_rules: dict[Primitive, PartialEvalCustomRule] = {}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the policy return only bool, ad_checkpoint.Recompute, ad_checkpoint.Saveable, or ad_checkpoint.Offloadable(...) on every path
2. Replace `return ad_checkpoint.Offloadable` with `return ad_checkpoint.Offloadable(ziel=...)` (an instantiated object)
3. Add a final `return False` (recompute) fallback so no path returns None

Example fix

# before
def my_policy(prim, *args):
    if prim == 'dot_general':
        return jax.ad_checkpoint.Offloadable

# after
def my_policy(prim, *args):
    if prim == 'dot_general':
        return jax.ad_checkpoint.Offloadable(ziel='iot')
    return False

from jax._src.ad_checkpoint import Offloadable
def check_policy_result(case):
    assert isinstance(case, bool) or hasattr(case, 'recompute'), case

def is_valid_policy_result(c) -> bool:
    return isinstance(c, bool) or isinstance(c, (jax.ad_checkpoint.Recompute, jax.ad_checkpoint.Saveable)) or isinstance(c, jax.ad_checkpoint.Offloadable)

When it happens

Trigger: Writing a custom remat policy (policy=lambda prim, inner_axis, *args: ...) that returns None, an int, a string, or the class jax.ad_checkpoint.Offloadable rather than Offloadable(...) instantiated. Called via _partial_eval_jaxpr_custom_cached during tracing of a checkpointed function.

Common situations: Custom remat policies with a missing return path (implicit None); copy-paste code returning ad_checkpoint.Offloadable without parentheses; returning numpy bools or other truthy objects instead of Python bools.

Related errors

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.
• Cannot interpret value of type {typ} as an abstract array; i

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/4cb25914a4dc1b01. Report an issue: GitHub.