xai-org/x-algorithm · error · NotImplementedError

Please override this method for specific attention impl.

Error message

Please override this method for specific attention impl.

What it means

Attention.call_attn is the abstract dispatch method of the Attention base class in xrex.models.attention; the base implementation unconditionally raises NotImplementedError. Every concrete attention subclass (JaxAttention, PallasAttention, CustomAttention, ...) must override call_attn, and __call__ routes into it, so hitting this means a raw or incomplete Attention subclass was instantiated and invoked.

Source

Thrown at phoenix/xrex/models/attention.py:76

        value: Optional[jax.Array],
        segment_ids: Optional[jax.Array],
        segment_ids_k: Optional[jax.Array],
        temp: Optional[jax.Array],
        **kwargs,
    ):
        return self.call_attn(query, key, value, segment_ids, segment_ids_k, temp, **kwargs)

    def call_attn(
        self,
        query: jax.Array,
        key: Optional[jax.Array],
        value: Optional[jax.Array],
        segment_ids: Optional[jax.Array],
        segment_ids_k: Optional[jax.Array],
        temp: Optional[jax.Array],
        **kwargs,
    ):
        raise NotImplementedError("Please override this method for specific attention impl.")


class CustomAttention(Attention):
    def sharded_custom_op_with_extra_args(
        self, **kwargs
    ) -> tuple[Callable, dict[str, Callable[[str], NamedShape]]]:
        raise NotImplementedError

    def call_attn(self, query, key, value, segment_ids, segment_ids_k, temp, **kwargs):
        assert self.sharding_context is not None

        body_fn, extra_arg_shape_fns = self.sharded_custom_op_with_extra_args(**kwargs)
        sharding_rule = functools.partial(
            self.sharding_context.sharding_specs, namespace=self.config.attn_sharding_namespace
        )

        extra_args = []
        extra_arg_shapes = []

View on GitHub (pinned to 24c60942c5)

Solutions

1. Override call_attn(self, query, key, value, segment_ids, segment_ids_k, temp, **kwargs) in your subclass.
2. Check the subclass actually inherits from the intended base and the method signature matches.
3. Make the base class abstract (abc.abstractmethod) so instantiation fails earlier with a clearer error.

Example fix

# before
class MyAttn(Attention):
    def attention(self, q, k, v, **kw): ...

# after
class MyAttn(Attention):
    def call_attn(self, query, key, value, segment_ids, segment_ids_k, temp, **kwargs): ...

from xrex.models.attention import Attention
assert type(attn).call_attn is not Attention.call_attn, "call_attn not overridden"

def has_call_attn_override(cls) -> bool:
    return cls.call_attn.__func__ is not Attention.__dict__["call_attn"]

When it happens

Trigger: Instantiating Attention (or a subclass that forgot to override call_attn) and calling it via __call__ with query/key/value arrays; registering a custom attention class whose method is named differently (e.g. call_attention) so the base stub runs.

Common situations: Adding a new attention kernel and misspelling the override method name; refactoring that renames call_attn in one subclass but not another; passing the base class as attn_class by accident in _get_attn_impl.

Related errors

• Mask dimensionality {mask.ndim} must match logits dimensiona
• Unknown remat policy: {policy}
• use_async_emb requires an embedding optimizer implementing A
• Invalid backward pass implementation: {backward_pass_impl}
• cap_method must be in [tanh, soft_sign, none], got {cap_meth

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