google-research/timesfm · error · ValueError

Activation: {config.activation} not supported.

Error message

Activation: {config.activation} not supported.

What it means

The MLP block's __init__ only supports activation names 'relu', 'swish', and 'none' (mapping to jax.nn functions or identity). Any other config.activation string falls through to this ValueError. It is a config validation error raised at module construction time, before any forward pass.

Source

Thrown at src/timesfm/flax/dense.py:64

      in_features=config.hidden_dims,
      out_features=config.output_dims,
      use_bias=config.use_bias,
      rngs=rngs,
    )
    self.residual_layer = nnx.Linear(
      in_features=config.input_dims,
      out_features=config.output_dims,
      use_bias=config.use_bias,
      rngs=rngs,
    )
    if config.activation == "relu":
      self.activation = jax.nn.relu
    elif config.activation == "swish":
      self.activation = jax.nn.swish
    elif config.activation == "none":
      self.activation = lambda x: x
    else:
      raise ValueError(f"Activation: {config.activation} not supported.")

  def __call__(self, x: Float[Array, "b ... i"]) -> Float[Array, "b ... o"]:
    return self.output_layer(
      self.activation(self.hidden_layer(x))
    ) + self.residual_layer(x)


class RandomFourierFeatures(nnx.Module):
  """Random Fourier features layer."""

  __data__ = ("phrase_shifts",)

  def __init__(self, config: RandomFourierFeaturesConfig, *, rngs=nnx.Rngs(42)):
    self.config = config

    if config.output_dims % 4 != 0:
      raise ValueError(
        f"Output dims must be a multiple of 4: {config.output_dims} % 4 != 0."

View on GitHub (pinned to 331c6d33cb)

Solutions

1. Set config.activation to one of the supported values: 'relu', 'swish', or 'none'.
2. If you need another activation, add an elif branch mapping it to the corresponding jax.nn function in dense.py.
3. Check for casing/typo issues ('relu' vs 'ReLU') in the config source (YAML/JSON/CLI flag).

Example fix

// before
config = DenseConfig(activation="gelu")  # ValueError
// after
config = DenseConfig(activation="relu")  # or "swish" / "none"

SUPPORTED = {"relu", "swish", "none"}
if config.activation not in SUPPORTED:
    raise ValueError(f"activation must be one of {sorted(SUPPORTED)}, got {config.activation!r}")

def is_supported_activation(a: object) -> bool:
    return isinstance(a, str) and a in {"relu", "swish", "none"}

try:
    layer = DenseModule(config)
except ValueError as e:
    if "not supported" in str(e) and "Activation" in str(e):
        config.activation = "relu"
        layer = DenseModule(config)
    else:
        raise

When it happens

Trigger: Constructing the MLP/dense module with a config whose activation field is set to an unsupported string such as 'gelu', 'tanh', 'ReLU', 'sigmoid', or a None/typo value instead of 'relu', 'swish', or 'none'.

Common situations: Copying a config from another framework (e.g. HF transformers uses 'gelu' by default), typos or casing mistakes in the activation name, or hand-editing a config dict to an activation the TimesFM flax backend does not implement.

Related errors

• Output dims must be a multiple of 4: {config.output_dims} %
• Layer norm: {config.attention_norm} not supported.
• Layer norm: {config.feedforward_norm} not supported.
• Activation: {config.ff_activation} not supported.
• The embedding dims of the rotary position embeddingmust matc

AI-assisted analysis of google-research/timesfm@331c6d33cb (2026-08-29). Data as JSON: /api/errors/374d134c84839c1b. Report an issue: GitHub.