jax-ml/jax · error · NotImplementedError

serializing DisabledSafetyCheck: {check}

Error message

serializing DisabledSafetyCheck: {check}

What it means

JAX's export serialization writes each disabled safety check into a flatbuffer; only two kinds are supported: per-custom-call-target checks and the platform-wide check. If a DisabledSafetyCheck object is neither of those (e.g. a future or custom subclass), the serializer raises NotImplementedError.

Source

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

        )

  raise NotImplementedError(
      f"cannot deserialize effect type {effect_type_name}"
  )


def _serialize_disabled_safety_check(
    builder: flatbuffers.Builder, check: _export.DisabledSafetyCheck
) -> int:
  custom_call_target_str = check.is_custom_call()
  custom_call_target = None
  if custom_call_target_str is not None:
    kind = ser_flatbuf.DisabledSafetyCheckKind.custom_call
    custom_call_target = builder.CreateString(custom_call_target_str)
  elif check == _export.DisabledSafetyCheck.platform():
    kind = ser_flatbuf.DisabledSafetyCheckKind.platform
  else:
    raise NotImplementedError(f"serializing DisabledSafetyCheck: {check}")

  ser_flatbuf.DisabledSafetyCheckStart(builder)
  ser_flatbuf.DisabledSafetyCheckAddKind(builder, kind)
  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")
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create disabled checks only via the public API: jax.export.export(..., export_platforms=...) or export_disabled_checks with a known custom call target or platform=True
2. If subclassing DisabledSafetyCheck, serialize it as custom_call by setting the custom call target name instead
3. Match JAX versions between the process that created the export and the one serializing it

Example fix

// before
class MyCheck(jax.export.DisabledSafetyCheck): ...
checks = [MyCheck()]
// after
checks = [jax.export.DisabledSafetyCheck.custom_call('my_custom_target')]
// or the platform check
checks = [jax.export.DisabledSafetyCheck.platform()]

from jax._src import export as _export
ok = all(c.custom_call_target is not None or c == _export.DisabledSafetyCheck.platform() for c in checks)

def is_serializable_check(c) -> bool:
    return getattr(c, 'custom_call_target', None) is not None or c == jax.export.DisabledSafetyCheck.platform()

try:
    jax.export.save(path, exported)
except NotImplementedError as e:
    if 'DisabledSafetyCheck' in str(e): drop_or_replace_offending_checks()

When it happens

Trigger: Calling jax.export.save with an ExportedModules bundle whose disabled_checks contain a DisabledSafetyCheck that has no custom_call_target_str and is not DisabledSafetyCheck.platform() (e.g. a user-defined subclass or manually crafted object).

Common situations: Version mismatch: an ExportedModules object created or monkey-patched with a newer/older JAX that has additional check kinds; constructing DisabledSafetyCheck via internal APIs instead of export_disabled_checks(context=...) or platform=True.

Related errors

• No VJP is available
• Duplicate serialization registration for type `{nodetype}`.
• Duplicate serialization registration for serialized_name `{s
• multi-platform lowering for buffer_callback
• The error occurred in the __reduce__ method, which may indic

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