jax-ml/jax · error · ValueError

Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'

Error message

Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'.

What it means

jax.nn.standardize supports only two variance algorithm strings: 'fast' (uses the less accurate E[x^2]-E[x]^2 formula) and 'stable' (two-pass variance). Any other string for the `algorithm` parameter raises this ValueError.

Source

Thrown at jax/_src/nn/functions.py:710

  if mean is None:
    mean = jnp.mean(x, axis, keepdims=True, where=where)
  if variance is None:
    if algorithm == "stable":
      variance = jnp.mean(
          jnp.square(jnp.subtract(x, mean)), axis, keepdims=True, where=where)
    elif algorithm == "fast":
      # This definition is traditionally seen as less accurate than the
      # two-pass mean((x - mean(x))**2) but may be faster and even, given
      # typical activation distributions and low-precision arithmetic, more
      # accurate when used in neural network normalization layers.
      variance = jnp.mean(
          jnp.square(x), axis, keepdims=True, where=where) - jnp.square(mean)
      # Because we're using a less accurate variance definition, it may
      # return negative values. This is problematic for the rsqrt, so we
      # clip to 0.
      variance = jnp.clip(variance, 0)
    else:
      raise ValueError(
          f"Unknown algorithm '{algorithm}'. Expected 'fast' or 'stable'.")
  return jnp.subtract(x, mean) * lax.rsqrt(variance + epsilon)

# TODO(slebedev): Change the type of `x` to `ArrayLike`.
@api.jit(static_argnames=("num_classes", "dtype", "axis", "out_sharding"))
def _one_hot(x: Array, num_classes: int, *,
             dtype: DTypeLike, axis: int | AxisName,
             out_sharding: NamedSharding | None) -> Array:
  num_classes = core.concrete_dim_or_error(
      num_classes,
      "The error arose in jax.nn.one_hot argument `num_classes`.")
  try:
    out_axis = util.canonicalize_axis(axis, x.ndim + 1)  # pyrefly: ignore[bad-argument-type]
  except TypeError:
    axis_size = lax.axis_size(axis)
    if num_classes != axis_size:
      raise ValueError(f"Expected num_classes to match the size of axis {axis}, "
                       f"but {num_classes} != {axis_size}") from None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass algorithm='fast' or algorithm='stable' exactly (lowercase)
If you wanted unbiased/biased variance behavior, pick 'stable' for the numerically safer two-pass computation

Example fix

// before
y = jax.nn.standardize(x, algorithm='unbiased')

// after
y = jax.nn.standardize(x, algorithm='stable')

Defensive patterns

Strategy: validation

Validate before calling

import jax.numpy as jnp

def standardize(x, algorithm):
    if algorithm not in ('fast', 'stable'):
        raise ValueError(f"algorithm must be 'fast' or 'stable', got {algorithm!r}")
    return jax.nn.standardize(x, algorithm=algorithm)

Type guard

def is_standardize_algorithm(a) -> bool: return a in ('fast', 'stable')

Prevention

Use Literal['fast','stable'] type hints for wrapper functions
Centralize algorithm strings as module constants instead of free-form strings

When it happens

Trigger: Calling jax.nn.standardize(x, algorithm='unbiased'), algorithm='population', algorithm=None, or a typo like 'Stable' (case-sensitive).

Common situations: Porting torch/TF normalization code where variance parameter names differ (torch uses unbiased=True/False); passing None expecting a default; version drift if custom algorithm names were removed/renamed.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/a524c69f03f9b4ed. Report an issue: GitHub.