xai-org/x-algorithm · error · ValueError

execution_devices must be a list of xc.Device. Got: {executi

Error message

execution_devices must be a list of xc.Device. Got: {execution_devices}

What it means

_JaxPjrtPickler replaces device objects with index references during pickling; it requires every element of execution_devices to be a jaxlib xla_extension (xc.Device) instance. Passing numpy objects, strings, ints, or devices from another abstraction raises this ValueError before pickling starts.

Source

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

    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,
        )


class _JaxPjrtPickler(pickle.Pickler):
    def __init__(self, file, execution_devices: Sequence[xc.Device]):
        super().__init__(file)
        if not all(isinstance(d, xc.Device) for d in execution_devices):
            raise ValueError(
                f"execution_devices must be a list of xc.Device. Got: {execution_devices}"
            )
        self.device_to_index: dict[int, int] = {d.id: i for i, d in enumerate(execution_devices)}

    def persistent_id(self, obj):
        if isinstance(obj, xc.LoadedExecutable):
            return ("exec", obj.client.serialize_executable(obj))
        if isinstance(obj, xc._xla.Executable):
            return ("exec", obj.serialize())
        if isinstance(obj, xc.Device):
            if obj.id not in self.device_to_index:
                raise pickle.PicklingError(f"Unknown device: {obj}")
            return ("device", self.device_to_index[obj.id])
        if isinstance(obj, xc.Client):
            return ("client",)
        if is_host_callbacks(obj):
            return ("callbacks",)
        if isinstance(obj, xc.CompileOptions):

View on GitHub (pinned to 24c60942c5)

Solutions

1. Pass exactly the output of jax.devices() (or a slice of it) as execution_devices
2. Filter out non-Device entries before calling, e.g. [d for d in devices if isinstance(d, xc.Device)]
3. Fix upstream helpers that were returning device indices

Example fix

# before
execution_devices = [0, 1, 2, 3]
# after
import jax
execution_devices = jax.devices()

from jax._src import xla_extension as xc
assert all(isinstance(d, xc.Device) for d in execution_devices)

def all_xc_devices(devs: Sequence) -> bool:
    return all(isinstance(d, xc.Device) for d in devs)

When it happens

Trigger: Calling the AOT serialization path with execution_devices built from e.g. jax.devices() results mixed with None/int placeholders, or devices from a mock/stub in tests, or a plain list of device ids instead of Device objects.

Common situations: Test fixtures substituting fake devices; building the device list from device indices or platform strings; API changes where a helper starts returning device IDs rather than Device objects.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• `axis` should be an int, slice or iterable of ints.
• Compilation does not support serialization
• Execution devices belong to a client other than `backend`. G
• Unsupported type: {type(self.jitted)}.
• Symbols are not supported in strato JSON

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