{"record":{"id":"f2f2aa26eb1c954e","repo":"xai-org/x-algorithm","slug":"cannot-remap-xla-compile-options-device-assignment","errorCode":null,"errorMessage":"Cannot remap XLA compile options device assignment: assignment expects {expected_devices} devices ({replica_count=} x {computation_count=}), but got {len(execution_devices)} execution devices.","messagePattern":"Cannot remap XLA compile options device assignment: assignment expects (.+?) devices \\((.+?) x (.+?)\\), but got (.+?) execution devices\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/utils/aot.py","lineNumber":115,"sourceCode":"    return compile_options\n\n\ndef remap_compile_options_device_assignment(\n    compile_options: xc.CompileOptions | None,\n    execution_devices: Sequence[xc.Device],\n) -> xc.CompileOptions | None:\n    if compile_options is None:\n        return None\n\n    device_assignment = compile_options.device_assignment\n    if device_assignment is None:\n        return compile_options\n\n    replica_count = device_assignment.replica_count()\n    computation_count = device_assignment.computation_count()\n    expected_devices = replica_count * computation_count\n    if len(execution_devices) != expected_devices:\n        raise ValueError(\n            \"Cannot remap XLA compile options device assignment: \"\n            f\"assignment expects {expected_devices} devices \"\n            f\"({replica_count=} x {computation_count=}), \"\n            f\"but got {len(execution_devices)} execution devices.\"\n        )\n\n    current_device_ids = np.asarray([d.id for d in execution_devices], dtype=np.int64)\n    compile_options.device_assignment = xc.DeviceAssignment.create(\n        current_device_ids.reshape(replica_count, computation_count, order=\"F\")\n    )\n    return compile_options\n\n\ndef load_compiled_from_serialized(\n    lowered: Lowered,\n    partially_serialized: PartiallySerialized,\n    execution_devices: Sequence[xc.Device],\n):","sourceCodeStart":97,"sourceCodeEnd":133,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/utils/aot.py#L97-L133","documentation":"When loading a serialized AOT-compiled executable onto a new set of devices, the stored XLA device assignment (replica_count x computation_count) must match the number of execution devices supplied. This ValueError is raised when remapping the compile options' device assignment because the device topology changed between compile time and load time.","triggerScenarios":"Calling load_compiled_from_serialized or persistent_load with an execution_devices list whose length differs from replica_count * computation_count of the compiled artifact, e.g. loading a cache compiled for 8 devices onto 4 devices.","commonSituations":"Sharing an AOT cache directory across machines with different GPU counts; changing JAX device count (CUDA_VISIBLE_DEVICES) between compile and load; single-host cache reused on multi-host jobs.","solutions":["Recompile (delete/invalidate the AOT cache entry) on the new topology so the device assignment matches","Ensure the same number of visible devices at load time as at compile time (e.g. same CUDA_VISIBLE_DEVICES)","Point aot_cache_dir at a per-topology cache directory keyed by device count"],"exampleFix":"# before\ncompiled = load_compiled_from_serialized(cache_dir, devices=devices_4)  # cache built for 8\n# after\ncompiled = compile_or_load(..., aot_cache_dir=cache_dir_for_this_topology)  # or clear stale cache","handlingStrategy":"fallback","validationCode":"n = jax.device_count()\n# only reuse a cache compiled for the same device count\ncache_key = f\"{aot_cache_dir}/devices_{n}\"\nif not (Path(cache_key) / 'meta.json').exists():\n    aot_cache_dir = None  # force fresh compile","typeGuard":null,"tryCatchPattern":"try:\n    exe = load_compiled_from_serialized(path, devices)\nexcept ValueError as e:\n    if 'execution devices' in str(e):\n        exe = compile_fresh(devices)  # fallback recompile\n    else:\n        raise","preventionTips":["Namespace AOT cache dirs by device count/topology","Keep CUDA_VISIBLE_DEVICES consistent between compile and load jobs","Never share one cache across heterogeneous machines"],"tags":["jax","xla","aot","device-topology","cache"],"backgroundTag":"device-count-mismatch","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}