xai-org/x-algorithm · error · ValueError

Compilation does not support serialization

Error message

Compilation does not support serialization

What it means

partially_serialize relies on a private _unloaded_executable attribute on the Compiled object, which only exists for executables produced by the XLA AOT path that supports serialization. If the Compiled object came from an ordinary jax jit compile (or a jaxlib version without that attribute), serialization is unsupported.

Source

Thrown at phoenix/xrex/utils/aot.py:191

                unloaded_executable.output_shardings,
                partially_serialized.out_shardings_mem_kinds,
            )
        ]
        unloaded_executable.output_shardings = restored_out_shardings

    return jax.stages.Compiled(
        unloaded_executable.load(),
        [],
        lowered.args_info,
        lowered.out_tree,
        no_kwargs=partially_serialized.no_kwargs,
    )


def partially_serialize(compiled: Compiled, execution_devices: Sequence[xc.Device]):
    unloaded_exec = getattr(compiled._executable, "_unloaded_executable", None)
    if unloaded_exec is None:
        raise ValueError("Compilation does not support serialization")

    in_shardings_mem_kinds = [s.memory_kind for s in unloaded_exec.input_shardings]
    out_shardings_mem_kinds = [s.memory_kind for s in unloaded_exec.output_shardings]

    compile_options = getattr(unloaded_exec, "compile_options", None)
    serialized_compile_options = None
    if compile_options is not None:
        serialized_compile_options = compile_options.SerializeAsString()

    with io.BytesIO() as file:
        _JaxPjrtPickler(file, execution_devices).dump(unloaded_exec)
        return PartiallySerialized(
            file.getvalue(),
            in_shardings_mem_kinds,
            out_shardings_mem_kinds,
            compiled._no_kwargs,
            serialized_compile_options,
        )

View on GitHub (pinned to 24c60942c5)

Solutions

1. Route compilation through the repo's AOT compile helper (_compile_or_load) so the executable retains serialization support
2. Pin/verify the jaxlib version this code was developed against
3. Skip caching (pass no aot_cache_dir) when the executable is not AOT-produced

from jax._src import xla_extension as xc
exe = compiled._executable
if getattr(exe, '_unloaded_executable', None) is None:
    skip_caching = True  # not serializable

def is_serializable(compiled) -> bool:
    return getattr(getattr(compiled, '_executable', None), '_unloaded_executable', None) is not None

try:
    blob = partially_serialize(compiled, devices)
except ValueError:
    logger.warning('skipping AOT cache write; not serializable')

When it happens

Trigger: Calling partially_serialize on a Compiled object obtained from standard jax.jit(...).lower(...).compile() rather than the repo's AOT compile path, or after the executable was already fully loaded in a way that dropped the unloaded handle.

Common situations: Mixing normal jit compilation with the AOT caching utilities; upgrading jaxlib so the private _unloaded_executable attribute disappears; serializing a compiled function obtained from a checkpoint restored via a different code path.

Related errors

• Cannot remap XLA compile options device assignment: assignme
• execution_devices must be a list of xc.Device. Got: {executi
• Execution devices belong to a client other than `backend`. G
• AOT is not implemented for non-XLA lowerings.
• Cannot lower a compiled jit function.

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/0ecf71846d59c7cc. Report an issue: GitHub.