xai-org/x-algorithm · error · ValueError

`axis` should be an int, slice or iterable of ints.

Error message

`axis` should be an int, slice or iterable of ints.

What it means

RMSNorm/LayerNorm-style __init__ normalizes the axis argument: it accepts a slice, an int, or an iterable whose elements are all ints, storing it as a tuple. Anything else (a float, a string, a single-element non-int iterable like [1.5] or np.array with float dtype) raises ValueError with the accepted forms.

Source

Thrown at phoenix/xrex/models/normalization.py:79

        self,
        axis: Union[int, Sequence[int], slice],
        eps: float = 1e-5,
        scale_init: Optional[hk.initializers.Initializer] = None,
        name: Optional[str] = None,
        create_scale: bool = True,
        pspec: Optional[P] = P(None),
        lr_multiplier: float = 1.0,
        weight_decay_mask: float = 0.0,
    ):
        super().__init__(name=name)
        if isinstance(axis, slice):
            self.axis = axis
        elif isinstance(axis, int):
            self.axis = (axis,)
        elif isinstance(axis, abc.Iterable) and all(isinstance(ax, int) for ax in axis):
            self.axis = tuple(axis)
        else:
            raise ValueError("`axis` should be an int, slice or iterable of ints.")

        self.eps = eps
        self.create_scale = create_scale
        self.reparameterize = weight_decay_mask > 0
        if scale_init is None:
            scale_init = jnp.zeros if self.reparameterize else jnp.ones
        if self.reparameterize and scale_init is not jnp.zeros:
            raise ValueError(
                "RMSNorm: when weight_decay_mask > 0 the layer is "
                "re-parameterized as (1 + scale) * x and `scale_init` must be "
                f"jnp.zeros (got {scale_init!r}). Pass scale_init=jnp.zeros or "
                "leave it as None to use the default."
            )
        self.scale_init = scale_init
        self.pspec = pspec
        self.lr_multiplier = lr_multiplier
        self.weight_decay_mask = weight_decay_mask

View on GitHub (pinned to 24c60942c5)

Solutions

1. Pass axis as int or an iterable of ints, e.g. axis=-1 or axis=(1, 2).
2. Coerce numpy scalars with int(ax) before constructing.
3. Validate config-loaded axis values against this contract in your config parser.

Example fix

# before
norm = Norm(axis=np.float64(-1.0), ...)

# after
norm = Norm(axis=-1, ...)

axis = tuple(int(a) for a in axis) if not isinstance(axis, int) else axis

def is_valid_axis(axis) -> bool:
    return isinstance(axis, (int, slice)) or (
        isinstance(axis, abc.Iterable) and all(isinstance(a, int) for a in axis)
    )

When it happens

Trigger: Constructing the normalization layer with axis=1.0, axis="features", axis=[0, 1.0], or an unhashable/mixed iterable.

Common situations: Programmatically computed axis values that arrive as floats or numpy scalars; configs that store axis as a string.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• execution_devices must be a list of xc.Device. Got: {executi
• Unsupported type: {type(self.jitted)}.
• Invalid message type: {type(message)=}
• async_emb axis {axis!r} is not a mesh axis of {mesh}
• async_emb requires token shards to vary across the communica

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