jax-ml/jax · error · ValueError

Cannot deserialize DisabledSafetyCheck with unknown kind: {k

Error message

Cannot deserialize DisabledSafetyCheck with unknown kind: {kind}

What it means

When loading a serialized JAX export, the deserializer reads the flatbuffer tag identifying the kind of disabled safety check. An unrecognized tag means the flatbuffer was written by a newer JAX with new check kinds that this JAX version cannot represent.

Source

Thrown at jax/_src/export/serialization.py:941

  if custom_call_target is not None:
    ser_flatbuf.DisabledSafetyCheckAddCustomCallTarget(
        builder, custom_call_target
    )
  return ser_flatbuf.DisabledSafetyCheckEnd(builder)


def _deserialize_disabled_safety_check(
    sc: ser_flatbuf.DisabledSafetyCheck,
) -> _export.DisabledSafetyCheck:
  kind = sc.Kind()
  if kind == ser_flatbuf.DisabledSafetyCheckKind.custom_call:
    return _export.DisabledSafetyCheck.custom_call(
        sc.CustomCallTarget().decode("utf-8")
    )
  if kind == ser_flatbuf.DisabledSafetyCheckKind.platform:
    return _export.DisabledSafetyCheck.platform()

  raise ValueError(f"Cannot deserialize DisabledSafetyCheck with unknown kind: {kind}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Upgrade jax (and jaxlib) to at least the version that produced the export file
2. Re-export the model with the currently installed JAX version
3. Verify the file integrity / regenerate the export

Example fix

# before
m = jax.export.load('model.jax')  # older jax, newer file
# after
pip install -U jax jaxlib
m = jax.export.load('model.jax')

import jax, jaxlib
# before loading an export produced elsewhere
assert version_tuple(jax.__version__) >= version_tuple(required_version)

def loadable_by_this_jax(path) -> bool:
    return jax.__version__ >= export_producer_version(path)

try:
    em = jax.export.load(path)
except ValueError as e:
    if 'unknown kind' in str(e): upgrade_or_reexport(path)

When it happens

Trigger: jax.export.load(path) on a file serialized by a newer JAX version that added a new DisabledSafetyCheckKind; or a corrupted/hand-edited flatbuffer with an invalid enum value.

Common situations: Downgrading JAX after an export file was produced with a newer release; reading a .jax file generated on a different toolchain.

Related errors

• stop_gradient only works on valid JAX arrays, but input argu
• Unknown GPU platform for __dlpack__: {platform_version}
• multi-platform lowering for buffer_callback
• mla requires cudnn version >= 9.10 and at least hopper arch.
• mTLS for the JAX distributed service requires jaxlib 0.11.2

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