xai-org/x-algorithm · critical · NotImplementedError

use_async_emb requires an embedding optimizer implementing…

Error message

use_async_emb requires an embedding optimizer implementing AsyncEmbOptimizer, and {type(self).__name__} has no fused table update

What it means

When use_async_emb is enabled, the trainer calls gradient_update_start on the embedding optimizer to perform a fused embedding-table update. The default implementation raises NotImplementedError, so any optimizer class that does not implement AsyncEmbOptimizer's fused update fails here, with the class name embedded in the message.

Solutions

Use an optimizer class that implements AsyncEmbOptimizer (fused table update) for the embeddings when use_async_emb=True
Subclass and implement gradient_update_start (and gradient_update_done) on your optimizer
Disable use_async_emb if fused async embedding updates are not needed

Example fix

# before
cfg = TrainConfig(use_async_emb=True)  # optimizer = default Adam
# after
cfg = TrainConfig(use_async_emb=False)
# or: optimizer = AsyncEmbAdam(...)  # implements AsyncEmbOptimizer

Defensive patterns

Strategy: type-guard

Validate before calling

from phoenix.xrex.optimizers.recsys import async_emb
if cfg.use_async_emb:
    assert isinstance(emb_opt, async_emb.AsyncEmbOptimizer)

Type guard

def supports_async_emb(opt) -> bool:
    return (hasattr(opt, 'gradient_update_start')
            and type(opt).gradient_update_start is not AsyncEmbGradientUpdate.gradient_update_start)

Try / catch

try:
    upd = emb_opt.gradient_update_start(ctx, update, table, state, gate)
except NotImplementedError:
    raise ConfigError('use_async_emb=True requires an AsyncEmbOptimizer') from None

Prevention

Gate use_async_emb on an isinstance check in config validation
Keep a registry of async-emb-capable optimizers

When it happens

Trigger: Enabling use_async_emb in the recsys training config while using a plain optimizer (e.g. the standard Adam wrapper) for the embedding tables; instantiating AsyncEmbGradientUpdate and calling gradient_update_start on a subclass that only overrides gradient_update_done.

Common situations: Turning on the async embedding update feature of xrex without swapping the embedding optimizer for an AsyncEmbOptimizer implementation; upgrading where the async-emb API gained new required methods.

Related errors

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

Appendix: source

Thrown at phoenix/xrex/optimizers/recsys/async_emb_gradient_update.py:32

class AsyncEmbGradientUpdate(NamedTuple):
    unique_tokens: jax.Array
    grads: jax.Array
    segment_ids: jax.Array
    pending: jax.Array


@runtime_checkable
class AsyncEmbOptimizer(Protocol):
    def gradient_update_start(
        self,
        context: async_emb.AsyncEmbContextHandle,
        update: AsyncEmbGradientUpdate,
        table: jax.Array,
        state: Any,
        gate: jax.Array,
    ) -> tuple[tuple[jax.Array, ...], jax.Array, Any, dict[str, jax.Array]]:
        raise NotImplementedError(
            f"use_async_emb requires an embedding optimizer implementing "
            f"AsyncEmbOptimizer, and {type(self).__name__} has no fused table update"
        )

    def gradient_update_done(
        self, context: async_emb.AsyncEmbContextHandle, gate: jax.Array
    ) -> tuple[jax.Array, jax.Array, jax.Array]:
        raise NotImplementedError(
            f"use_async_emb requires an embedding optimizer implementing "
            f"AsyncEmbOptimizer, and {type(self).__name__} has no fused table update"
        )

View on GitHub (pinned to 24c60942c5)