Repository: huggingface/transformers
File: src/transformers/models/timesfm2_5/modular_timesfm2_5.py (Python)

Code before:
    pass


class TimesFm2_5MLP(CLIPMLP):
    def __init__(self, config: TimesFm2_5Config):
        super().__init__()
        self.activation_fn = ACT2FN[config.activation]
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.use_bias)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.use_bias)


class TimesFm2_5ResidualBlock(TimesFmResidualBlock):
    """[`TimesFmResidualBlock`] variant with configurable `use_bias` and `activation`."""

    def __init__(self, config, input_dims: int, hidden_dims: int, output_dims: int, use_bias: bool | None = None):
        super().__init__(input_dims, hidden_dims, output_dims)
        use_bias = use_bias if use_bias is not None else config.use_bias
        self.input_layer = nn.Linear(input_dims, hidden_dims, bias=use_bias)
        self.output_layer = nn.Linear(hidden_dims, output_dims, bias=use_bias)
        self.residual_layer = nn.Linear(input_dims, output_dims, bias=use_bias)
        self.activation = ACT2FN[config.activation]

    def forward(self, x):
        # Align activations to block parameter dtype for mixed precision stability.
        if x.dtype != self.input_layer.weight.dtype:
            x = x.to(self.input_layer.weight.dtype)
        return super().forward(x)


class TimesFm2_5RMSNorm(LlamaRMSNorm):
    pass


class TimesFm2_5RotaryEmbedding(LlamaRotaryEmbedding):
    pass


class TimesFm2_5Attention(ApertusAttention):
    """TimesFM 2.5 attention with learnable per-dimension query scaling."""

    def __init__(self, config: TimesFm2_5Config, layer_idx: int):
        super().__init__(config, layer_idx)
        self.scaling = nn.Parameter(torch.empty((self.head_dim,)))

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        past_key_values=None,
        cache_position: torch.LongTensor | None = None,

Review comment (suggestion):
Should be a no op in most times and we don't need to check ourselves?
```suggestion
        # Align activations to block parameter dtype for mixed precision stability
        x = x.to(self.input_layer.weight.dtype)
```

Code after:
    pass


class TimesFm2_5MLP(CLIPMLP):
    def __init__(self, config: TimesFm2_5Config):
        super().__init__()
        self.activation_fn = ACT2FN[config.activation]
        self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size, bias=config.use_bias)
        self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size, bias=config.use_bias)


class TimesFm2_5ResidualBlock(TimesFmResidualBlock):
    """[`TimesFmResidualBlock`] variant with configurable `use_bias` and `activation`."""

    def __init__(self, config, input_dims: int, hidden_dims: int, output_dims: int, use_bias: bool | None = None):
        super().__init__(input_dims, hidden_dims, output_dims)
        use_bias = use_bias if use_bias is not None else config.use_bias
        self.input_layer = nn.Linear(input_dims, hidden_dims, bias=use_bias)
        self.output_layer = nn.Linear(hidden_dims, output_dims, bias=use_bias)
        self.residual_layer = nn.Linear(input_dims, output_dims, bias=use_bias)
        self.activation = ACT2FN[config.activation]

    def forward(self, x):
        # Align activations to block parameter dtype for mixed precision stability
        x = x.to(self.input_layer.weight.dtype)
        return super().forward(x)


class TimesFm2_5RMSNorm(LlamaRMSNorm):
    pass


class TimesFm2_5RotaryEmbedding(LlamaRotaryEmbedding):
    pass


class TimesFm2_5Attention(ApertusAttention):
    """TimesFM 2.5 attention with learnable per-dimension query scaling."""

    def __init__(self, config: TimesFm2_5Config, layer_idx: int):
        super().__init__(config, layer_idx)
        self.scaling = nn.Parameter(torch.empty((self.head_dim,)))

    def forward(
        self,
        hidden_states: torch.Tensor,
        position_embeddings: tuple[torch.Tensor, torch.Tensor],
        attention_mask: torch.Tensor | None,
        past_key_values=None,
        cache_position: torch.LongTensor | None = None,
        **kwargs: Unpack[TransformersKwargs],