{"record":{"id":"4cb25914a4dc1b01","repo":"jax-ml/jax","slug":"value-returned-by-a-remat-policy-should-be-a-bool","errorCode":null,"errorMessage":"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)}.","messagePattern":"Value returned by a remat policy should be a bool or `ad_checkpoint\\.Recompute`, `ad_checkpoint\\.Saveable` or `ad_checkpoint\\.Offloadable\\(\\.\\.\\.\\)`\\. Got (.+?) of type (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/interpreters/partial_eval.py","lineNumber":987,"sourceCode":"class Offloadable(NamedTuple):\n  src: MemoryKind\n  dst: MemoryKind\n\nRematCases = RecomputeType | SaveableType | Offloadable\nRematCases_ = RematCases | bool\n\ndef ensure_enum(case: bool | RematCases) -> RematCases:\n  if isinstance(case, bool):\n    return Saveable if case else Recompute\n  if not isinstance(case, (RecomputeType, SaveableType, Offloadable)):\n    msg = (\"Value returned by a remat policy should be a bool or\"\n           \" `ad_checkpoint.Recompute`, `ad_checkpoint.Saveable` or\"\n           \" `ad_checkpoint.Offloadable(...)`.\"\n           f\" Got {case} of type {type(case)}.\")\n    if isinstance(case, Offloadable):\n      msg += (\"Did you return `Offloadable` instead of an instantiated\"\n              \" `Offloadable(...)`?\")\n    raise TypeError(msg)\n  return case\n\n# A primitive rule for policy-driven partial evaluation returns a 5-tuple\n# with the components representing, respectively:\n#  * the JaxprEqn for the 'known' side (or None if there is no known component),\n#  * the JaxprEqn for the 'unknown' side (or None),\n#  * a list of booleans indicating which of the original outputs are unknown,\n#  * a list of booleans indicating which of the original outputs are\n#    instantiated (i.e. available) in the 'unknown' side,\n#  * a list of Var instances representing residuals to be added (i.e. to be\n#    plumbed as outputs of the 'known' side jaxpr and added as input binders to\n#    the 'unknown' jaxpr).\nPartialEvalCustomResult = tuple[JaxprEqn | None, JaxprEqn | None,\n                                Sequence[bool], Sequence[bool], list[Var]]\nPartialEvalCustomRule = Callable[\n    [Callable[..., RematCases_], Sequence[bool], Sequence[bool], JaxprEqn],\n    PartialEvalCustomResult]\npartial_eval_jaxpr_custom_rules: dict[Primitive, PartialEvalCustomRule] = {}","sourceCodeStart":969,"sourceCodeEnd":1005,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/interpreters/partial_eval.py#L969-L1005","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Make the policy return only bool, ad_checkpoint.Recompute, ad_checkpoint.Saveable, or ad_checkpoint.Offloadable(...) on every path","Replace `return ad_checkpoint.Offloadable` with `return ad_checkpoint.Offloadable(ziel=...)` (an instantiated object)","Add a final `return False` (recompute) fallback so no path returns None"],"exampleFix":"# before\ndef my_policy(prim, *args):\n    if prim == 'dot_general':\n        return jax.ad_checkpoint.Offloadable\n\n# after\ndef my_policy(prim, *args):\n    if prim == 'dot_general':\n        return jax.ad_checkpoint.Offloadable(ziel='iot')\n    return False","handlingStrategy":"type-guard","validationCode":"from jax._src.ad_checkpoint import Offloadable\ndef check_policy_result(case):\n    assert isinstance(case, bool) or hasattr(case, 'recompute'), case","typeGuard":"def is_valid_policy_result(c) -> bool:\n    return isinstance(c, bool) or isinstance(c, (jax.ad_checkpoint.Recompute, jax.ad_checkpoint.Saveable)) or isinstance(c, jax.ad_checkpoint.Offloadable)","tryCatchPattern":null,"preventionTips":["End every custom policy with `return False`","Instantiate Offloadable with parentheses: Offloadable(...)"],"tags":["jax","remat","checkpoint-policy","type-error","ad-checkpoint"],"backgroundTag":"invalid-return-type-from-callback","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}