xai-org/x-algorithm · error · NotImplementedError

AOT is not implemented for non-XLA lowerings.

Error message

AOT is not implemented for non-XLA lowerings.

What it means

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.

Source

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

        ]
    )


def get_cache_key(lowered: Lowered, extra_info: list[str] = None, add_loc: bool = False):
    hash_obj = hashlib.sha256()
    _hash_string(hash_obj, AOT_CACHE_VERSION)
    environment = get_environment_info()
    _hash_string(hash_obj, environment)

    _hash_xla_flags(hash_obj, [])
    sanitized_ir = get_sanitized_ir_text(lowered, add_loc)
    _hash_string(hash_obj, sanitized_ir)
    _hash_string(hash_obj, str(lowered.args_info))
    _hash_string(hash_obj, str(lowered.out_tree))

    lowering = lowered._lowering
    if not isinstance(lowering, jax.interpreters.pxla.MeshComputation):
        raise NotImplementedError("AOT is not implemented for non-XLA lowerings.")

    compile_args = lowering.compile_args
    if "global_in_avals" in compile_args:
        _hash_string(hash_obj, str(compile_args["global_in_avals"]))
    if "global_out_avals" in compile_args:
        _hash_string(hash_obj, str(compile_args["global_out_avals"]))
    if "in_shardings" in compile_args:
        _hash_string(hash_obj, str(compile_args["in_shardings"]))
    if "out_shardings" in compile_args:
        _hash_string(hash_obj, str(compile_args["out_shardings"]))
    if "in_layouts" in compile_args:
        _hash_string(hash_obj, str(compile_args["in_layouts"]))
    if "out_layouts" in compile_args:
        _hash_string(hash_obj, str(compile_args["out_layouts"]))

    additional_content = ":".join(extra_info) if extra_info else None
    if additional_content:
        _hash_string(hash_obj, additional_content)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Run with the standard XLA backend (don't set JAX_PLATFORMS/disable-XLA flags) when using the AOT cache
2. Pin the jax version the AOT code supports until the private-API check is updated
3. Bypass caching (skip compile_or_load's cache path) for non-XLA backends

import jax.interpreters.pxla as pxla
ok = isinstance(lowered._lowering, pxla.MeshComputation)

def is_xla_lowering(lowered) -> bool:
    import jax.interpreters.pxla as pxla
    return isinstance(lowered._lowering, pxla.MeshComputation)

try:
    key = get_cache_key(lowered)
except NotImplementedError:
    key = None  # skip AOT cache, compile directly

When it happens

Trigger: 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.

Common situations: Debugging with jax disabled/XLA off; new jax release renaming pxla.MeshComputation; custom backends (iree, tpu-plugin variants) that don't use MeshComputation.

Related errors

• Cannot remap XLA compile options device assignment: assignme
• Please override this method for specific attention impl.
• Unknown remat policy: {policy}
• use_async_emb requires an embedding optimizer implementing A
• Compilation does not support serialization

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