google-research/timesfm · error · ValueError

Activation: {config.ff_activation} not supported.

Error message

Activation: {config.ff_activation} not supported.

What it means

The feed-forward network only supports activation names 'relu', 'swish', and 'none'; any other config.ff_activation value falls through to this ValueError in __init__. Same pattern as the dense-layer activation check but keyed on ff_activation.

Source

Thrown at src/timesfm/flax/transformer.py:336

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

  def __call__(
    self,
    input_embeddings: Float[Array, "b n d"],
    patch_mask: Bool[Array, "b n"],
    decode_cache: DecodeCache | None = None,
  ) -> tuple[Float[Array, "b n d"], DecodeCache | None]:
    attn_output, decode_cache = self.attn(
      inputs_q=self.pre_attn_ln(input_embeddings),
      decode_cache=decode_cache,
      patch_mask=patch_mask,
      sow_weights=False,
      deterministic=True,
    )
    attn_output = self.post_attn_ln(attn_output) + input_embeddings
    output_embeddings = (
      self.post_ff_ln(self.ff1(self.activation(self.ff0(self.pre_ff_ln(attn_output)))))
      + attn_output

View on GitHub (pinned to 331c6d33cb)

Solutions

1. Set config.ff_activation to 'relu', 'swish', or 'none'.
2. To use another activation, add an elif branch mapping it to the jax.nn function in transformer.py.
3. Check the config file/default dict for the injected value when the key is omitted.

Example fix

// before
config = TransformerConfig(ff_activation="gelu")  # ValueError
// after
config = TransformerConfig(ff_activation="swish")

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

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

try:
    block = TransformerBlock(config)
except ValueError as e:
    if "ff_activation" in str(e) or "Activation" in str(e):
        config.ff_activation = "swish"
        block = TransformerBlock(config)
    else:
        raise

When it happens

Trigger: Constructing the transformer with config.ff_activation set to 'gelu', 'tanh', 'elu', 'gelu_new', or None instead of 'relu', 'swish', or 'none'.

Common situations: Porting configs from HF/other frameworks whose default FF activation is 'gelu'; casing mistakes ('ReLU'); missing field causing a None default.

Related errors

• Activation: {config.activation} not supported.
• 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.
• Memory dimension ({self.qkv_features}) must be divisible by

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