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

Code before:
        count = past_values.new_zeros(batch_size)
        mean = past_values.new_zeros(batch_size)
        std = past_values.new_zeros(batch_size)
        mean_history: list[torch.Tensor] = []
        std_history: list[torch.Tensor] = []

        for i in range(patched_inputs.shape[1]):
            count, mean, std = self._update_running_stats(
                count, mean, std, patched_inputs[:, i, :], patched_masks_bool[:, i, :]
            )
            mean_history.append(mean)
            std_history.append(std)

        if mean_history:
            context_mu = torch.stack(mean_history, dim=1)
            context_sigma = torch.stack(std_history, dim=1)
        else:
            context_mu = mean.unsqueeze(1)
            context_sigma = std.unsqueeze(1)

        normed_inputs = self._revin(patched_inputs, context_mu, context_sigma, reverse=False, mask=patched_masks_bool)

        tokenizer_inputs = torch.cat(
            [normed_inputs, patched_masks_bool.to(dtype=normed_inputs.dtype)],
            dim=-1,
        ).to(dtype=self.dtype)
        input_embeddings = self.input_ff_layer(tokenizer_inputs)

        patch_padding = patched_masks_bool[..., -1]

        sequence_length = input_embeddings.shape[1]
        num_masked = patch_padding.to(torch.int32).sum(dim=-1, keepdim=True)
        position_ids = torch.arange(sequence_length, device=input_embeddings.device).unsqueeze(0) - num_masked

        padding_mask = (~patch_padding).to(torch.int64)
        cache_position = torch.arange(sequence_length, device=input_embeddings.device)
        attention_mask = create_causal_mask(
            self.config, input_embeddings, padding_mask, cache_position, past_key_values=None
        )
        position_embeddings = self.rotary_emb(input_embeddings, position_ids)

        hidden_states = input_embeddings

        for layer in self.layers:
            hidden_states = layer(
                hidden_states,
                position_embeddings=position_embeddings,
                attention_mask=attention_mask,
                position_ids=position_ids,
                **kwargs,
            )

Review comment (suggestion):
This is a bit weird, can we achieve this in a different way? Would like to avoid using `self.dtype`

Code after:

        patched_inputs = past_values.view(batch_size, -1, patch_len)
        patched_masks = past_values_padding[:, :seq_len].view(batch_size, -1, patch_len)
        patched_masks_bool = patched_masks >= 0.5

        count = past_values.new_zeros(batch_size)
        mean = past_values.new_zeros(batch_size)
        std = past_values.new_zeros(batch_size)
        mean_history: list[torch.Tensor] = []
        std_history: list[torch.Tensor] = []

        for i in range(patched_inputs.shape[1]):
            count, mean, std = self._update_running_stats(
                count, mean, std, patched_inputs[:, i, :], patched_masks_bool[:, i, :]
            )
            mean_history.append(mean)
            std_history.append(std)

        if mean_history:
            context_mu = torch.stack(mean_history, dim=1)
            context_sigma = torch.stack(std_history, dim=1)
        else:
            context_mu = mean.unsqueeze(1)
            context_sigma = std.unsqueeze(1)

        normed_inputs = self._revin(patched_inputs, context_mu, context_sigma, reverse=False, mask=patched_masks_bool)

        tokenizer_inputs = torch.cat(
            [normed_inputs, patched_masks_bool.to(dtype=normed_inputs.dtype)],
            dim=-1,
        )
        input_embeddings = self.input_ff_layer(tokenizer_inputs)

        patch_padding = patched_masks_bool[..., -1]

        sequence_length = input_embeddings.shape[1]
        num_masked = patch_padding.to(torch.int32).sum(dim=-1, keepdim=True)
        position_ids = torch.arange(sequence_length, device=input_embeddings.device).unsqueeze(0) - num_masked

        padding_mask = (~patch_padding).to(torch.int64)
        cache_position = torch.arange(sequence_length, device=input_embeddings.device)
        attention_mask = create_causal_mask(
            self.config, input_embeddings, padding_mask, cache_position, past_key_values=None
        )
        position_embeddings = self.rotary_emb(input_embeddings, position_ids)

        hidden_states = input_embeddings

        for layer in self.layers:
            hidden_states = layer(
                hidden_states,