xai-org/x-algorithm · error · NotImplementedError

Invalid attention implementation: {self.config.attn_impl}

Error message

Invalid attention implementation: {self.config.attn_impl}

What it means

_get_attn_impl maps config.attn_impl strings to attention classes via a match statement ('jax_attn', 'pallas_attn', 'flash_attn', 'pallas_ranker_attn', 'cutedsl_ranker_attn', ...). Any string not covered falls into `case _` and raises NotImplementedError with the invalid attn_impl value.

Source

Thrown at phoenix/xrex/models/layers.py:438

                    "LengthDistribution so block_sparse is constructed."
                )
                attn_class = CutedslRankerVarlenAttention

                def _add_h_axis(arr):
                    arr = jnp.asarray(arr, dtype=jnp.int32)
                    return jnp.broadcast_to(
                        jnp.expand_dims(arr, axis=1),
                        (arr.shape[0], self.config.num_q_heads) + arr.shape[1:],
                    )

                bs = seqpack_layout.block_sparse
                extra_attn_kwargs.update(
                    {f"bsp_{f.name}": _add_h_axis(getattr(bs, f.name)) for f in dc_fields(bs)}
                )
            case "cutedsl_ranker_attn":
                attn_class = CutedslRankerAttention
            case _:
                raise NotImplementedError(
                    f"Invalid attention implementation: {self.config.attn_impl}"
                )

        return attn_class, extra_attn_kwargs

    @hk.transparent
    def _out_projection(
        self,
        x: jax.Array,
        output_size: int,
        pspec: Optional[P] = None,
        lr_multiplier: float = 1.0,
        init_scale: float = 1.0,
        name: Optional[str] = None,
        rms_clip_axes=(-2, -1),
        with_bias: bool = False,
        w_init=None,
    ) -> jax.Array:

View on GitHub (pinned to 24c60942c5)

Solutions

Check the match arms in layers.py _get_attn_impl for the exact supported strings and use one.
Fix typos, casing, and stray whitespace in the config value.
If you added a new attention class, register it with its own match arm.

Example fix

# before
config.attn_impl = "pallas_rank_attn"

# after
config.attn_impl = "pallas_ranker_attn"

Defensive patterns

Strategy: validation

Validate before calling

SUPPORTED_ATTN_IMPLS = {"jax_attn", "pallas_attn", "flash_attn", "pallas_ranker_attn", "pallas_ranker_attn_infer", "cutedsl_ranker_attn"}
assert config.attn_impl in SUPPORTED_ATTN_IMPLS, config.attn_impl

Type guard

def is_valid_attn_impl(name: str) -> bool:
    return name in SUPPORTED_ATTN_IMPLS

Prevention

Fail fast on config load with the list of valid values.
Strip/normalize whitespace and case in config strings before use.

When it happens

Trigger: Typos in attn_impl (e.g. 'pallas_rank_attn', 'jax-attention'); requesting an impl that exists in another fork but not here; whitespace/case differences in the config string.

Common situations: See trigger scenarios.

Understand the failure class

Background: "Invalid value" and "allowed values are" config errors: what your library rejected and how to fix it — this error's family across 41 libraries.

Related errors

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