{"record":{"id":"d0e0a9fda59ed634","repo":"xai-org/x-algorithm","slug":"aot-is-not-implemented-for-non-xla-lowerings","errorCode":null,"errorMessage":"AOT is not implemented for non-XLA lowerings.","messagePattern":"AOT is not implemented for non-XLA lowerings\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/utils/aot.py","lineNumber":355,"sourceCode":"        ]\n    )\n\n\ndef get_cache_key(lowered: Lowered, extra_info: list[str] = None, add_loc: bool = False):\n    hash_obj = hashlib.sha256()\n    _hash_string(hash_obj, AOT_CACHE_VERSION)\n    environment = get_environment_info()\n    _hash_string(hash_obj, environment)\n\n    _hash_xla_flags(hash_obj, [])\n    sanitized_ir = get_sanitized_ir_text(lowered, add_loc)\n    _hash_string(hash_obj, sanitized_ir)\n    _hash_string(hash_obj, str(lowered.args_info))\n    _hash_string(hash_obj, str(lowered.out_tree))\n\n    lowering = lowered._lowering\n    if not isinstance(lowering, jax.interpreters.pxla.MeshComputation):\n        raise NotImplementedError(\"AOT is not implemented for non-XLA lowerings.\")\n\n    compile_args = lowering.compile_args\n    if \"global_in_avals\" in compile_args:\n        _hash_string(hash_obj, str(compile_args[\"global_in_avals\"]))\n    if \"global_out_avals\" in compile_args:\n        _hash_string(hash_obj, str(compile_args[\"global_out_avals\"]))\n    if \"in_shardings\" in compile_args:\n        _hash_string(hash_obj, str(compile_args[\"in_shardings\"]))\n    if \"out_shardings\" in compile_args:\n        _hash_string(hash_obj, str(compile_args[\"out_shardings\"]))\n    if \"in_layouts\" in compile_args:\n        _hash_string(hash_obj, str(compile_args[\"in_layouts\"]))\n    if \"out_layouts\" in compile_args:\n        _hash_string(hash_obj, str(compile_args[\"out_layouts\"]))\n\n    additional_content = \":\".join(extra_info) if extra_info else None\n    if additional_content:\n        _hash_string(hash_obj, additional_content)","sourceCodeStart":337,"sourceCodeEnd":373,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/utils/aot.py#L337-L373","documentation":"The AOT cache key builder hashes XLA-specific internals of the lowering (MeshComputation compile_args). If the lowered object is not a pxla.MeshComputation — i.e. it came from a non-XLA backend or a different lowering pipeline — it cannot be hashed for AOT caching and NotImplementedError is raised.","triggerScenarios":"Running the AOT compile/cache path on a non-XLA backend (e.g. a plugin backend or eager/debug interpreter), or on a jax version where the private lowered._lowering type changed so the isinstance check fails.","commonSituations":"Debugging with jax disabled/XLA off; new jax release renaming pxla.MeshComputation; custom backends (iree, tpu-plugin variants) that don't use MeshComputation.","solutions":["Run with the standard XLA backend (don't set JAX_PLATFORMS/disable-XLA flags) when using the AOT cache","Pin the jax version the AOT code supports until the private-API check is updated","Bypass caching (skip compile_or_load's cache path) for non-XLA backends"],"exampleFix":null,"handlingStrategy":"try-catch","validationCode":"import jax.interpreters.pxla as pxla\nok = isinstance(lowered._lowering, pxla.MeshComputation)","typeGuard":"def is_xla_lowering(lowered) -> bool:\n    import jax.interpreters.pxla as pxla\n    return isinstance(lowered._lowering, pxla.MeshComputation)","tryCatchPattern":"try:\n    key = get_cache_key(lowered)\nexcept NotImplementedError:\n    key = None  # skip AOT cache, compile directly","preventionTips":["Don't disable XLA or use plugin backends with the AOT cache","Pin jax; private _lowering/pxla internals change between versions","Make cache usage optional so unsupported lowerings fall back to normal compile"],"tags":["jax","xla","aot","not-implemented","backend"],"backgroundTag":"unsupported-backend-operation","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}