xai-org/x-algorithm · error · NotImplementedError

unknown optim={optim!r}

Error message

unknown optim={optim!r}

What it means

The Optimizer config's make() resolves the optimizer name through _OPTIM_ALIASES and only accepts (aliases of) 'adam'. Any other name raises NotImplementedError, telling you this entry point currently supports only the AdamW path.

Source

Thrown at phoenix/xrex/optimizers/optim.py:139

    return optax.GradientTransformation(init_empty_state, update_fn)


_OPTIM_ALIASES: dict[str, str] = {}


@configclass
class OptimConfig(Config):
    optim: str
    learning_rate: BaseSchedule | float = 1.0
    b1: float = 0.9
    b2: float = 0.99
    weight_decay: float = 0.0
    clip_by_global_norm: float = 1.0

    def make(self):
        optim = _OPTIM_ALIASES.get(self.optim, self.optim)
        if optim not in ("adam",):
            raise NotImplementedError(f"unknown optim={optim!r}")

        @inject_hyperparams
        def schedule_optim(learning_rate, b1, b2, weight_decay):
            core = optax.adamw(learning_rate, b1=b1, b2=b2, weight_decay=weight_decay)
            return optax.chain(
                optax.clip_by_global_norm(self.clip_by_global_norm),
                core,
                scale_by_lr_multiplier(),
            )

        schedule_args = map(_instantiate, (self.learning_rate, self.b1, self.b2, self.weight_decay))
        return schedule_optim(*schedule_args)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Use 'adam' (or an alias present in _OPTIM_ALIASES) since only the AdamW path is implemented
2. If you need another optimizer, build the optax transformation directly and pass it in instead of using the string config
3. Add your optimizer to _OPTIM_ALIASES and extend make() with a matching branch

Example fix

# before
optim_cfg = OptimizerConfig(optim="adafactor")
# after
optim_cfg = OptimizerConfig(optim="adam")

from phoenix.xrex.optimizers.optim import _OPTIM_ALIASES
assert cfg.optim in _OPTIM_ALIASES or cfg.optim == 'adam', 'only adam supported'

def is_supported_optim(name: str) -> bool:
    return _OPTIM_ALIASES.get(name, name) == 'adam'

try:
    opt = cfg.make()
except NotImplementedError:
    opt = my_custom_optax_optimizer()  # fallback

When it happens

Trigger: Setting config.optim = 'sgd', 'adafactor', 'lion', etc. and calling make() (typically via _instantiate during trainer construction).

Common situations: Porting training configs from other frameworks expecting a string optimizer registry; assuming a generic optimizer enum exists in xrex.

Related errors

• Unknown remat policy: {policy}
• `static_args` must specify a subset of `inner_factory`'s par
• use_async_emb requires an embedding optimizer implementing A
• Invalid argument {arg!r}, not a key=value replacement and no
• Expected bool [True, true, False, false], got {val!r}

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