Add files using upload-large-folder tool

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/encoder_decoder/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_encoder_decoder import *
+    from .modeling_encoder_decoder import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/encoder_decoder/configuration_encoder_decoder.py

@@ -0,0 +1,102 @@
+# Copyright 2020 The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from huggingface_hub.dataclasses import strict
+
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+from ..auto import AutoConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="")
+@strict
+class EncoderDecoderConfig(PreTrainedConfig):
+    r"""
+    Examples:
+
+    ```python
+    >>> from transformers import BertConfig, EncoderDecoderConfig, EncoderDecoderModel
+
+    >>> # Initializing a BERT google-bert/bert-base-uncased style configuration
+    >>> config_encoder = BertConfig()
+    >>> config_decoder = BertConfig()
+
+    >>> config = EncoderDecoderConfig.from_encoder_decoder_configs(config_encoder, config_decoder)
+
+    >>> # Initializing a Bert2Bert model (with random weights) from the google-bert/bert-base-uncased style configurations
+    >>> model = EncoderDecoderModel(config=config)
+
+    >>> # Accessing the model configuration
+    >>> config_encoder = model.config.encoder
+    >>> config_decoder = model.config.decoder
+    >>> # set decoder config to causal lm
+    >>> config_decoder.is_decoder = True
+    >>> config_decoder.add_cross_attention = True
+
+    >>> # Saving the model, including its configuration
+    >>> model.save_pretrained("my-model")
+
+    >>> # loading model and config from pretrained folder
+    >>> encoder_decoder_config = EncoderDecoderConfig.from_pretrained("my-model")
+    >>> model = EncoderDecoderModel.from_pretrained("my-model", config=encoder_decoder_config)
+    ```"""
+
+    model_type = "encoder-decoder"
+    sub_configs = {"encoder": AutoConfig, "decoder": AutoConfig}
+    has_no_defaults_at_init = True
+
+    pad_token_id: int | None = None
+    decoder_start_token_id: int | None = None
+    is_encoder_decoder: bool | None = True
+
+    def __post_init__(self, **kwargs):
+        if "encoder" not in kwargs or "decoder" not in kwargs:
+            raise ValueError(
+                f"A configuration of type {self.model_type} cannot be instantiated because not both `encoder` and"
+                f" `decoder` sub-configurations are passed, but only {kwargs}"
+            )
+
+        encoder_config = kwargs.pop("encoder")
+        encoder_model_type = encoder_config.pop("model_type")
+        decoder_config = kwargs.pop("decoder")
+        decoder_model_type = decoder_config.pop("model_type")
+
+        self.encoder = AutoConfig.for_model(encoder_model_type, **encoder_config)
+        self.decoder = AutoConfig.for_model(decoder_model_type, **decoder_config)
+        super().__post_init__(**kwargs)
+
+    @classmethod
+    def from_encoder_decoder_configs(
+        cls, encoder_config: PreTrainedConfig, decoder_config: PreTrainedConfig, **kwargs
+    ) -> PreTrainedConfig:
+        r"""
+        Instantiate a [`EncoderDecoderConfig`] (or a derived class) from a pre-trained encoder model configuration and
+        decoder model configuration.
+
+        Returns:
+            [`EncoderDecoderConfig`]: An instance of a configuration object
+        """
+        logger.info("Set `config.is_decoder=True` and `config.add_cross_attention=True` for decoder_config")
+        decoder_config.is_decoder = True
+        decoder_config.add_cross_attention = True
+
+        return cls(encoder=encoder_config.to_dict(), decoder=decoder_config.to_dict(), **kwargs)
+
+
+__all__ = ["EncoderDecoderConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/encoder_decoder/modeling_encoder_decoder.py

@@ -0,0 +1,471 @@
+# Copyright 2018 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Classes to support Encoder-Decoder architectures"""
+
+import inspect
+import warnings
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from ...cache_utils import Cache
+from ...configuration_utils import PreTrainedConfig
+from ...generation import GenerationMixin
+from ...modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
+from ...modeling_utils import PreTrainedModel
+from ...utils import auto_docstring, logging
+from ...utils.generic import can_return_tuple
+from ..auto.configuration_auto import AutoConfig
+from ..auto.modeling_auto import AutoModel, AutoModelForCausalLM
+from .configuration_encoder_decoder import EncoderDecoderConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+DEPRECATION_WARNING = (
+    "Version v4.12.0 introduces a better way to train encoder-decoder models by computing the loss inside the"
+    " encoder-decoder framework rather than in the decoder itself. You may observe training discrepancies if"
+    " fine-tuning a model trained with versions anterior to 4.12.0. The decoder_input_ids are now created based on the"
+    " labels, no need to pass them yourself anymore."
+)
+
+
+def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):
+    """
+    Shift input ids one token to the right.
+    """
+    shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+    shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()
+    if decoder_start_token_id is None:
+        raise ValueError("Make sure to set the decoder_start_token_id attribute of the model's configuration.")
+    shifted_input_ids[:, 0] = decoder_start_token_id
+
+    if pad_token_id is None:
+        raise ValueError("Make sure to set the pad_token_id attribute of the model's configuration.")
+    # replace possible -100 values in labels by `pad_token_id`
+    shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)
+
+    return shifted_input_ids
+
+
+@auto_docstring
+class EncoderDecoderModel(PreTrainedModel, GenerationMixin):
+    r"""
+    [`EncoderDecoderModel`] is a generic model class that will be instantiated as a transformer architecture with one
+    of the base model classes of the library as encoder and another one as decoder when created with the
+    :meth*~transformers.AutoModel.from_pretrained* class method for the encoder and
+    :meth*~transformers.AutoModelForCausalLM.from_pretrained* class method for the decoder.
+    """
+
+    config: EncoderDecoderConfig
+    base_model_prefix = "encoder_decoder"
+    main_input_name = "input_ids"
+    supports_gradient_checkpointing = True
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    def __init__(
+        self,
+        config: PreTrainedConfig | None = None,
+        encoder: PreTrainedModel | None = None,
+        decoder: PreTrainedModel | None = None,
+    ):
+        r"""
+        encoder (`PreTrainedModel`, *optional*):
+            The encoder model to use.
+        decoder (`PreTrainedModel`, *optional*):
+            The decoder model to use.
+        """
+        if config is None and (encoder is None or decoder is None):
+            raise ValueError("Either a configuration or an encoder and a decoder has to be provided.")
+        if config is None:
+            config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config)
+        else:
+            if not isinstance(config, self.config_class):
+                raise ValueError(f"Config: {config} has to be of type {self.config_class}")
+
+        if getattr(config.decoder, "cross_attention_hidden_size", None) is not None:
+            if config.decoder.cross_attention_hidden_size != config.encoder.hidden_size:
+                raise ValueError(
+                    "If `cross_attention_hidden_size` is specified in the decoder's configuration, it has to be equal"
+                    f" to the encoder's `hidden_size`. Got {config.decoder.cross_attention_hidden_size} for"
+                    f" `config.decoder.cross_attention_hidden_size` and {config.encoder.hidden_size} for"
+                    " `config.encoder.hidden_size`."
+                )
+
+        # initialize with config
+        super().__init__(config)
+
+        if encoder is None:
+            from ..auto.modeling_auto import AutoModel
+
+            encoder = AutoModel.from_config(config.encoder)
+
+        if decoder is None:
+            from ..auto.modeling_auto import AutoModelForCausalLM
+
+            decoder = AutoModelForCausalLM.from_config(config.decoder)
+
+        self.encoder = encoder
+        self.decoder = decoder
+
+        if self.encoder.config.to_dict() != self.config.encoder.to_dict():
+            logger.warning(
+                f"Config of the encoder: {self.encoder.__class__} is overwritten by shared encoder config:"
+                f" {self.config.encoder}"
+            )
+        if self.decoder.config.to_dict() != self.config.decoder.to_dict():
+            logger.warning(
+                f"Config of the decoder: {self.decoder.__class__} is overwritten by shared decoder config:"
+                f" {self.config.decoder}"
+            )
+
+        # make sure that the individual model's config refers to the shared config
+        # so that the updates to the config will be synced
+        # update `_attn_implementation` because the attn is set in a deepcopied config within PreTrainedModel
+        self.config.encoder._attn_implementation = self.encoder.config._attn_implementation
+        self.config.decoder._attn_implementation = self.decoder.config._attn_implementation
+        self.encoder.config = self.config.encoder
+        self.decoder.config = self.config.decoder
+
+        # encoder outputs might need to be projected to different dimension for decoder
+        if (
+            self.encoder.config.hidden_size != self.decoder.config.hidden_size
+            and getattr(self.decoder.config, "cross_attention_hidden_size", None) is None
+        ):
+            self.enc_to_dec_proj = nn.Linear(self.encoder.config.hidden_size, self.decoder.config.hidden_size)
+
+        if self.encoder.get_output_embeddings() is not None:
+            raise ValueError(
+                f"The encoder {self.encoder} should not have a LM Head. Please use a model without LM Head"
+            )
+
+        decoder_signature = set(inspect.signature(self.decoder.forward).parameters.keys())
+        if "encoder_hidden_states" not in decoder_signature:
+            raise ValueError(
+                "The selected decoder is not prepared for the encoder hidden states to be passed. Please see the "
+                "following discussion on GitHub: https://github.com/huggingface/transformers/issues/23350"
+            )
+
+        self.post_init()
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        if module in self.encoder.modules():
+            self.encoder._init_weights(module)
+        elif module in self.decoder.modules():
+            self.decoder._init_weights(module)
+
+    def get_input_embeddings(self):
+        return self.encoder.get_input_embeddings()
+
+    def get_output_embeddings(self):
+        return self.decoder.get_output_embeddings()
+
+    def set_output_embeddings(self, new_embeddings):
+        return self.decoder.set_output_embeddings(new_embeddings)
+
+    @classmethod
+    def from_encoder_decoder_pretrained(
+        cls,
+        encoder_pretrained_model_name_or_path: str | None = None,
+        decoder_pretrained_model_name_or_path: str | None = None,
+        *model_args,
+        **kwargs,
+    ) -> PreTrainedModel:
+        r"""
+        Instantiate an encoder and a decoder from one or two base classes of the library from pretrained model
+        checkpoints.
+
+
+        The model is set in evaluation mode by default using `model.eval()` (Dropout modules are deactivated). To train
+        the model, you need to first set it back in training mode with `model.train()`.
+
+        Params:
+            encoder_pretrained_model_name_or_path (`str`, *optional*):
+                Information necessary to initiate the encoder. Can be either:
+
+                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+                    - A path to a *directory* containing model weights saved using
+                      [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+
+            decoder_pretrained_model_name_or_path (`str`, *optional*, defaults to `None`):
+                Information necessary to initiate the decoder. Can be either:
+
+                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
+                    - A path to a *directory* containing model weights saved using
+                      [`~PreTrainedModel.save_pretrained`], e.g., `./my_model_directory/`.
+
+            model_args (remaining positional arguments, *optional*):
+                All remaining positional arguments will be passed to the underlying model's `__init__` method.
+
+            kwargs (remaining dictionary of keyword arguments, *optional*):
+                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
+                `output_attentions=True`).
+
+                - To update the encoder configuration, use the prefix *encoder_* for each configuration parameter.
+                - To update the decoder configuration, use the prefix *decoder_* for each configuration parameter.
+                - To update the parent model configuration, do not use a prefix for each configuration parameter.
+
+                Behaves differently depending on whether a `config` is provided or automatically loaded.
+
+        Example:
+
+        ```python
+        >>> from transformers import EncoderDecoderModel
+
+        >>> # initialize a bert2bert from two pretrained BERT models. Note that the cross-attention layers will be randomly initialized
+        >>> model = EncoderDecoderModel.from_encoder_decoder_pretrained("google-bert/bert-base-uncased", "google-bert/bert-base-uncased")
+        >>> # saving model after fine-tuning
+        >>> model.save_pretrained("./bert2bert")
+        >>> # load fine-tuned model
+        >>> model = EncoderDecoderModel.from_pretrained("./bert2bert")
+        ```"""
+
+        kwargs_encoder = {
+            argument[len("encoder_") :]: value for argument, value in kwargs.items() if argument.startswith("encoder_")
+        }
+
+        kwargs_decoder = {
+            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+        }
+
+        # remove encoder, decoder kwargs from kwargs
+        for key in kwargs_encoder:
+            del kwargs["encoder_" + key]
+        for key in kwargs_decoder:
+            del kwargs["decoder_" + key]
+
+        # Load and initialize the encoder and decoder
+        # The distinction between encoder and decoder at the model level is made
+        # by the value of the flag `is_decoder` that we need to set correctly.
+        encoder = kwargs_encoder.pop("model", None)
+        if encoder is None:
+            if encoder_pretrained_model_name_or_path is None:
+                raise ValueError(
+                    "If `encoder_model` is not defined as an argument, a `encoder_pretrained_model_name_or_path` has "
+                    "to be defined."
+                )
+
+            if "config" not in kwargs_encoder:
+                encoder_config, kwargs_encoder = AutoConfig.from_pretrained(
+                    encoder_pretrained_model_name_or_path, **kwargs_encoder, return_unused_kwargs=True
+                )
+
+                if getattr(encoder_config, "is_decoder", False) or getattr(
+                    encoder_config, "add_cross_attention", False
+                ):
+                    logger.info(
+                        f"Initializing {encoder_pretrained_model_name_or_path} as a encoder model "
+                        "from a decoder model. Cross-attention and causal mask are disabled."
+                    )
+                    encoder_config.is_decoder = False
+                    encoder_config.add_cross_attention = False
+
+                kwargs_encoder["config"] = encoder_config
+
+            encoder = AutoModel.from_pretrained(encoder_pretrained_model_name_or_path, *model_args, **kwargs_encoder)
+
+        decoder = kwargs_decoder.pop("model", None)
+        if decoder is None:
+            if decoder_pretrained_model_name_or_path is None:
+                raise ValueError(
+                    "If `decoder_model` is not defined as an argument, a `decoder_pretrained_model_name_or_path` has "
+                    "to be defined."
+                )
+
+            if "config" not in kwargs_decoder:
+                decoder_config, kwargs_decoder = AutoConfig.from_pretrained(
+                    decoder_pretrained_model_name_or_path, **kwargs_decoder, return_unused_kwargs=True
+                )
+            else:
+                decoder_config = kwargs_decoder["config"]
+
+            if (
+                getattr(decoder_config, "is_decoder", None) is False
+                or getattr(decoder_config, "add_cross_attention", None) is False
+            ):
+                logger.info(
+                    f"Initializing {decoder_pretrained_model_name_or_path} as a decoder model. Cross attention"
+                    f" layers are added to {decoder_pretrained_model_name_or_path} and randomly initialized if"
+                    f" {decoder_pretrained_model_name_or_path}'s architecture allows for cross attention layers."
+                )
+                decoder_config.is_decoder = True
+                decoder_config.add_cross_attention = True
+
+            kwargs_decoder["config"] = decoder_config
+            decoder = AutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)
+
+        # instantiate config with corresponding kwargs
+        config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
+        return cls(encoder=encoder, decoder=decoder, config=config)
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.FloatTensor | None = None,
+        decoder_input_ids: torch.LongTensor | None = None,
+        decoder_attention_mask: torch.BoolTensor | None = None,
+        encoder_outputs: tuple[torch.FloatTensor] | None = None,
+        past_key_values: Cache | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        decoder_inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        use_cache: bool | None = None,
+        **kwargs,
+    ) -> tuple | Seq2SeqLMOutput:
+        r"""
+        decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Indices of decoder input sequence tokens in the vocabulary.
+
+            Indices can be obtained using [`PreTrainedTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+
+            If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see
+            `past_key_values`).
+
+            For training, `decoder_input_ids` are automatically created by the model by shifting the `labels` to the
+            right, replacing -100 by the `pad_token_id` and prepending them with the `decoder_start_token_id`.
+        decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):
+            Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also
+            be used by default.
+        decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded
+            representation. This is useful if you want more control over how to convert `decoder_input_ids` indices
+            into associated vectors than the model's internal embedding lookup matrix.
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the masked language modeling loss for the decoder. Indices should be in `[-100, 0,
+            ..., config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored
+            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
+
+        Examples:
+
+        ```python
+        >>> from transformers import EncoderDecoderModel, BertTokenizer
+        >>> import torch
+
+        >>> tokenizer = BertTokenizer.from_pretrained("google-bert/bert-base-uncased")
+        >>> model = EncoderDecoderModel.from_encoder_decoder_pretrained(
+        ...     "google-bert/bert-base-uncased", "google-bert/bert-base-uncased"
+        ... )  # initialize Bert2Bert from pre-trained checkpoints
+
+        >>> # training
+        >>> model.config.decoder_start_token_id = tokenizer.cls_token_id
+        >>> model.config.pad_token_id = tokenizer.pad_token_id
+        >>> model.config.vocab_size = model.config.decoder.vocab_size
+
+        >>> input_ids = tokenizer("This is a really long text", return_tensors="pt").input_ids
+        >>> labels = tokenizer("This is the corresponding summary", return_tensors="pt").input_ids
+        >>> outputs = model(input_ids=input_ids, labels=labels)
+        >>> loss, logits = outputs.loss, outputs.logits
+
+        >>> # save and load from pretrained
+        >>> model.save_pretrained("bert2bert")
+        >>> model = EncoderDecoderModel.from_pretrained("bert2bert")
+
+        >>> # generation
+        >>> generated = model.generate(input_ids)
+        ```"""
+        # `record outputs` can rely on the absence of the kwarg to retrieve whether the config should be used or not
+        # Hence, we use this workaround to allow for defaults to work as expected
+        kwargs_shared = {key: kwargs[key] for key in ["output_attentions", "output_hidden_states"] if key in kwargs}
+
+        kwargs_encoder = {argument: value for argument, value in kwargs.items() if not argument.startswith("decoder_")}
+        kwargs_encoder = kwargs_encoder | kwargs_shared
+
+        kwargs_decoder = {
+            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
+        }
+        if "num_items_in_batch" in kwargs_encoder:
+            kwargs_decoder["num_items_in_batch"] = kwargs_encoder.pop("num_items_in_batch", None)
+        kwargs_decoder = kwargs_decoder | kwargs_shared
+
+        if encoder_outputs is None:
+            encoder_outputs = self.encoder(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                return_dict=True,
+                **kwargs_encoder,
+            )
+        elif isinstance(encoder_outputs, tuple):
+            encoder_outputs = BaseModelOutput(*encoder_outputs)
+
+        encoder_hidden_states = encoder_outputs[0]
+
+        # optionally project encoder_hidden_states
+        if (
+            self.encoder.config.hidden_size != self.decoder.config.hidden_size
+            and getattr(self.decoder.config, "cross_attention_hidden_size", None) is None
+        ):
+            encoder_hidden_states = self.enc_to_dec_proj(encoder_hidden_states)
+
+        if (labels is not None) and (decoder_input_ids is None and decoder_inputs_embeds is None):
+            decoder_input_ids = shift_tokens_right(
+                labels, self.config.pad_token_id, self.config.decoder_start_token_id
+            )
+            if decoder_attention_mask is None:
+                decoder_attention_mask = (decoder_input_ids != self.config.pad_token_id).to(decoder_input_ids.dtype)
+
+        # Decode
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=attention_mask,
+            inputs_embeds=decoder_inputs_embeds,
+            use_cache=use_cache,
+            past_key_values=past_key_values,
+            return_dict=True,
+            **kwargs_decoder,
+        )
+
+        # Compute loss independent from decoder (as some shift the logits inside them)
+        loss = None
+        if labels is not None:
+            warnings.warn(DEPRECATION_WARNING, FutureWarning)
+            logits = decoder_outputs.logits
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(logits.reshape(-1, self.decoder.config.vocab_size), labels.view(-1))
+
+        return Seq2SeqLMOutput(
+            loss=loss,
+            logits=decoder_outputs.logits,
+            past_key_values=decoder_outputs.past_key_values,
+            decoder_hidden_states=decoder_outputs.hidden_states,
+            decoder_attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+            encoder_last_hidden_state=encoder_outputs.last_hidden_state,
+            encoder_hidden_states=encoder_outputs.hidden_states,
+            encoder_attentions=encoder_outputs.attentions,
+        )
+
+    def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):
+        return shift_tokens_right(labels, self.config.pad_token_id, self.config.decoder_start_token_id)
+
+    def resize_token_embeddings(self, *args, **kwargs):
+        raise NotImplementedError(
+            "Resizing the embedding layers via the EncoderDecoderModel directly is not supported. Please use the"
+            " respective methods of the wrapped objects (model.encoder.resize_token_embeddings(...) or"
+            " model.decoder.resize_token_embeddings(...))"
+        )
+
+
+__all__ = ["EncoderDecoderModel"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/lw_detr/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_lw_detr import *
+    from .modeling_lw_detr import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/lw_detr/configuration_lw_detr.py

@@ -0,0 +1,226 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/lw_detr/modular_lw_detr.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_lw_detr.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+
+from huggingface_hub.dataclasses import strict
+
+from ...backbone_utils import BackboneConfigMixin, consolidate_backbone_kwargs_to_config
+from ...configuration_utils import PreTrainedConfig
+from ...utils import auto_docstring, logging
+from ..auto import AutoConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="AnnaZhang/lwdetr_small_60e_coco")
+@strict
+class LwDetrViTConfig(BackboneConfigMixin, PreTrainedConfig):
+    r"""
+    pretrain_image_size (`int`, *optional*, defaults to 224):
+        The size (resolution) of each image during pretraining.
+    window_block_indices (`list[int]`, *optional*, defaults to `[]`):
+        List of indices of blocks that should have window attention instead of regular global self-attention.
+    use_absolute_position_embeddings (`bool`, *optional*, defaults to `True`):
+        Whether to add absolute position embeddings to the patch embeddings.
+    cae_init_values (`float`, *optional*, defaults to 0.1):
+        Initialization value for CAE parameters when `use_cae` is enabled.
+    num_windows (`int`, *optional*, defaults to 16):
+        Number of windows for window-based attention. Must be a perfect square and the image size must be
+        divisible by the square root of this value. This enables efficient window-major feature map organization.
+
+    Example:
+
+    ```python
+    >>> from transformers import LwDetrViTConfig, LwDetrViTModel
+
+    >>> # Initializing a LW-DETR ViT configuration
+    >>> configuration = LwDetrViTConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = LwDetrViTModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "lw_detr_vit"
+
+    hidden_size: int = 768
+    num_hidden_layers: int = 12
+    num_attention_heads: int = 12
+    mlp_ratio: int = 4
+    hidden_act: str = "gelu"
+    dropout_prob: float | int = 0.0
+    initializer_range: float = 0.02
+    layer_norm_eps: float = 1e-6
+
+    image_size: int | list[int] | tuple[int, int] = 256
+    pretrain_image_size: int | list[int] | tuple[int, int] = 224
+    patch_size: int | list[int] | tuple[int, int] = 16
+    num_channels: int = 3
+    qkv_bias: bool = True
+    window_block_indices: list[int] | tuple[int, ...] = ()
+    use_absolute_position_embeddings: bool = True
+    _out_features: list[str] | None = None
+    _out_indices: list[int] | None = None
+    cae_init_values: float = 0.1
+    num_windows: int = 16
+
+    def __post_init__(self, **kwargs):
+        self.num_windows_side = int(math.sqrt(self.num_windows))
+        self.stage_names = ["stem"] + [f"stage{idx}" for idx in range(1, self.num_hidden_layers + 1)]
+        self.set_output_features_output_indices(
+            out_indices=kwargs.pop("out_indices", None), out_features=kwargs.pop("out_features", None)
+        )
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of `@strict`-powered validation. Validates the architecture of the config."""
+        if self.num_windows % math.sqrt(self.num_windows) != 0:
+            raise ValueError(
+                f"`num_windows` has to be a perfect square, where num_windows % math.sqrt(num_windows) != 0, but got {self.num_windows}."
+            )
+        if self.image_size / self.num_windows % math.sqrt(self.num_windows) != 0:
+            raise ValueError(
+                f"`image_size` has to be divisible by `num_windows`, where image_size / num_windows % math.sqrt(num_windows) != 0,but got {self.image_size} and {self.num_windows}."
+            )
+
+
+@auto_docstring(checkpoint="AnnaZhang/lwdetr_small_60e_coco")
+@strict
+class LwDetrConfig(PreTrainedConfig):
+    r"""
+    projector_scale_factors (`list[float]`, *optional*, defaults to `[]`):
+        Scale factors for the feature pyramid network. Each scale factor determines the resolution of features
+        at different levels. Supported values are 0.5, 1.0, and 2.0.
+    hidden_expansion (`float`, *optional*, defaults to 0.5):
+        Expansion factor for hidden dimensions in the projector layers.
+    c2f_num_blocks (`int`, *optional*, defaults to 3):
+        Number of blocks in the C2F layer.
+    activation_function (`str`, *optional*, defaults to `"silu"`):
+        The non-linear activation function in the projector. Supported values are `"silu"`, `"relu"`, `"gelu"`.
+    batch_norm_eps (`float`, *optional*, defaults to 1e-05):
+        The epsilon value for batch normalization layers.
+    decoder_ffn_dim (`int`, *optional*, defaults to 2048):
+        Dimension of the "intermediate" (often named feed-forward) layer in decoder.
+    decoder_n_points (`int`, *optional*, defaults to 4):
+        The number of sampled keys in each feature level for each attention head in the decoder.
+    decoder_self_attention_heads (`int`, *optional*, defaults to 8):
+        Number of attention heads for each attention layer in the decoder self-attention.
+    decoder_cross_attention_heads (`int`, *optional*, defaults to 16):
+        Number of attention heads for each attention layer in the decoder cross-attention.
+    decoder_activation_function (`str`, *optional*, defaults to `"relu"`):
+        The non-linear activation function in the decoder. Supported values are `"relu"`, `"silu"`, `"gelu"`.
+    num_queries (`int`, *optional*, defaults to 300):
+        Number of object queries, i.e. detection slots. This is the maximal number of objects
+        [`LwDetrModel`] can detect in a single image.
+    group_detr (`int`, *optional*, defaults to 13):
+        Number of groups for Group DETR attention mechanism, which helps reduce computational complexity.
+    disable_custom_kernels (`bool`, *optional*, defaults to `True`):
+        Disable the use of custom CUDA and CPU kernels. This option is necessary for the ONNX export, as custom
+        kernels are not supported by PyTorch ONNX export.
+    class_loss_coefficient (`float`, *optional*, defaults to 1):
+        Relative weight of the classification loss in the Hungarian matching cost.
+
+    Examples:
+
+    ```python
+    >>> from transformers import LwDetrConfig, LwDetrModel
+
+    >>> # Initializing a LW-DETR AnnaZhang/lwdetr_small_60e_coco style configuration
+    >>> configuration = LwDetrConfig()
+
+    >>> # Initializing a model (with random weights) from the AnnaZhang/lwdetr_small_60e_coco style configuration
+    >>> model = LwDetrModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "lw_detr"
+    sub_configs = {"backbone_config": AutoConfig}
+
+    backbone_config: dict | PreTrainedConfig | None = None
+    projector_scale_factors: list[float] | tuple[float, ...] = ()
+    hidden_expansion: float = 0.5
+    c2f_num_blocks: int = 3
+    activation_function: str = "silu"
+    batch_norm_eps: float = 1e-5
+    dropout: float | int = 0.0
+    decoder_ffn_dim: int = 2048
+    decoder_n_points: int = 4
+    decoder_layers: int = 3
+    decoder_self_attention_heads: int = 8
+    decoder_cross_attention_heads: int = 16
+    decoder_activation_function: str = "relu"
+    num_queries: int = 300
+    attention_bias: bool = True
+    attention_dropout: float | int = 0.0
+    activation_dropout: float | int = 0.0
+    group_detr: int = 13
+    init_std: float = 0.02
+    disable_custom_kernels: bool = True
+    class_cost: int | float = 2
+    bbox_cost: int | float = 5
+    giou_cost: int | float = 2
+    class_loss_coefficient: int | float = 1
+    dice_loss_coefficient: int | float = 1
+    bbox_loss_coefficient: int | float = 5
+    giou_loss_coefficient: int | float = 2
+    eos_coefficient: float = 0.1
+    focal_alpha: float = 0.25
+    auxiliary_loss: bool = True
+    d_model: int = 256
+
+    def __post_init__(self, **kwargs):
+        if "mask_loss_coefficient" in kwargs:
+            logger.warning_once(
+                "The parameter `mask_loss_coefficient` was renamed to `class_loss_coefficient` in LW-DETR. "
+                "Please use `class_loss_coefficient` instead. `mask_loss_coefficient` will be removed in a future version."
+            )
+            self.class_loss_coefficient = kwargs.pop("mask_loss_coefficient")
+
+        self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config(
+            backbone_config=self.backbone_config,
+            default_config_type="lw_detr_vit",
+            default_config_kwargs={
+                "image_size": 1024,
+                "hidden_size": 192,
+                "num_hidden_layers": 10,
+                "window_block_indices": [0, 1, 3, 6, 7, 9],
+                "out_indices": [2, 4, 5, 9],
+            },
+            **kwargs,
+        )
+
+        self.projector_in_channels = [self.d_model] * len(self.projector_scale_factors)
+        self.projector_out_channels = self.d_model
+        self.num_feature_levels = len(self.projector_scale_factors)
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of `@strict`-powered validation. Validates the architecture of the config."""
+        for scale in self.projector_scale_factors:
+            if scale not in [0.5, 1.0, 2.0]:
+                raise ValueError(f"Unsupported scale factor: {scale}")
+
+
+__all__ = ["LwDetrConfig", "LwDetrViTConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/lw_detr/modeling_lw_detr.py

@@ -0,0 +1,1673 @@
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+#           This file was automatically generated from src/transformers/models/lw_detr/modular_lw_detr.py.
+#               Do NOT edit this file manually as any edits will be overwritten by the generation of
+#             the file from the modular. If any change should be done, please apply the change to the
+#                          modular_lw_detr.py file directly. One of our CI enforces this.
+#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import collections.abc
+import math
+import warnings
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+from ... import initialization as init
+from ...activations import ACT2CLS, ACT2FN
+from ...backbone_utils import BackboneMixin, filter_output_hidden_states
+from ...integrations import use_kernel_forward_from_hub
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BackboneOutput, BaseModelOutput, BaseModelOutputWithCrossAttentions
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, torch_compilable_check
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from .configuration_lw_detr import LwDetrConfig, LwDetrViTConfig
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    key_states = repeat_kv(key, module.num_key_value_groups)
+    value_states = repeat_kv(value, module.num_key_value_groups)
+
+    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+    attn_output = torch.matmul(attn_weights, value_states)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """
+    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
+    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class LwDetrViTAttention(nn.Module):
+    """LwDetr ViT attention with k_proj bias=False and dropout from config.dropout_prob."""
+
+    def __init__(self, config: LwDetrViTConfig):
+        super().__init__()
+        self.config = config
+        self.num_attention_heads = config.num_attention_heads
+        self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
+        self.attention_dropout = config.dropout_prob
+        self.scaling = self.head_dim**-0.5
+        self.is_causal = False
+
+        self.q_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.qkv_bias)
+        self.k_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.qkv_bias)
+        self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=True)
+        self.num_key_value_groups = 1
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        return attn_output, attn_weights
+
+
+class LwDetrViTMlp(nn.Module):
+    def __init__(self, config, in_features: int, hidden_features: int) -> None:
+        super().__init__()
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = ACT2FN[config.hidden_act]
+        self.fc2 = nn.Linear(hidden_features, in_features)
+        self.drop = nn.Dropout(config.dropout_prob)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+
+        return x
+
+
+class LwDetrViTLayer(GradientCheckpointingLayer):
+    def __init__(
+        self,
+        config: LwDetrViTConfig,
+        layer_idx,
+    ) -> None:
+        super().__init__()
+
+        dim = config.hidden_size
+        self.attention = LwDetrViTAttention(config)
+        self.intermediate = LwDetrViTMlp(config=config, in_features=dim, hidden_features=int(dim * config.mlp_ratio))
+        self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        self.gamma_1 = nn.Parameter(torch.Tensor(dim), requires_grad=True)
+        self.gamma_2 = nn.Parameter(torch.Tensor(dim), requires_grad=True)
+
+        self.window = layer_idx in config.window_block_indices
+        self.num_windows = config.num_windows
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        batch_size, seq_len, channels = hidden_states.shape
+        hidden_states_norm = self.layernorm_before(hidden_states)
+
+        if not self.window:
+            hidden_states_norm = hidden_states_norm.reshape(
+                batch_size // self.num_windows, self.num_windows * seq_len, channels
+            )
+
+        attention_output, _ = self.attention(hidden_states_norm, **kwargs)
+        attention_output = attention_output * self.gamma_1
+
+        if not self.window:
+            attention_output = attention_output.reshape(batch_size, seq_len, channels)
+
+        hidden_states = hidden_states + attention_output
+
+        layer_output = self.layernorm_after(hidden_states)
+        layer_output = self.intermediate(layer_output)
+        layer_output = layer_output * self.gamma_2
+
+        hidden_states = hidden_states + layer_output
+
+        return hidden_states
+
+
+class LwDetrViTEmbeddings(nn.Module):
+    """
+    This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
+    `hidden_states` (patch embeddings) to be consumed by a Transformer.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        image_size, patch_size = config.pretrain_image_size, config.patch_size
+        num_channels, hidden_size = config.num_channels, config.hidden_size
+
+        image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
+        patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
+        num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_channels = num_channels
+        self.num_patches = num_patches
+
+        if config.use_absolute_position_embeddings:
+            # Initialize absolute positional embedding with pretrain image size.
+            num_positions = num_patches + 1
+            self.position_embeddings = nn.Parameter(torch.zeros(1, num_positions, config.hidden_size))
+        else:
+            self.position_embeddings = None
+
+        self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
+
+    def get_absolute_positions(self, abs_pos_embeddings, has_cls_token, height, width):
+        """
+        Calculate absolute positional embeddings. If needed, resize embeddings and remove cls_token dimension for the
+        original embeddings.
+
+        Args:
+            abs_pos_embeddings (`torch.Tensor`):
+                Absolute positional embeddings with (1, num_position, num_channels).
+            has_cls_token (`bool`):
+                If true, has 1 embedding in abs_pos_embeddings for cls token.
+            height (`int`):
+                Height of input image tokens.
+            width (`int`):
+                Width of input image tokens.
+
+        Returns:
+            Absolute positional embeddings after processing with shape (1, height, width, num_channels)
+        """
+        if has_cls_token:
+            abs_pos_embeddings = abs_pos_embeddings[:, 1:]
+        num_position = abs_pos_embeddings.shape[1]
+        size = int(math.sqrt(num_position))  # This is a constant and can be recorded as such in the ONNX export.
+        if size * size != num_position:
+            raise ValueError("Absolute position embeddings must be a square number.")
+
+        if torch.jit.is_tracing() or (size != height or size != width):
+            # nn.functional.interpolate is a noop in case size == height and size == width - we need to always capture this path with jit.trace.
+            new_abs_pos_embeddings = nn.functional.interpolate(
+                abs_pos_embeddings.reshape(1, size, size, -1).permute(0, 3, 1, 2),
+                size=(height, width),
+                mode="bicubic",
+                align_corners=False,
+            )
+
+            return new_abs_pos_embeddings.permute(0, 2, 3, 1)
+        else:
+            return abs_pos_embeddings.reshape(1, height, width, -1)
+
+    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:
+        num_channels = pixel_values.shape[1]
+        if num_channels != self.num_channels:
+            raise ValueError(
+                "Make sure that the channel dimension of the pixel values match with the one set in the configuration."
+                f" Expected {self.num_channels} but got {num_channels}."
+            )
+        embeddings = self.projection(pixel_values)
+
+        if self.position_embeddings is not None:
+            # (batch_size, num_channels, height, width) -> (batch_size, height, width, num_channels)
+            embeddings = embeddings.permute(0, 2, 3, 1)
+            # add position embeddings
+            embeddings = embeddings + self.get_absolute_positions(
+                self.position_embeddings, True, embeddings.shape[1], embeddings.shape[2]
+            )
+            # (batch_size, height, width, num_channels) -> (batch_size, num_channels, height, width)
+            embeddings = embeddings.permute(0, 3, 1, 2)
+
+        return embeddings
+
+
+@auto_docstring
+class LwDetrViTPreTrainedModel(PreTrainedModel):
+    config: LwDetrViTConfig
+    base_model_prefix = "lw_detr_vit"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LwDetrViTEmbeddings", "LwDetrViTLayer"]
+    _supports_sdpa = True
+    _supports_flash_attn = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": LwDetrViTLayer,
+        "attentions": LwDetrViTAttention,
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module) -> None:
+        """Initialize the weights"""
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            init.trunc_normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            init.zeros_(module.bias)
+            init.ones_(module.weight)
+        elif isinstance(module, LwDetrViTEmbeddings):
+            init.trunc_normal_(module.position_embeddings, mean=0.0, std=self.config.initializer_range)
+        if isinstance(module, LwDetrViTLayer):
+            init.constant_(module.gamma_1, self.config.cae_init_values)
+            init.constant_(module.gamma_2, self.config.cae_init_values)
+
+
+class LwDetrViTEncoder(LwDetrViTPreTrainedModel):
+    def __init__(self, config: LwDetrViTConfig):
+        super().__init__(config)
+        self.layer = nn.ModuleList([LwDetrViTLayer(config, idx) for idx in range(config.num_hidden_layers)])
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutput:
+        for layer_module in self.layer:
+            hidden_states = layer_module(hidden_states, **kwargs)
+
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+@auto_docstring()
+class LwDetrViTBackbone(BackboneMixin, LwDetrViTPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.embeddings = LwDetrViTEmbeddings(config)
+        self.encoder = LwDetrViTEncoder(config)
+        self.num_features = [config.hidden_size for _ in range(config.num_hidden_layers + 1)]
+
+        # initialize weights and apply final processing
+        self.post_init()
+
+    def get_input_embeddings(self) -> LwDetrViTEmbeddings:
+        return self.embeddings.projection
+
+    @can_return_tuple
+    @filter_output_hidden_states
+    @auto_docstring
+    def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BackboneOutput:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import LwDetrViTConfig, LwDetrViTBackbone
+        >>> import torch
+
+        >>> config = LwDetrViTConfig()
+        >>> model = LwDetrViTBackbone(config)
+
+        >>> pixel_values = torch.randn(1, 3, 224, 224)
+
+        >>> with torch.no_grad():
+        ...     outputs = model(pixel_values)
+
+        >>> feature_maps = outputs.feature_maps
+        >>> list(feature_maps[-1].shape)
+        [1, 768, 14, 14]
+        ```"""
+        embedding_output = self.embeddings(pixel_values)
+
+        batch_size, channels, height, width = embedding_output.shape
+        # (batch_size, channels, height, width) -> (batch_size, height, width, channels)
+        hidden_states = embedding_output.permute(0, 2, 3, 1)
+
+        window_height = height // self.config.num_windows_side
+        window_width = width // self.config.num_windows_side
+        # (batch_size, height, width, channels) -> (batch_size*num_windows_side**2, window_height*window_width, channels)
+        hidden_states = (
+            hidden_states.reshape(
+                batch_size,
+                self.config.num_windows_side,
+                window_height,
+                self.config.num_windows_side,
+                window_width,
+                channels,
+            )
+            .permute(0, 1, 3, 2, 4, 5)
+            .reshape(batch_size * self.config.num_windows_side**2, window_height * window_width, channels)
+        )
+
+        kwargs["output_hidden_states"] = True  # required to extract layers for the stages
+        output = self.encoder(hidden_states, **kwargs)
+
+        feature_maps = ()
+        hidden_states = output.hidden_states
+        for stage, hidden_state in zip(self.stage_names, hidden_states):
+            if stage in self.out_features:
+                hidden_state = (
+                    hidden_state.reshape(
+                        batch_size,
+                        self.config.num_windows_side,
+                        self.config.num_windows_side,
+                        window_height,
+                        window_width,
+                        channels,
+                    )
+                    .permute(0, 5, 1, 3, 2, 4)
+                    .reshape(batch_size, channels, height, width)
+                )
+                feature_maps += (hidden_state,)
+
+        return BackboneOutput(
+            feature_maps=feature_maps, hidden_states=output.hidden_states, attentions=output.attentions
+        )
+
+
+class LwDetrConvNormLayer(nn.Module):
+    def __init__(
+        self,
+        config: LwDetrConfig,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+        activation: str | None = None,
+    ):
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding=kernel_size // 2,
+            bias=False,
+        )
+        self.norm = nn.BatchNorm2d(out_channels, config.batch_norm_eps)
+        self.activation = nn.Identity() if activation is None else ACT2CLS[activation]()
+
+    def forward(self, hidden_state):
+        hidden_state = self.conv(hidden_state)
+        hidden_state = self.norm(hidden_state)
+        hidden_state = self.activation(hidden_state)
+        return hidden_state
+
+
+class LwDetrRepVggBlock(nn.Module):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__()
+        hidden_channels = int(config.d_model * config.hidden_expansion)
+        self.conv1 = LwDetrConvNormLayer(
+            config, hidden_channels, hidden_channels, 3, 1, activation=config.activation_function
+        )
+        self.conv2 = LwDetrConvNormLayer(
+            config, hidden_channels, hidden_channels, 3, 1, activation=config.activation_function
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.conv1(x)
+        y = self.conv2(y)
+        return y
+
+
+class LwDetrC2FLayer(nn.Module):
+    # Inspired by RTDetrCSPRepLayer
+    def __init__(self, config: LwDetrConfig, in_channels: int):
+        super().__init__()
+        num_blocks = config.c2f_num_blocks
+        activation = config.activation_function
+        out_channels = config.d_model
+
+        self.hidden_channels = int(out_channels * config.hidden_expansion)
+
+        conv1_out_channels = 2 * self.hidden_channels
+        self.conv1 = LwDetrConvNormLayer(config, in_channels, conv1_out_channels, 1, 1, activation=activation)
+
+        conv2_in_channels = (2 + num_blocks) * self.hidden_channels
+        self.conv2 = LwDetrConvNormLayer(config, conv2_in_channels, out_channels, 1, 1, activation=activation)
+
+        self.bottlenecks = nn.ModuleList(LwDetrRepVggBlock(config) for _ in range(num_blocks))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv1(hidden_states)
+        all_hidden_states = list(hidden_states.split(self.hidden_channels, 1))
+        hidden_states = all_hidden_states[-1]
+        hidden_states = hidden_states.contiguous()
+
+        for bottleneck in self.bottlenecks:
+            hidden_states = bottleneck(hidden_states)
+            all_hidden_states.append(hidden_states)
+
+        hidden_states = torch.cat(all_hidden_states, 1)
+        hidden_states = self.conv2(hidden_states)
+        return hidden_states
+
+
+class LwDetrLayerNorm(nn.LayerNorm):
+    r"""LayerNorm that supports two data formats: channels_last (default) or channels_first.
+    The ordering of the dimensions in the inputs. channels_last corresponds to inputs with shape (batch_size, height,
+    width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width).
+    """
+
+    def __init__(self, normalized_shape, *, eps=1e-6, data_format="channels_last", **kwargs):
+        super().__init__(normalized_shape, eps=eps, **kwargs)
+        if data_format not in ["channels_last", "channels_first"]:
+            raise NotImplementedError(f"Unsupported data format: {data_format}")
+        self.data_format = data_format
+
+    def forward(self, features: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            features: Tensor of shape (batch_size, channels, height, width) OR (batch_size, height, width, channels)
+        """
+        if self.data_format == "channels_first":
+            features = features.permute(0, 2, 3, 1)
+            features = super().forward(features)
+            features = features.permute(0, 3, 1, 2)
+        else:
+            features = super().forward(features)
+        return features
+
+
+class LwDetrSamplingLayer(nn.Module):
+    def __init__(self, config: LwDetrConfig, channel_size: int, scale: float):
+        super().__init__()
+
+        self.scale = scale
+        self.channel_size = channel_size
+
+        layers = []
+        if scale == 2.0:
+            if channel_size > 512:
+                layers.append(LwDetrConvNormLayer(config, channel_size, channel_size // 2, 1, 1, activation="relu"))
+                layers.append(nn.ConvTranspose2d(channel_size // 2, channel_size // 4, kernel_size=2, stride=2))
+            else:
+                layers.append(nn.ConvTranspose2d(channel_size, channel_size // 2, 2, 2))
+        elif scale == 0.5:
+            layers.append(LwDetrConvNormLayer(config, channel_size, channel_size, 3, 2, activation="relu"))
+        self.layers = nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+        return hidden_states
+
+
+class LwDetrScaleProjector(nn.Module):
+    def __init__(self, config: LwDetrConfig, scale: float):
+        super().__init__()
+
+        intermediate_dims = [config.backbone_config.hidden_size] * len(config.backbone_config.out_indices)
+        sampling_layers = []
+        for channel_size in intermediate_dims:
+            sampling_layers.append(LwDetrSamplingLayer(config, channel_size, scale))
+        self.sampling_layers = nn.ModuleList(sampling_layers)
+
+        intermediate_dim = intermediate_dims[-1]
+        if scale == 2.0:
+            if intermediate_dim > 512:
+                intermediate_dim = intermediate_dim // 4
+            else:
+                intermediate_dim = intermediate_dim // 2
+        projector_input_dim = intermediate_dim * len(intermediate_dims)
+
+        self.projector_layer = LwDetrC2FLayer(config, projector_input_dim)
+        self.layer_norm = LwDetrLayerNorm(config.d_model, data_format="channels_first")
+
+    def forward(self, hidden_states_tuple: tuple[torch.Tensor]) -> torch.Tensor:
+        sampled_hidden_states = []
+        for sampling_layer, hidden_states in zip(self.sampling_layers, hidden_states_tuple):
+            hidden_states = sampling_layer(hidden_states)
+            sampled_hidden_states.append(hidden_states)
+        hidden_states = torch.cat(sampled_hidden_states, dim=1)
+        hidden_states = self.projector_layer(hidden_states)
+        hidden_states = self.layer_norm(hidden_states)
+        return hidden_states
+
+
+class LwDetrMultiScaleProjector(nn.Module):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__()
+
+        self.config = config
+        scale_factors = config.projector_scale_factors
+
+        self.scale_layers = nn.ModuleList([LwDetrScaleProjector(config, scale) for scale in scale_factors])
+
+    def forward(self, hidden_states: tuple[torch.Tensor]) -> list[torch.Tensor]:
+        output_hidden_states = []
+        for scale_layer in self.scale_layers:
+            output_hidden_states.append(scale_layer(hidden_states))
+        return output_hidden_states
+
+
+class LwDetrConvEncoder(nn.Module):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__()
+        self.backbone = LwDetrViTBackbone(config.backbone_config)
+        self.projector = LwDetrMultiScaleProjector(config)
+
+    def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
+        # send pixel_values through the model to get list of feature maps
+        features = self.backbone(pixel_values).feature_maps
+        features = self.projector(features)
+        out = []
+        for feature_map in features:
+            # downsample pixel_mask to match shape of corresponding feature_map
+            mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
+            out.append((feature_map, mask))
+        return out
+
+
+class LwDetrAttention(nn.Module):
+    """LW-DETR self-attention with group-DETR training technique."""
+
+    def __init__(self, config: LwDetrConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.d_model // config.decoder_self_attention_heads)
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False
+        self.num_key_value_groups = 1
+
+        self.q_proj = nn.Linear(
+            config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.decoder_self_attention_heads * self.head_dim, config.d_model, bias=config.attention_bias
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        hidden_states_original = hidden_states
+        if position_embeddings is not None:
+            hidden_states = hidden_states if position_embeddings is None else hidden_states + position_embeddings
+
+        if self.training:
+            # at training, we use group detr technique to add more supervision by using multiple weight-sharing decoders at once for faster convergence
+            # at inference, we only use one decoder
+            hidden_states_original = torch.cat(
+                hidden_states_original.split(seq_len // self.config.group_detr, dim=1), dim=0
+            )
+            hidden_states = torch.cat(hidden_states.split(seq_len // self.config.group_detr, dim=1), dim=0)
+
+        attention_input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*attention_input_shape, -1, self.head_dim)
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states_original).view(hidden_shape).transpose(1, 2)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask=None,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(*attention_input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        if self.training:
+            attn_output = torch.cat(torch.split(attn_output, batch_size, dim=0), dim=1)
+
+        return attn_output, attn_weights
+
+
+@use_kernel_forward_from_hub("MultiScaleDeformableAttention")
+class MultiScaleDeformableAttention(nn.Module):
+    def forward(
+        self,
+        value: Tensor,
+        value_spatial_shapes: Tensor,
+        value_spatial_shapes_list: list[tuple],
+        level_start_index: Tensor,
+        sampling_locations: Tensor,
+        attention_weights: Tensor,
+        im2col_step: int,
+    ):
+        batch_size, _, num_heads, hidden_dim = value.shape
+        _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
+        value_list = value.split([height * width for height, width in value_spatial_shapes_list], dim=1)
+        sampling_grids = 2 * sampling_locations - 1
+        sampling_value_list = []
+        for level_id, (height, width) in enumerate(value_spatial_shapes_list):
+            # batch_size, height*width, num_heads, hidden_dim
+            # -> batch_size, height*width, num_heads*hidden_dim
+            # -> batch_size, num_heads*hidden_dim, height*width
+            # -> batch_size*num_heads, hidden_dim, height, width
+            value_l_ = (
+                value_list[level_id]
+                .flatten(2)
+                .transpose(1, 2)
+                .reshape(batch_size * num_heads, hidden_dim, height, width)
+            )
+            # batch_size, num_queries, num_heads, num_points, 2
+            # -> batch_size, num_heads, num_queries, num_points, 2
+            # -> batch_size*num_heads, num_queries, num_points, 2
+            sampling_grid_l_ = sampling_grids[:, :, :, level_id].transpose(1, 2).flatten(0, 1)
+            # batch_size*num_heads, hidden_dim, num_queries, num_points
+            sampling_value_l_ = nn.functional.grid_sample(
+                value_l_,
+                sampling_grid_l_,
+                mode="bilinear",
+                padding_mode="zeros",
+                align_corners=False,
+            )
+            sampling_value_list.append(sampling_value_l_)
+        # (batch_size, num_queries, num_heads, num_levels, num_points)
+        # -> (batch_size, num_heads, num_queries, num_levels, num_points)
+        # -> (batch_size, num_heads, 1, num_queries, num_levels*num_points)
+        attention_weights = attention_weights.transpose(1, 2).reshape(
+            batch_size * num_heads, 1, num_queries, num_levels * num_points
+        )
+        output = (
+            (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights)
+            .sum(-1)
+            .view(batch_size, num_heads * hidden_dim, num_queries)
+        )
+        return output.transpose(1, 2).contiguous()
+
+
+class LwDetrMultiscaleDeformableAttention(nn.Module):
+    """
+    Multiscale deformable attention as proposed in Deformable DETR.
+    """
+
+    def __init__(self, config: LwDetrConfig, num_heads: int, n_points: int):
+        super().__init__()
+
+        self.attn = MultiScaleDeformableAttention()
+
+        if config.d_model % num_heads != 0:
+            raise ValueError(
+                f"embed_dim (d_model) must be divisible by num_heads, but got {config.d_model} and {num_heads}"
+            )
+        dim_per_head = config.d_model // num_heads
+        # check if dim_per_head is power of 2
+        if not ((dim_per_head & (dim_per_head - 1) == 0) and dim_per_head != 0):
+            warnings.warn(
+                "You'd better set embed_dim (d_model) in LwDetrMultiscaleDeformableAttention to make the"
+                " dimension of each attention head a power of 2 which is more efficient in the authors' CUDA"
+                " implementation."
+            )
+
+        self.im2col_step = 64
+
+        self.d_model = config.d_model
+        self.n_levels = config.num_feature_levels
+        self.n_heads = num_heads
+        self.n_points = n_points
+
+        self.sampling_offsets = nn.Linear(config.d_model, num_heads * self.n_levels * n_points * 2)
+        self.attention_weights = nn.Linear(config.d_model, num_heads * self.n_levels * n_points)
+        self.value_proj = nn.Linear(config.d_model, config.d_model)
+        self.output_proj = nn.Linear(config.d_model, config.d_model)
+
+        self.disable_custom_kernels = config.disable_custom_kernels
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        position_embeddings: torch.Tensor | None = None,
+        reference_points=None,
+        spatial_shapes=None,
+        spatial_shapes_list=None,
+        level_start_index=None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # add position embeddings to the hidden states before projecting to queries and keys
+        if position_embeddings is not None:
+            hidden_states = hidden_states + position_embeddings
+
+        batch_size, num_queries, _ = hidden_states.shape
+        batch_size, sequence_length, _ = encoder_hidden_states.shape
+        total_elements = sum(height * width for height, width in spatial_shapes_list)
+        torch_compilable_check(
+            total_elements == sequence_length,
+            "Make sure to align the spatial shapes with the sequence length of the encoder hidden states",
+        )
+
+        value = self.value_proj(encoder_hidden_states)
+        if attention_mask is not None:
+            # we invert the attention_mask
+            value = value.masked_fill(~attention_mask[..., None], float(0))
+        value = value.view(batch_size, sequence_length, self.n_heads, self.d_model // self.n_heads)
+        sampling_offsets = self.sampling_offsets(hidden_states).view(
+            batch_size, num_queries, self.n_heads, self.n_levels, self.n_points, 2
+        )
+        attention_weights = self.attention_weights(hidden_states).view(
+            batch_size, num_queries, self.n_heads, self.n_levels * self.n_points
+        )
+        attention_weights = F.softmax(attention_weights, -1).view(
+            batch_size, num_queries, self.n_heads, self.n_levels, self.n_points
+        )
+        # batch_size, num_queries, n_heads, n_levels, n_points, 2
+        num_coordinates = reference_points.shape[-1]
+        if num_coordinates == 2:
+            offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1)
+            sampling_locations = (
+                reference_points[:, :, None, :, None, :]
+                + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
+            )
+        elif num_coordinates == 4:
+            sampling_locations = (
+                reference_points[:, :, None, :, None, :2]
+                + sampling_offsets / self.n_points * reference_points[:, :, None, :, None, 2:] * 0.5
+            )
+        else:
+            raise ValueError(f"Last dim of reference_points must be 2 or 4, but got {reference_points.shape[-1]}")
+
+        output = self.attn(
+            value,
+            spatial_shapes,
+            spatial_shapes_list,
+            level_start_index,
+            sampling_locations,
+            attention_weights,
+            self.im2col_step,
+        )
+
+        output = self.output_proj(output)
+
+        return output, attention_weights
+
+
+class LwDetrMLP(nn.Module):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__()
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.decoder_activation_function]
+        self.fc1 = nn.Linear(config.d_model, config.decoder_ffn_dim)
+        self.fc2 = nn.Linear(config.decoder_ffn_dim, config.d_model)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class LwDetrDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: LwDetrConfig, layer_idx: int):
+        nn.Module.__init__(self)
+
+        # self-attention
+        self.self_attn = LwDetrAttention(config, layer_idx=layer_idx)
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.decoder_activation_function]
+        self.activation_dropout = config.activation_dropout
+        self.self_attn_layer_norm = nn.LayerNorm(config.d_model)
+
+        # cross-attention
+        self.cross_attn = LwDetrMultiscaleDeformableAttention(
+            config,
+            num_heads=config.decoder_cross_attention_heads,
+            n_points=config.decoder_n_points,
+        )
+        self.cross_attn_layer_norm = nn.LayerNorm(config.d_model)
+
+        # mlp
+        self.mlp = LwDetrMLP(config)
+        self.layer_norm = nn.LayerNorm(config.d_model)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: torch.Tensor | None = None,
+        reference_points=None,
+        spatial_shapes=None,
+        spatial_shapes_list=None,
+        level_start_index=None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        self_attention_output, self_attn_weights = self.self_attn(
+            hidden_states, position_embeddings=position_embeddings, **kwargs
+        )
+
+        self_attention_output = nn.functional.dropout(self_attention_output, p=self.dropout, training=self.training)
+        hidden_states = hidden_states + self_attention_output
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+
+        cross_attention_output, cross_attn_weights = self.cross_attn(
+            hidden_states=hidden_states,
+            attention_mask=encoder_attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            position_embeddings=position_embeddings,
+            reference_points=reference_points,
+            spatial_shapes=spatial_shapes,
+            spatial_shapes_list=spatial_shapes_list,
+            level_start_index=level_start_index,
+            **kwargs,
+        )
+        cross_attention_output = nn.functional.dropout(cross_attention_output, p=self.dropout, training=self.training)
+        hidden_states = hidden_states + cross_attention_output
+        hidden_states = self.cross_attn_layer_norm(hidden_states)
+
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.layer_norm(hidden_states)
+
+        return hidden_states
+
+
+@auto_docstring
+class LwDetrPreTrainedModel(PreTrainedModel):
+    config: LwDetrConfig
+    base_model_prefix = "model"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    _no_split_modules = [
+        r"LwDetrConvEncoder",
+        r"LwDetrDecoderLayer",
+    ]
+    _supports_sdpa = True
+    _supports_flash_attn = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "attentions": [LwDetrAttention, LwDetrMultiscaleDeformableAttention],
+        "hidden_states": [LwDetrDecoderLayer],
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+
+        if isinstance(module, LwDetrMultiscaleDeformableAttention):
+            init.constant_(module.sampling_offsets.weight, 0.0)
+            thetas = torch.arange(module.n_heads, dtype=torch.int64).float() * (2.0 * math.pi / module.n_heads)
+            grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+            grid_init = (
+                (grid_init / grid_init.abs().max(-1, keepdim=True)[0])
+                .view(module.n_heads, 1, 1, 2)
+                .repeat(1, module.n_levels, module.n_points, 1)
+            )
+            for i in range(module.n_points):
+                grid_init[:, :, i, :] *= i + 1
+
+            init.copy_(module.sampling_offsets.bias, grid_init.view(-1))
+            init.constant_(module.attention_weights.weight, 0.0)
+            init.constant_(module.attention_weights.bias, 0.0)
+            init.xavier_uniform_(module.value_proj.weight)
+            init.constant_(module.value_proj.bias, 0.0)
+            init.xavier_uniform_(module.output_proj.weight)
+            init.constant_(module.output_proj.bias, 0.0)
+        if hasattr(module, "level_embed"):
+            init.normal_(module.level_embed)
+        if hasattr(module, "refpoint_embed") and module.refpoint_embed is not None:
+            init.constant_(module.refpoint_embed.weight, 0)
+        if hasattr(module, "class_embed") and module.class_embed is not None:
+            prior_prob = 0.01
+            bias_value = -math.log((1 - prior_prob) / prior_prob)
+            init.constant_(module.class_embed.bias, bias_value)
+        if hasattr(module, "bbox_embed") and module.bbox_embed is not None:
+            init.constant_(module.bbox_embed.layers[-1].weight, 0)
+            init.constant_(module.bbox_embed.layers[-1].bias, 0)
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for outputs of the LwDetrDecoder. This class adds two attributes to
+    BaseModelOutputWithCrossAttentions, namely:
+    - a stacked tensor of intermediate decoder hidden states (i.e. the output of each decoder layer)
+    - a stacked tensor of intermediate reference points.
+    """
+)
+@dataclass
+class LwDetrDecoderOutput(BaseModelOutputWithCrossAttentions):
+    r"""
+    cross_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` and `config.add_cross_attention=True` is passed or when `config.output_attentions=True`):
+        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+        sequence_length)`. Attentions weights of the decoder's cross-attention layer, after the attention softmax,
+        used to compute the weighted average in the cross-attention heads.
+    intermediate_hidden_states (`torch.FloatTensor` of shape `(config.decoder_layers, batch_size, num_queries, hidden_size)`, *optional*, returned when `config.auxiliary_loss=True`):
+        Intermediate decoder activations, i.e. the output of each decoder layer, each of them gone through a
+        layernorm.
+    intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, sequence_length, hidden_size)`):
+        Stacked intermediate reference points (reference points of each layer of the decoder).
+    """
+
+    intermediate_hidden_states: torch.FloatTensor | None = None
+
+    intermediate_reference_points: torch.FloatTensor | None = None
+
+
+def encode_sinusoidal_position_embedding(
+    pos_tensor: torch.Tensor,
+    num_pos_feats: int = 128,
+    temperature: int = 10000,
+) -> torch.Tensor:
+    """Sinusoidal position embeddings from normalized anchor coordinates.
+
+    Each coordinate in `pos_tensor` is independently encoded with ``num_pos_feats``
+    interleaved sin/cos components; per-coordinate embeddings are concatenated.
+    Handles 2-D ``(x, y)`` and N-D ``(x, y, w, h)`` inputs. For 2-D+ inputs the
+    x and y embeddings are swapped to follow the DETR ``[pos_y, pos_x, ...]`` convention.
+
+    Args:
+        pos_tensor: Normalized coordinates in ``[0, 1]``, shape ``(..., n_coords)``.
+        num_pos_feats: Embedding dimension per coordinate.
+        temperature: Base for the frequency decay.
+
+    Returns:
+        Tensor of shape ``(..., n_coords * num_pos_feats)``, same dtype as input.
+    """
+    scale = 2 * math.pi
+    dim_t = torch.arange(num_pos_feats, dtype=torch.float32, device=pos_tensor.device)
+    dim_t = temperature ** (2 * torch.div(dim_t, 2, rounding_mode="floor") / num_pos_feats)
+
+    coords = pos_tensor.unbind(-1)  # list of (...,) tensors
+    embeddings = [coord[..., None] * scale / dim_t for coord in coords]  # each (..., num_pos_feats)
+    embeddings = [
+        torch.stack((e[..., 0::2].sin(), e[..., 1::2].cos()), dim=-1).flatten(-2) for e in embeddings
+    ]  # each (..., num_pos_feats)
+
+    if len(embeddings) >= 2:
+        embeddings[0], embeddings[1] = embeddings[1], embeddings[0]
+
+    return torch.cat(embeddings, dim=-1).to(pos_tensor.dtype)
+
+
+class LwDetrDecoder(LwDetrPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DeformableDetrDecoderLayer`].
+
+    The decoder updates the query embeddings through multiple self-attention and deformable cross-attention layers.
+
+    Some tweaks for LwDetr:
+
+    - it uses group detr technique at training for faster convergence.
+
+    Args:
+        config: LwDetrConfig
+    """
+
+    _can_record_outputs = {
+        "hidden_states": LwDetrDecoderLayer,
+        "attentions": OutputRecorder(LwDetrAttention, layer_name="self_attn", index=1),
+        "cross_attentions": OutputRecorder(LwDetrMultiscaleDeformableAttention, layer_name="cross_attn", index=1),
+    }
+
+    def __init__(self, config: LwDetrConfig):
+        super().__init__(config)
+        self.dropout = config.dropout
+        self.layers = nn.ModuleList([LwDetrDecoderLayer(config, i) for i in range(config.decoder_layers)])
+        self.layernorm = nn.LayerNorm(config.d_model)
+
+        self.gradient_checkpointing = False
+
+        self.ref_point_head = LwDetrMLPPredictionHead(2 * config.d_model, config.d_model, config.d_model, num_layers=2)
+
+        self.post_init()
+
+    def get_reference(self, reference_points, valid_ratios):
+        # batch_size, num_queries, batch_size, 4
+        obj_center = reference_points[..., :4]
+
+        # batch_size, num_queries, num_levels, 4
+        reference_points_inputs = obj_center[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[:, None]
+
+        # batch_size, num_queries, d_model * 2
+        query_sine_embed = encode_sinusoidal_position_embedding(
+            reference_points_inputs[:, :, 0, :], num_pos_feats=self.config.d_model // 2
+        )
+
+        # batch_size, num_queries, d_model
+        query_pos = self.ref_point_head(query_sine_embed)
+        return reference_points_inputs, query_pos
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor | None = None,
+        reference_points: torch.Tensor | None = None,
+        spatial_shapes: torch.Tensor | None = None,
+        spatial_shapes_list: torch.Tensor | None = None,
+        level_start_index: torch.Tensor | None = None,
+        valid_ratios: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        intermediate = ()
+        intermediate_reference_points = (reference_points,)
+
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+
+        reference_points_inputs, query_pos = self.get_reference(reference_points, valid_ratios)
+
+        for idx, decoder_layer in enumerate(self.layers):
+            hidden_states = decoder_layer(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                position_embeddings=query_pos,
+                reference_points=reference_points_inputs,
+                spatial_shapes=spatial_shapes,
+                spatial_shapes_list=spatial_shapes_list,
+                level_start_index=level_start_index,
+                **kwargs,
+            )
+            intermediate_hidden_states = self.layernorm(hidden_states)
+            intermediate += (intermediate_hidden_states,)
+
+        intermediate = torch.stack(intermediate)
+        last_hidden_state = intermediate[-1]
+        intermediate_reference_points = torch.stack(intermediate_reference_points)
+
+        return LwDetrDecoderOutput(
+            last_hidden_state=last_hidden_state,
+            intermediate_hidden_states=intermediate,
+            intermediate_reference_points=intermediate_reference_points,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for outputs of the LwDetr backbone-decoder model.
+    """
+)
+@dataclass
+class LwDetrModelOutput(ModelOutput):
+    r"""
+    init_reference_points (`torch.FloatTensor` of shape  `(batch_size, num_queries, 4)`):
+        Initial reference points sent through the Transformer decoder.
+    intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
+        Stacked intermediate hidden states (output of each layer of the decoder).
+    intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+        Stacked intermediate reference points (reference points of each layer of the decoder).
+    enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+        Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
+        picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
+        foreground and background).
+    enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+        Logits of predicted bounding boxes coordinates in the first stage.
+    """
+
+    init_reference_points: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    intermediate_hidden_states: torch.FloatTensor | None = None
+    intermediate_reference_points: torch.FloatTensor | None = None
+    enc_outputs_class: torch.FloatTensor | None = None
+    enc_outputs_coord_logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    cross_attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+def refine_bboxes(reference_points, deltas):
+    reference_points = reference_points.to(deltas.device)
+    new_reference_points_cxcy = deltas[..., :2] * reference_points[..., 2:] + reference_points[..., :2]
+    new_reference_points_wh = deltas[..., 2:].exp() * reference_points[..., 2:]
+    new_reference_points = torch.cat((new_reference_points_cxcy, new_reference_points_wh), -1)
+    return new_reference_points
+
+
+@auto_docstring(
+    custom_intro="""
+    The bare LW Detr Model (consisting of a backbone and decoder Transformer) outputting raw
+    hidden-states without any specific head on top.
+    """
+)
+class LwDetrModel(LwDetrPreTrainedModel):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__(config)
+
+        # Create backbone + positional encoding
+        self.backbone = LwDetrConvEncoder(config)
+
+        self.group_detr = config.group_detr
+        self.num_queries = config.num_queries
+        hidden_dim = config.d_model
+        self.reference_point_embed = nn.Embedding(self.num_queries * self.group_detr, 4)
+        self.query_feat = nn.Embedding(self.num_queries * self.group_detr, hidden_dim)
+
+        self.decoder = LwDetrDecoder(config)
+
+        self.enc_output = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim) for _ in range(self.group_detr)])
+        self.enc_output_norm = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(self.group_detr)])
+        # Should normally be None and then instantiated in the ForObjectDetection class
+        self.enc_out_bbox_embed = nn.ModuleList(
+            [LwDetrMLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3) for _ in range(self.group_detr)]
+        )
+        self.enc_out_class_embed = nn.ModuleList(
+            [nn.Linear(config.d_model, config.num_labels) for _ in range(self.group_detr)]
+        )
+
+        self.post_init()
+
+    def freeze_backbone(self):
+        for name, param in self.backbone.model.named_parameters():
+            param.requires_grad_(False)
+
+    def unfreeze_backbone(self):
+        for name, param in self.backbone.model.named_parameters():
+            param.requires_grad_(True)
+
+    def get_valid_ratio(self, mask, dtype=torch.float32):
+        """Get the valid ratio of all feature maps."""
+
+        _, height, width = mask.shape
+        valid_height = torch.sum(mask[:, :, 0], 1)
+        valid_width = torch.sum(mask[:, 0, :], 1)
+        valid_ratio_height = valid_height.to(dtype) / height
+        valid_ratio_width = valid_width.to(dtype) / width
+        valid_ratio = torch.stack([valid_ratio_width, valid_ratio_height], -1)
+        return valid_ratio
+
+    def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes):
+        """Generate the encoder output proposals from encoded enc_output.
+
+        Args:
+            enc_output (Tensor[batch_size, sequence_length, hidden_size]): Output of the encoder.
+            padding_mask (Tensor[batch_size, sequence_length]): Padding mask for `enc_output`.
+            spatial_shapes (list[tuple[int, int]]): Spatial shapes of the feature maps.
+
+        Returns:
+            `tuple(torch.FloatTensor)`: A tuple of feature map and bbox prediction.
+                - object_query (Tensor[batch_size, sequence_length, hidden_size]): Object query features. Later used to
+                  directly predict a bounding box. (without the need of a decoder)
+                - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals in [0, 1] space.
+                  Invalid positions (padding or out-of-bounds) are filled with 0.
+                - invalid_mask (Tensor[batch_size, sequence_length, 1]): Boolean mask that is True for invalid positions
+                  (padded pixels or proposals whose coordinates fall outside (0.01, 0.99)).
+        """
+        batch_size = enc_output.shape[0]
+        proposals = []
+        _cur = 0
+        for level, (height, width) in enumerate(spatial_shapes):
+            mask_flatten_ = padding_mask[:, _cur : (_cur + height * width)].view(batch_size, height, width, 1)
+            valid_height = torch.sum(~mask_flatten_[:, :, 0, 0], 1)
+            valid_width = torch.sum(~mask_flatten_[:, 0, :, 0], 1)
+
+            grid_y, grid_x = torch.meshgrid(
+                torch.linspace(
+                    0,
+                    height - 1,
+                    height,
+                    dtype=enc_output.dtype,
+                    device=enc_output.device,
+                ),
+                torch.linspace(
+                    0,
+                    width - 1,
+                    width,
+                    dtype=enc_output.dtype,
+                    device=enc_output.device,
+                ),
+                indexing="ij",
+            )
+            grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1)
+
+            scale = torch.cat([valid_width.unsqueeze(-1), valid_height.unsqueeze(-1)], 1).view(batch_size, 1, 1, 2)
+            grid = (grid.unsqueeze(0).expand(batch_size, -1, -1, -1) + 0.5) / scale
+            width_height = torch.ones_like(grid) * 0.05 * (2.0**level)
+            proposal = torch.cat((grid, width_height), -1).view(batch_size, -1, 4)
+            proposals.append(proposal)
+            _cur += height * width
+        output_proposals = torch.cat(proposals, 1)
+        output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True)
+        invalid_mask = padding_mask.unsqueeze(-1) | ~output_proposals_valid
+        output_proposals = output_proposals.masked_fill(invalid_mask, float(0))
+
+        # assign each pixel as an object query
+        object_query = enc_output
+        object_query = object_query.masked_fill(invalid_mask, float(0))
+        return object_query, output_proposals, invalid_mask
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor = None,
+        pixel_mask: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> LwDetrModelOutput:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, DeformableDetrModel
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> image_processor = AutoImageProcessor.from_pretrained("AnnaZhang/lwdetr_small_60e_coco")
+        >>> model = DeformableDetrModel.from_pretrained("AnnaZhang/lwdetr_small_60e_coco")
+
+        >>> inputs = image_processor(images=image, return_tensors="pt")
+
+        >>> outputs = model(**inputs)
+
+        >>> last_hidden_states = outputs.last_hidden_state
+        >>> list(last_hidden_states.shape)
+        [1, 300, 256]
+        ```"""
+        batch_size, num_channels, height, width = pixel_values.shape
+        device = pixel_values.device
+
+        if pixel_mask is None:
+            pixel_mask = torch.ones(((batch_size, height, width)), dtype=torch.long, device=device)
+
+        # Extract multi-scale feature maps of same resolution `config.d_model` (cf Figure 4 in paper)
+        # First, sent pixel_values + pixel_mask through Backbone to obtain the features
+        # which is a list of tuples
+        features = self.backbone(pixel_values, pixel_mask)
+
+        # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
+        sources = []
+        masks = []
+        for level, (source, mask) in enumerate(features):
+            sources.append(source)
+            masks.append(mask)
+            if mask is None:
+                raise ValueError("No attention mask was provided")
+
+        if self.training:
+            reference_points = self.reference_point_embed.weight
+            query_feat = self.query_feat.weight
+        else:
+            # only use one group in inference
+            reference_points = self.reference_point_embed.weight[: self.num_queries]
+            query_feat = self.query_feat.weight[: self.num_queries]
+
+        # Prepare encoder inputs (by flattening)
+        source_flatten = []
+        mask_flatten = []
+        spatial_shapes_list = []
+        for source, mask in zip(sources, masks):
+            batch_size, num_channels, height, width = source.shape
+            spatial_shape = (height, width)
+            spatial_shapes_list.append(spatial_shape)
+            source = source.flatten(2).transpose(1, 2)
+            mask = mask.flatten(1)
+            source_flatten.append(source)
+            mask_flatten.append(mask)
+        source_flatten = torch.cat(source_flatten, 1)
+        mask_flatten = torch.cat(mask_flatten, 1)
+        spatial_shapes = torch.as_tensor(spatial_shapes_list, dtype=torch.long, device=source_flatten.device)
+        level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
+        valid_ratios = torch.stack([self.get_valid_ratio(m, dtype=source_flatten.dtype) for m in masks], 1)
+
+        target = query_feat.unsqueeze(0).expand(batch_size, -1, -1)
+        reference_points = reference_points.unsqueeze(0).expand(batch_size, -1, -1)
+
+        object_query_embedding, output_proposals, invalid_mask = self.gen_encoder_output_proposals(
+            source_flatten, ~mask_flatten, spatial_shapes_list
+        )
+
+        group_detr = self.group_detr if self.training else 1
+        topk = self.num_queries
+        topk_coords_logits = []
+        topk_coords_logits_undetach = []
+        object_query_undetach = []
+
+        for group_id in range(group_detr):
+            group_object_query = self.enc_output[group_id](object_query_embedding)
+            group_object_query = self.enc_output_norm[group_id](group_object_query)
+
+            group_enc_outputs_class = self.enc_out_class_embed[group_id](group_object_query)
+            group_enc_outputs_class = group_enc_outputs_class.masked_fill(invalid_mask, float("-inf"))
+            group_delta_bbox = self.enc_out_bbox_embed[group_id](group_object_query)
+            group_enc_outputs_coord = refine_bboxes(output_proposals, group_delta_bbox)
+
+            group_topk_proposals = torch.topk(group_enc_outputs_class.max(-1)[0], topk, dim=1)[1]
+            group_topk_coords_logits_undetach = torch.gather(
+                group_enc_outputs_coord,
+                1,
+                group_topk_proposals.unsqueeze(-1).repeat(1, 1, 4),
+            )
+            group_topk_coords_logits = group_topk_coords_logits_undetach.detach()
+            group_object_query_undetach = torch.gather(
+                group_object_query, 1, group_topk_proposals.unsqueeze(-1).repeat(1, 1, self.config.d_model)
+            )
+
+            topk_coords_logits.append(group_topk_coords_logits)
+            topk_coords_logits_undetach.append(group_topk_coords_logits_undetach)
+            object_query_undetach.append(group_object_query_undetach)
+
+        topk_coords_logits = torch.cat(topk_coords_logits, 1)
+        topk_coords_logits_undetach = torch.cat(topk_coords_logits_undetach, 1)
+        object_query_undetach = torch.cat(object_query_undetach, 1)
+
+        enc_outputs_class = object_query_undetach
+        enc_outputs_coord_logits = topk_coords_logits_undetach
+
+        reference_points = refine_bboxes(topk_coords_logits, reference_points)
+
+        init_reference_points = reference_points
+        decoder_outputs = self.decoder(
+            inputs_embeds=target,
+            reference_points=reference_points,
+            spatial_shapes=spatial_shapes,
+            spatial_shapes_list=spatial_shapes_list,
+            level_start_index=level_start_index,
+            valid_ratios=valid_ratios,
+            encoder_hidden_states=source_flatten,
+            encoder_attention_mask=mask_flatten,
+            **kwargs,
+        )
+
+        return LwDetrModelOutput(
+            init_reference_points=init_reference_points,
+            last_hidden_state=decoder_outputs.last_hidden_state,
+            intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
+            intermediate_reference_points=decoder_outputs.intermediate_reference_points,
+            enc_outputs_class=enc_outputs_class,
+            enc_outputs_coord_logits=enc_outputs_coord_logits,
+            hidden_states=decoder_outputs.hidden_states,
+            attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+        )
+
+
+class LwDetrMLPPredictionHead(nn.Module):
+    """
+    Very simple multi-layer perceptron (MLP, also called FFN), used to predict the normalized center coordinates,
+    height and width of a bounding box w.r.t. an image.
+
+    """
+
+    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
+        super().__init__()
+        self.num_layers = num_layers
+        h = [hidden_dim] * (num_layers - 1)
+        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
+
+    def forward(self, x):
+        for i, layer in enumerate(self.layers):
+            x = nn.functional.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
+        return x
+
+
+@auto_docstring(
+    custom_intro="""
+    Output type of [`LwDetrForObjectDetection`].
+    """
+)
+@dataclass
+class LwDetrObjectDetectionOutput(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
+        Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
+        bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
+        scale-invariant IoU loss.
+    loss_dict (`Dict`, *optional*):
+        A dictionary containing the individual losses. Useful for logging.
+    logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
+        Classification logits (including no-object) for all queries.
+    pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
+        Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
+        values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
+        possible padding). You can use [`~DeformableDetrProcessor.post_process_object_detection`] to retrieve the
+        unnormalized bounding boxes.
+    auxiliary_outputs (`list[Dict]`, *optional*):
+        Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
+        and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
+        `pred_boxes`) for each decoder layer.
+    init_reference_points (`torch.FloatTensor` of shape  `(batch_size, num_queries, 4)`):
+        Initial reference points sent through the Transformer decoder.
+    intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
+        Stacked intermediate hidden states (output of each layer of the decoder).
+    intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+        Stacked intermediate reference points (reference points of each layer of the decoder).
+    enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+        Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
+        picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
+        foreground and background).
+    enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+        Logits of predicted bounding boxes coordinates in the first stage.
+    """
+
+    loss: torch.FloatTensor | None = None
+    loss_dict: dict | None = None
+    logits: torch.FloatTensor | None = None
+    pred_boxes: torch.FloatTensor | None = None
+    auxiliary_outputs: list[dict] | None = None
+    init_reference_points: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    intermediate_hidden_states: torch.FloatTensor | None = None
+    intermediate_reference_points: torch.FloatTensor | None = None
+    enc_outputs_class: Any = None
+    enc_outputs_coord_logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    cross_attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    LW DETR Model (consisting of a backbone and decoder Transformer) with object detection heads on
+    top, for tasks such as COCO detection.
+    """
+)
+class LwDetrForObjectDetection(LwDetrPreTrainedModel):
+    # When using clones, all layers > 0 will be clones, but layer 0 *is* required
+    # We can't initialize the model on meta device as some weights are modified during the initialization
+    _no_split_modules = None
+    _tied_weights_keys = None
+
+    def __init__(self, config: LwDetrConfig):
+        super().__init__(config)
+        self.model = LwDetrModel(config)
+        self.class_embed = nn.Linear(config.d_model, config.num_labels)
+        self.bbox_embed = LwDetrMLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3)
+
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor = None,
+        pixel_mask: torch.LongTensor | None = None,
+        labels: list[dict] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> LwDetrObjectDetectionOutput:
+        r"""
+        labels (`list[Dict]` of len `(batch_size,)`, *optional*):
+            Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
+            following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
+            respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
+            in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.
+
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, LwDetrForObjectDetection
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> image_processor = AutoImageProcessor.from_pretrained("AnnaZhang/lwdetr_small_60e_coco")
+        >>> model = LwDetrForObjectDetection.from_pretrained("AnnaZhang/lwdetr_small_60e_coco")
+
+        >>> inputs = image_processor(images=image, return_tensors="pt")
+        >>> outputs = model(**inputs)
+
+        >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
+        >>> target_sizes = torch.tensor([image.size[::-1]])
+        >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[
+        ...     0
+        ... ]
+        >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        ...     box = [round(i, 2) for i in box.tolist()]
+        ...     print(
+        ...         f"Detected {model.config.id2label[label.item()]} with confidence "
+        ...         f"{round(score.item(), 3)} at location {box}"
+        ...     )
+        Detected cat with confidence 0.8 at location [16.5, 52.84, 318.25, 470.78]
+        Detected cat with confidence 0.789 at location [342.19, 24.3, 640.02, 372.25]
+        Detected remote with confidence 0.633 at location [40.79, 72.78, 176.76, 117.25]
+        ```"""
+        outputs = self.model(
+            pixel_values,
+            pixel_mask=pixel_mask,
+            **kwargs,
+        )
+
+        last_hidden_states = outputs.last_hidden_state
+        intermediate_reference_points = outputs.intermediate_reference_points
+        enc_outputs_class_logits = outputs.enc_outputs_class
+        enc_outputs_boxes_logits = outputs.enc_outputs_coord_logits
+
+        logits = self.class_embed(last_hidden_states)
+        pred_boxes_delta = self.bbox_embed(last_hidden_states)
+        pred_boxes = refine_bboxes(intermediate_reference_points[-1], pred_boxes_delta)
+
+        enc_outputs_class_logits_list = enc_outputs_class_logits.split(self.config.num_queries, dim=1)
+        pred_class = []
+        group_detr = self.config.group_detr if self.training else 1
+        for group_index in range(group_detr):
+            group_pred_class = self.model.enc_out_class_embed[group_index](enc_outputs_class_logits_list[group_index])
+            pred_class.append(group_pred_class)
+        enc_outputs_class_logits = torch.cat(pred_class, dim=1)
+
+        loss, loss_dict, auxiliary_outputs = None, None, None
+        if labels is not None:
+            outputs_class, outputs_coord = None, None
+            if self.config.auxiliary_loss:
+                intermediate_hidden_states = outputs.intermediate_hidden_states
+                outputs_coord_delta = self.bbox_embed(intermediate_hidden_states)
+                outputs_coord = refine_bboxes(intermediate_reference_points, outputs_coord_delta)
+                outputs_class = self.class_embed(intermediate_hidden_states)
+
+            loss, loss_dict, auxiliary_outputs = self.loss_function(
+                logits,
+                labels,
+                self.device,
+                pred_boxes,
+                self.config,
+                outputs_class,
+                outputs_coord,
+                enc_outputs_class_logits,
+                enc_outputs_boxes_logits,
+            )
+
+        return LwDetrObjectDetectionOutput(
+            loss=loss,
+            loss_dict=loss_dict,
+            logits=logits,
+            pred_boxes=pred_boxes,
+            auxiliary_outputs=auxiliary_outputs,
+            last_hidden_state=outputs.last_hidden_state,
+            intermediate_hidden_states=outputs.intermediate_hidden_states,
+            intermediate_reference_points=outputs.intermediate_reference_points,
+            init_reference_points=outputs.init_reference_points,
+            enc_outputs_class=enc_outputs_class_logits,
+            enc_outputs_coord_logits=enc_outputs_boxes_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
+
+__all__ = [
+    "LwDetrPreTrainedModel",
+    "LwDetrModel",
+    "LwDetrForObjectDetection",
+    "LwDetrViTPreTrainedModel",
+    "LwDetrViTBackbone",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/lw_detr/modular_lw_detr.py

@@ -0,0 +1,1398 @@
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from huggingface_hub.dataclasses import strict
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2FN
+from ...backbone_utils import consolidate_backbone_kwargs_to_config
+from ...configuration_utils import PreTrainedConfig
+from ...modeling_layers import GradientCheckpointingLayer
+from ...modeling_outputs import BackboneOutput, BaseModelOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging
+from ...utils.generic import can_return_tuple, merge_with_config_defaults
+from ...utils.output_capturing import OutputRecorder, capture_outputs
+from ..auto import AutoConfig
+from ..conditional_detr.modeling_conditional_detr import encode_sinusoidal_position_embedding
+from ..convnext.modeling_convnext import ConvNextLayerNorm
+from ..deformable_detr.modeling_deformable_detr import (
+    DeformableDetrDecoderOutput,
+    DeformableDetrForObjectDetection,
+    DeformableDetrMLPPredictionHead,
+    DeformableDetrModel,
+    DeformableDetrMultiscaleDeformableAttention,
+)
+from ..llama.modeling_llama import eager_attention_forward
+from ..rt_detr.modeling_rt_detr import RTDetrConvNormLayer
+from ..vit.modeling_vit import ViTAttention
+from ..vitdet.configuration_vitdet import VitDetConfig
+from ..vitdet.modeling_vitdet import (
+    VitDetBackbone,
+    VitDetEmbeddings,
+    VitDetMlp,
+    VitDetPreTrainedModel,
+)
+
+
+logger = logging.get_logger(__name__)
+
+
+@auto_docstring(checkpoint="AnnaZhang/lwdetr_small_60e_coco")
+@strict
+class LwDetrViTConfig(VitDetConfig):
+    r"""
+    pretrain_image_size (`int`, *optional*, defaults to 224):
+        The size (resolution) of each image during pretraining.
+    window_block_indices (`list[int]`, *optional*, defaults to `[]`):
+        List of indices of blocks that should have window attention instead of regular global self-attention.
+    use_absolute_position_embeddings (`bool`, *optional*, defaults to `True`):
+        Whether to add absolute position embeddings to the patch embeddings.
+    cae_init_values (`float`, *optional*, defaults to 0.1):
+        Initialization value for CAE parameters when `use_cae` is enabled.
+    num_windows (`int`, *optional*, defaults to 16):
+        Number of windows for window-based attention. Must be a perfect square and the image size must be
+        divisible by the square root of this value. This enables efficient window-major feature map organization.
+
+    Example:
+
+    ```python
+    >>> from transformers import LwDetrViTConfig, LwDetrViTModel
+
+    >>> # Initializing a LW-DETR ViT configuration
+    >>> configuration = LwDetrViTConfig()
+
+    >>> # Initializing a model (with random weights) from the configuration
+    >>> model = LwDetrViTModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "lw_detr_vit"
+
+    image_size: int | list[int] | tuple[int, int] = 256
+    cae_init_values: float = 0.1
+    num_windows: int = 16
+
+    residual_block_indices = AttributeError()
+    use_relative_position_embeddings = AttributeError()
+    window_size = AttributeError()
+    drop_path_rate = AttributeError()
+
+    def __post_init__(self, **kwargs):
+        self.num_windows_side = int(math.sqrt(self.num_windows))
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of `@strict`-powered validation. Validates the architecture of the config."""
+        if self.num_windows % math.sqrt(self.num_windows) != 0:
+            raise ValueError(
+                f"`num_windows` has to be a perfect square, where num_windows % math.sqrt(num_windows) != 0, but got {self.num_windows}."
+            )
+        if self.image_size / self.num_windows % math.sqrt(self.num_windows) != 0:
+            raise ValueError(
+                f"`image_size` has to be divisible by `num_windows`, where image_size / num_windows % math.sqrt(num_windows) != 0,but got {self.image_size} and {self.num_windows}."
+            )
+
+
+@auto_docstring(checkpoint="AnnaZhang/lwdetr_small_60e_coco")
+@strict
+class LwDetrConfig(PreTrainedConfig):
+    r"""
+    projector_scale_factors (`list[float]`, *optional*, defaults to `[]`):
+        Scale factors for the feature pyramid network. Each scale factor determines the resolution of features
+        at different levels. Supported values are 0.5, 1.0, and 2.0.
+    hidden_expansion (`float`, *optional*, defaults to 0.5):
+        Expansion factor for hidden dimensions in the projector layers.
+    c2f_num_blocks (`int`, *optional*, defaults to 3):
+        Number of blocks in the C2F layer.
+    activation_function (`str`, *optional*, defaults to `"silu"`):
+        The non-linear activation function in the projector. Supported values are `"silu"`, `"relu"`, `"gelu"`.
+    batch_norm_eps (`float`, *optional*, defaults to 1e-05):
+        The epsilon value for batch normalization layers.
+    decoder_ffn_dim (`int`, *optional*, defaults to 2048):
+        Dimension of the "intermediate" (often named feed-forward) layer in decoder.
+    decoder_n_points (`int`, *optional*, defaults to 4):
+        The number of sampled keys in each feature level for each attention head in the decoder.
+    decoder_self_attention_heads (`int`, *optional*, defaults to 8):
+        Number of attention heads for each attention layer in the decoder self-attention.
+    decoder_cross_attention_heads (`int`, *optional*, defaults to 16):
+        Number of attention heads for each attention layer in the decoder cross-attention.
+    decoder_activation_function (`str`, *optional*, defaults to `"relu"`):
+        The non-linear activation function in the decoder. Supported values are `"relu"`, `"silu"`, `"gelu"`.
+    num_queries (`int`, *optional*, defaults to 300):
+        Number of object queries, i.e. detection slots. This is the maximal number of objects
+        [`LwDetrModel`] can detect in a single image.
+    group_detr (`int`, *optional*, defaults to 13):
+        Number of groups for Group DETR attention mechanism, which helps reduce computational complexity.
+    disable_custom_kernels (`bool`, *optional*, defaults to `True`):
+        Disable the use of custom CUDA and CPU kernels. This option is necessary for the ONNX export, as custom
+        kernels are not supported by PyTorch ONNX export.
+    class_loss_coefficient (`float`, *optional*, defaults to 1):
+        Relative weight of the classification loss in the Hungarian matching cost.
+
+    Examples:
+
+    ```python
+    >>> from transformers import LwDetrConfig, LwDetrModel
+
+    >>> # Initializing a LW-DETR AnnaZhang/lwdetr_small_60e_coco style configuration
+    >>> configuration = LwDetrConfig()
+
+    >>> # Initializing a model (with random weights) from the AnnaZhang/lwdetr_small_60e_coco style configuration
+    >>> model = LwDetrModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+
+    model_type = "lw_detr"
+    sub_configs = {"backbone_config": AutoConfig}
+
+    backbone_config: dict | PreTrainedConfig | None = None
+    projector_scale_factors: list[float] | tuple[float, ...] = ()
+    hidden_expansion: float = 0.5
+    c2f_num_blocks: int = 3
+    activation_function: str = "silu"
+    batch_norm_eps: float = 1e-5
+    dropout: float | int = 0.0
+    decoder_ffn_dim: int = 2048
+    decoder_n_points: int = 4
+    decoder_layers: int = 3
+    decoder_self_attention_heads: int = 8
+    decoder_cross_attention_heads: int = 16
+    decoder_activation_function: str = "relu"
+    num_queries: int = 300
+    attention_bias: bool = True
+    attention_dropout: float | int = 0.0
+    activation_dropout: float | int = 0.0
+    group_detr: int = 13
+    init_std: float = 0.02
+    disable_custom_kernels: bool = True
+    class_cost: int | float = 2
+    bbox_cost: int | float = 5
+    giou_cost: int | float = 2
+    class_loss_coefficient: int | float = 1
+    dice_loss_coefficient: int | float = 1
+    bbox_loss_coefficient: int | float = 5
+    giou_loss_coefficient: int | float = 2
+    eos_coefficient: float = 0.1
+    focal_alpha: float = 0.25
+    auxiliary_loss: bool = True
+    d_model: int = 256
+
+    def __post_init__(self, **kwargs):
+        if "mask_loss_coefficient" in kwargs:
+            logger.warning_once(
+                "The parameter `mask_loss_coefficient` was renamed to `class_loss_coefficient` in LW-DETR. "
+                "Please use `class_loss_coefficient` instead. `mask_loss_coefficient` will be removed in a future version."
+            )
+            self.class_loss_coefficient = kwargs.pop("mask_loss_coefficient")
+
+        self.backbone_config, kwargs = consolidate_backbone_kwargs_to_config(
+            backbone_config=self.backbone_config,
+            default_config_type="lw_detr_vit",
+            default_config_kwargs={
+                "image_size": 1024,
+                "hidden_size": 192,
+                "num_hidden_layers": 10,
+                "window_block_indices": [0, 1, 3, 6, 7, 9],
+                "out_indices": [2, 4, 5, 9],
+            },
+            **kwargs,
+        )
+
+        self.projector_in_channels = [self.d_model] * len(self.projector_scale_factors)
+        self.projector_out_channels = self.d_model
+        self.num_feature_levels = len(self.projector_scale_factors)
+        super().__post_init__(**kwargs)
+
+    def validate_architecture(self):
+        """Part of `@strict`-powered validation. Validates the architecture of the config."""
+        for scale in self.projector_scale_factors:
+            if scale not in [0.5, 1.0, 2.0]:
+                raise ValueError(f"Unsupported scale factor: {scale}")
+
+
+class LwDetrViTAttention(ViTAttention):
+    """LwDetr ViT attention with k_proj bias=False and dropout from config.dropout_prob."""
+
+    def __init__(self, config: LwDetrViTConfig):
+        super().__init__()
+        self.attention_dropout = config.dropout_prob
+        self.k_proj = nn.Linear(config.hidden_size, config.num_attention_heads * self.head_dim, bias=False)
+        self.num_key_value_groups = 1
+
+
+class LwDetrViTMlp(VitDetMlp):
+    pass
+
+
+class LwDetrViTLayer(GradientCheckpointingLayer):
+    def __init__(
+        self,
+        config: LwDetrViTConfig,
+        layer_idx,
+    ) -> None:
+        super().__init__()
+
+        dim = config.hidden_size
+        self.attention = LwDetrViTAttention(config)
+        self.intermediate = LwDetrViTMlp(config=config, in_features=dim, hidden_features=int(dim * config.mlp_ratio))
+        self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+        self.gamma_1 = nn.Parameter(torch.Tensor(dim), requires_grad=True)
+        self.gamma_2 = nn.Parameter(torch.Tensor(dim), requires_grad=True)
+
+        self.window = layer_idx in config.window_block_indices
+        self.num_windows = config.num_windows
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> torch.Tensor:
+        batch_size, seq_len, channels = hidden_states.shape
+        hidden_states_norm = self.layernorm_before(hidden_states)
+
+        if not self.window:
+            hidden_states_norm = hidden_states_norm.reshape(
+                batch_size // self.num_windows, self.num_windows * seq_len, channels
+            )
+
+        attention_output, _ = self.attention(hidden_states_norm, **kwargs)
+        attention_output = attention_output * self.gamma_1
+
+        if not self.window:
+            attention_output = attention_output.reshape(batch_size, seq_len, channels)
+
+        hidden_states = hidden_states + attention_output
+
+        layer_output = self.layernorm_after(hidden_states)
+        layer_output = self.intermediate(layer_output)
+        layer_output = layer_output * self.gamma_2
+
+        hidden_states = hidden_states + layer_output
+
+        return hidden_states
+
+
+class LwDetrViTEmbeddings(VitDetEmbeddings):
+    pass
+
+
+class LwDetrViTPreTrainedModel(VitDetPreTrainedModel):
+    config: LwDetrViTConfig
+    base_model_prefix = "lw_detr_vit"
+    main_input_name = "pixel_values"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["LwDetrViTEmbeddings", "LwDetrViTLayer"]
+    _supports_sdpa = True
+    _supports_flash_attn = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "hidden_states": LwDetrViTLayer,
+        "attentions": LwDetrViTAttention,
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module) -> None:
+        if isinstance(module, (nn.Linear, nn.Conv2d)):
+            init.trunc_normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                init.zeros_(module.bias)
+        elif isinstance(module, nn.LayerNorm):
+            init.zeros_(module.bias)
+            init.ones_(module.weight)
+        elif isinstance(module, LwDetrViTEmbeddings):
+            init.trunc_normal_(module.position_embeddings, mean=0.0, std=self.config.initializer_range)
+        if isinstance(module, LwDetrViTLayer):
+            init.constant_(module.gamma_1, self.config.cae_init_values)
+            init.constant_(module.gamma_2, self.config.cae_init_values)
+
+
+class LwDetrViTEncoder(LwDetrViTPreTrainedModel):
+    def __init__(self, config: LwDetrViTConfig):
+        super().__init__(config)
+        self.layer = nn.ModuleList([LwDetrViTLayer(config, idx) for idx in range(config.num_hidden_layers)])
+        self.post_init()
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> BaseModelOutput:
+        for layer_module in self.layer:
+            hidden_states = layer_module(hidden_states, **kwargs)
+
+        return BaseModelOutput(last_hidden_state=hidden_states)
+
+
+@auto_docstring()
+class LwDetrViTBackbone(VitDetBackbone):
+    def forward(self, pixel_values: torch.Tensor, **kwargs: Unpack[TransformersKwargs]) -> BackboneOutput:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import LwDetrViTConfig, LwDetrViTBackbone
+        >>> import torch
+
+        >>> config = LwDetrViTConfig()
+        >>> model = LwDetrViTBackbone(config)
+
+        >>> pixel_values = torch.randn(1, 3, 224, 224)
+
+        >>> with torch.no_grad():
+        ...     outputs = model(pixel_values)
+
+        >>> feature_maps = outputs.feature_maps
+        >>> list(feature_maps[-1].shape)
+        [1, 768, 14, 14]
+        ```"""
+        embedding_output = self.embeddings(pixel_values)
+
+        batch_size, channels, height, width = embedding_output.shape
+        # (batch_size, channels, height, width) -> (batch_size, height, width, channels)
+        hidden_states = embedding_output.permute(0, 2, 3, 1)
+
+        window_height = height // self.config.num_windows_side
+        window_width = width // self.config.num_windows_side
+        # (batch_size, height, width, channels) -> (batch_size*num_windows_side**2, window_height*window_width, channels)
+        hidden_states = (
+            hidden_states.reshape(
+                batch_size,
+                self.config.num_windows_side,
+                window_height,
+                self.config.num_windows_side,
+                window_width,
+                channels,
+            )
+            .permute(0, 1, 3, 2, 4, 5)
+            .reshape(batch_size * self.config.num_windows_side**2, window_height * window_width, channels)
+        )
+
+        kwargs["output_hidden_states"] = True  # required to extract layers for the stages
+        output = self.encoder(hidden_states, **kwargs)
+
+        feature_maps = ()
+        hidden_states = output.hidden_states
+        for stage, hidden_state in zip(self.stage_names, hidden_states):
+            if stage in self.out_features:
+                hidden_state = (
+                    hidden_state.reshape(
+                        batch_size,
+                        self.config.num_windows_side,
+                        self.config.num_windows_side,
+                        window_height,
+                        window_width,
+                        channels,
+                    )
+                    .permute(0, 5, 1, 3, 2, 4)
+                    .reshape(batch_size, channels, height, width)
+                )
+                feature_maps += (hidden_state,)
+
+        return BackboneOutput(
+            feature_maps=feature_maps, hidden_states=output.hidden_states, attentions=output.attentions
+        )
+
+
+class LwDetrConvNormLayer(RTDetrConvNormLayer):
+    def __init__(
+        self,
+        config: LwDetrConfig,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int,
+        activation: str | None = None,
+    ):
+        super().__init__(config, in_channels, out_channels, kernel_size, stride, activation)
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride,
+            padding=kernel_size // 2,
+            bias=False,
+        )
+
+
+class LwDetrRepVggBlock(nn.Module):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__()
+        hidden_channels = int(config.d_model * config.hidden_expansion)
+        self.conv1 = LwDetrConvNormLayer(
+            config, hidden_channels, hidden_channels, 3, 1, activation=config.activation_function
+        )
+        self.conv2 = LwDetrConvNormLayer(
+            config, hidden_channels, hidden_channels, 3, 1, activation=config.activation_function
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = self.conv1(x)
+        y = self.conv2(y)
+        return y
+
+
+class LwDetrC2FLayer(nn.Module):
+    # Inspired by RTDetrCSPRepLayer
+    def __init__(self, config: LwDetrConfig, in_channels: int):
+        super().__init__()
+        num_blocks = config.c2f_num_blocks
+        activation = config.activation_function
+        out_channels = config.d_model
+
+        self.hidden_channels = int(out_channels * config.hidden_expansion)
+
+        conv1_out_channels = 2 * self.hidden_channels
+        self.conv1 = LwDetrConvNormLayer(config, in_channels, conv1_out_channels, 1, 1, activation=activation)
+
+        conv2_in_channels = (2 + num_blocks) * self.hidden_channels
+        self.conv2 = LwDetrConvNormLayer(config, conv2_in_channels, out_channels, 1, 1, activation=activation)
+
+        self.bottlenecks = nn.ModuleList(LwDetrRepVggBlock(config) for _ in range(num_blocks))
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.conv1(hidden_states)
+        all_hidden_states = list(hidden_states.split(self.hidden_channels, 1))
+        hidden_states = all_hidden_states[-1]
+        hidden_states = hidden_states.contiguous()
+
+        for bottleneck in self.bottlenecks:
+            hidden_states = bottleneck(hidden_states)
+            all_hidden_states.append(hidden_states)
+
+        hidden_states = torch.cat(all_hidden_states, 1)
+        hidden_states = self.conv2(hidden_states)
+        return hidden_states
+
+
+class LwDetrLayerNorm(ConvNextLayerNorm):
+    pass
+
+
+class LwDetrSamplingLayer(nn.Module):
+    def __init__(self, config: LwDetrConfig, channel_size: int, scale: float):
+        super().__init__()
+
+        self.scale = scale
+        self.channel_size = channel_size
+
+        layers = []
+        if scale == 2.0:
+            if channel_size > 512:
+                layers.append(LwDetrConvNormLayer(config, channel_size, channel_size // 2, 1, 1, activation="relu"))
+                layers.append(nn.ConvTranspose2d(channel_size // 2, channel_size // 4, kernel_size=2, stride=2))
+            else:
+                layers.append(nn.ConvTranspose2d(channel_size, channel_size // 2, 2, 2))
+        elif scale == 0.5:
+            layers.append(LwDetrConvNormLayer(config, channel_size, channel_size, 3, 2, activation="relu"))
+        self.layers = nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+        return hidden_states
+
+
+class LwDetrScaleProjector(nn.Module):
+    def __init__(self, config: LwDetrConfig, scale: float):
+        super().__init__()
+
+        intermediate_dims = [config.backbone_config.hidden_size] * len(config.backbone_config.out_indices)
+        sampling_layers = []
+        for channel_size in intermediate_dims:
+            sampling_layers.append(LwDetrSamplingLayer(config, channel_size, scale))
+        self.sampling_layers = nn.ModuleList(sampling_layers)
+
+        intermediate_dim = intermediate_dims[-1]
+        if scale == 2.0:
+            if intermediate_dim > 512:
+                intermediate_dim = intermediate_dim // 4
+            else:
+                intermediate_dim = intermediate_dim // 2
+        projector_input_dim = intermediate_dim * len(intermediate_dims)
+
+        self.projector_layer = LwDetrC2FLayer(config, projector_input_dim)
+        self.layer_norm = LwDetrLayerNorm(config.d_model, data_format="channels_first")
+
+    def forward(self, hidden_states_tuple: tuple[torch.Tensor]) -> torch.Tensor:
+        sampled_hidden_states = []
+        for sampling_layer, hidden_states in zip(self.sampling_layers, hidden_states_tuple):
+            hidden_states = sampling_layer(hidden_states)
+            sampled_hidden_states.append(hidden_states)
+        hidden_states = torch.cat(sampled_hidden_states, dim=1)
+        hidden_states = self.projector_layer(hidden_states)
+        hidden_states = self.layer_norm(hidden_states)
+        return hidden_states
+
+
+class LwDetrMultiScaleProjector(nn.Module):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__()
+
+        self.config = config
+        scale_factors = config.projector_scale_factors
+
+        self.scale_layers = nn.ModuleList([LwDetrScaleProjector(config, scale) for scale in scale_factors])
+
+    def forward(self, hidden_states: tuple[torch.Tensor]) -> list[torch.Tensor]:
+        output_hidden_states = []
+        for scale_layer in self.scale_layers:
+            output_hidden_states.append(scale_layer(hidden_states))
+        return output_hidden_states
+
+
+class LwDetrConvEncoder(nn.Module):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__()
+        self.backbone = LwDetrViTBackbone(config.backbone_config)
+        self.projector = LwDetrMultiScaleProjector(config)
+
+    def forward(self, pixel_values: torch.Tensor, pixel_mask: torch.Tensor):
+        # send pixel_values through the model to get list of feature maps
+        features = self.backbone(pixel_values).feature_maps
+        features = self.projector(features)
+        out = []
+        for feature_map in features:
+            # downsample pixel_mask to match shape of corresponding feature_map
+            mask = nn.functional.interpolate(pixel_mask[None].float(), size=feature_map.shape[-2:]).to(torch.bool)[0]
+            out.append((feature_map, mask))
+        return out
+
+
+class LwDetrAttention(nn.Module):
+    """LW-DETR self-attention with group-DETR training technique."""
+
+    def __init__(self, config: LwDetrConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.head_dim = getattr(config, "head_dim", config.d_model // config.decoder_self_attention_heads)
+        self.scaling = self.head_dim**-0.5
+        self.attention_dropout = config.attention_dropout
+        self.is_causal = False
+        self.num_key_value_groups = 1
+
+        self.q_proj = nn.Linear(
+            config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.k_proj = nn.Linear(
+            config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.v_proj = nn.Linear(
+            config.d_model, config.decoder_self_attention_heads * self.head_dim, bias=config.attention_bias
+        )
+        self.o_proj = nn.Linear(
+            config.decoder_self_attention_heads * self.head_dim, config.d_model, bias=config.attention_bias
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        batch_size, seq_len, _ = hidden_states.shape
+
+        hidden_states_original = hidden_states
+        if position_embeddings is not None:
+            hidden_states = hidden_states if position_embeddings is None else hidden_states + position_embeddings
+
+        if self.training:
+            # at training, we use group detr technique to add more supervision by using multiple weight-sharing decoders at once for faster convergence
+            # at inference, we only use one decoder
+            hidden_states_original = torch.cat(
+                hidden_states_original.split(seq_len // self.config.group_detr, dim=1), dim=0
+            )
+            hidden_states = torch.cat(hidden_states.split(seq_len // self.config.group_detr, dim=1), dim=0)
+
+        attention_input_shape = hidden_states.shape[:-1]
+        hidden_shape = (*attention_input_shape, -1, self.head_dim)
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_states = self.v_proj(hidden_states_original).view(hidden_shape).transpose(1, 2)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask=None,
+            dropout=0.0 if not self.training else self.attention_dropout,
+            scaling=self.scaling,
+            **kwargs,
+        )
+        attn_output = attn_output.reshape(*attention_input_shape, -1).contiguous()
+        attn_output = self.o_proj(attn_output)
+
+        if self.training:
+            attn_output = torch.cat(torch.split(attn_output, batch_size, dim=0), dim=1)
+
+        return attn_output, attn_weights
+
+
+class LwDetrMultiscaleDeformableAttention(DeformableDetrMultiscaleDeformableAttention):
+    pass
+
+
+class LwDetrMLP(nn.Module):
+    def __init__(self, config: LwDetrConfig):
+        super().__init__()
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.decoder_activation_function]
+        self.fc1 = nn.Linear(config.d_model, config.decoder_ffn_dim)
+        self.fc2 = nn.Linear(config.decoder_ffn_dim, config.d_model)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        residual = hidden_states
+        hidden_states = self.fc1(hidden_states)
+        hidden_states = self.activation_fn(hidden_states)
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = self.fc2(hidden_states)
+        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
+        hidden_states = residual + hidden_states
+        return hidden_states
+
+
+class LwDetrDecoderLayer(GradientCheckpointingLayer):
+    def __init__(self, config: LwDetrConfig, layer_idx: int):
+        nn.Module.__init__(self)
+
+        # self-attention
+        self.self_attn = LwDetrAttention(config, layer_idx=layer_idx)
+        self.dropout = config.dropout
+        self.activation_fn = ACT2FN[config.decoder_activation_function]
+        self.activation_dropout = config.activation_dropout
+        self.self_attn_layer_norm = nn.LayerNorm(config.d_model)
+
+        # cross-attention
+        self.cross_attn = LwDetrMultiscaleDeformableAttention(
+            config,
+            num_heads=config.decoder_cross_attention_heads,
+            n_points=config.decoder_n_points,
+        )
+        self.cross_attn_layer_norm = nn.LayerNorm(config.d_model)
+
+        # mlp
+        self.mlp = LwDetrMLP(config)
+        self.layer_norm = nn.LayerNorm(config.d_model)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        position_embeddings: torch.Tensor | None = None,
+        reference_points=None,
+        spatial_shapes=None,
+        spatial_shapes_list=None,
+        level_start_index=None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        self_attention_output, self_attn_weights = self.self_attn(
+            hidden_states, position_embeddings=position_embeddings, **kwargs
+        )
+
+        self_attention_output = nn.functional.dropout(self_attention_output, p=self.dropout, training=self.training)
+        hidden_states = hidden_states + self_attention_output
+        hidden_states = self.self_attn_layer_norm(hidden_states)
+
+        cross_attention_output, cross_attn_weights = self.cross_attn(
+            hidden_states=hidden_states,
+            attention_mask=encoder_attention_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            position_embeddings=position_embeddings,
+            reference_points=reference_points,
+            spatial_shapes=spatial_shapes,
+            spatial_shapes_list=spatial_shapes_list,
+            level_start_index=level_start_index,
+            **kwargs,
+        )
+        cross_attention_output = nn.functional.dropout(cross_attention_output, p=self.dropout, training=self.training)
+        hidden_states = hidden_states + cross_attention_output
+        hidden_states = self.cross_attn_layer_norm(hidden_states)
+
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = self.layer_norm(hidden_states)
+
+        return hidden_states
+
+
+@auto_docstring
+class LwDetrPreTrainedModel(PreTrainedModel):
+    config: LwDetrConfig
+    base_model_prefix = "model"
+    main_input_name = "pixel_values"
+    input_modalities = ("image",)
+    _no_split_modules = [
+        r"LwDetrConvEncoder",
+        r"LwDetrDecoderLayer",
+    ]
+    _supports_sdpa = True
+    _supports_flash_attn = True
+    _supports_flex_attn = True
+    _supports_attention_backend = True
+    _can_record_outputs = {
+        "attentions": [LwDetrAttention, LwDetrMultiscaleDeformableAttention],
+        "hidden_states": [LwDetrDecoderLayer],
+    }
+
+    @torch.no_grad()
+    def _init_weights(self, module):
+        super()._init_weights(module)
+
+        if isinstance(module, LwDetrMultiscaleDeformableAttention):
+            init.constant_(module.sampling_offsets.weight, 0.0)
+            thetas = torch.arange(module.n_heads, dtype=torch.int64).float() * (2.0 * math.pi / module.n_heads)
+            grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
+            grid_init = (
+                (grid_init / grid_init.abs().max(-1, keepdim=True)[0])
+                .view(module.n_heads, 1, 1, 2)
+                .repeat(1, module.n_levels, module.n_points, 1)
+            )
+            for i in range(module.n_points):
+                grid_init[:, :, i, :] *= i + 1
+
+            init.copy_(module.sampling_offsets.bias, grid_init.view(-1))
+            init.constant_(module.attention_weights.weight, 0.0)
+            init.constant_(module.attention_weights.bias, 0.0)
+            init.xavier_uniform_(module.value_proj.weight)
+            init.constant_(module.value_proj.bias, 0.0)
+            init.xavier_uniform_(module.output_proj.weight)
+            init.constant_(module.output_proj.bias, 0.0)
+        if hasattr(module, "level_embed"):
+            init.normal_(module.level_embed)
+        if hasattr(module, "refpoint_embed") and module.refpoint_embed is not None:
+            init.constant_(module.refpoint_embed.weight, 0)
+        if hasattr(module, "class_embed") and module.class_embed is not None:
+            prior_prob = 0.01
+            bias_value = -math.log((1 - prior_prob) / prior_prob)
+            init.constant_(module.class_embed.bias, bias_value)
+        if hasattr(module, "bbox_embed") and module.bbox_embed is not None:
+            init.constant_(module.bbox_embed.layers[-1].weight, 0)
+            init.constant_(module.bbox_embed.layers[-1].bias, 0)
+
+
+def refine_bboxes(reference_points, deltas):
+    reference_points = reference_points.to(deltas.device)
+    new_reference_points_cxcy = deltas[..., :2] * reference_points[..., 2:] + reference_points[..., :2]
+    new_reference_points_wh = deltas[..., 2:].exp() * reference_points[..., 2:]
+    new_reference_points = torch.cat((new_reference_points_cxcy, new_reference_points_wh), -1)
+    return new_reference_points
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for outputs of the LwDetrDecoder. This class adds two attributes to
+    BaseModelOutputWithCrossAttentions, namely:
+    - a stacked tensor of intermediate decoder hidden states (i.e. the output of each decoder layer)
+    - a stacked tensor of intermediate reference points.
+    """
+)
+@dataclass
+class LwDetrDecoderOutput(DeformableDetrDecoderOutput):
+    pass
+
+
+class LwDetrDecoder(LwDetrPreTrainedModel):
+    """
+    Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`DeformableDetrDecoderLayer`].
+
+    The decoder updates the query embeddings through multiple self-attention and deformable cross-attention layers.
+
+    Some tweaks for LwDetr:
+
+    - it uses group detr technique at training for faster convergence.
+
+    Args:
+        config: LwDetrConfig
+    """
+
+    _can_record_outputs = {
+        "hidden_states": LwDetrDecoderLayer,
+        "attentions": OutputRecorder(LwDetrAttention, layer_name="self_attn", index=1),
+        "cross_attentions": OutputRecorder(LwDetrMultiscaleDeformableAttention, layer_name="cross_attn", index=1),
+    }
+
+    def __init__(self, config: LwDetrConfig):
+        super().__init__(config)
+        self.dropout = config.dropout
+        self.layers = nn.ModuleList([LwDetrDecoderLayer(config, i) for i in range(config.decoder_layers)])
+        self.layernorm = nn.LayerNorm(config.d_model)
+
+        self.gradient_checkpointing = False
+
+        self.ref_point_head = LwDetrMLPPredictionHead(2 * config.d_model, config.d_model, config.d_model, num_layers=2)
+
+        self.post_init()
+
+    def get_reference(self, reference_points, valid_ratios):
+        # batch_size, num_queries, batch_size, 4
+        obj_center = reference_points[..., :4]
+
+        # batch_size, num_queries, num_levels, 4
+        reference_points_inputs = obj_center[:, :, None] * torch.cat([valid_ratios, valid_ratios], -1)[:, None]
+
+        # batch_size, num_queries, d_model * 2
+        query_sine_embed = encode_sinusoidal_position_embedding(
+            reference_points_inputs[:, :, 0, :], num_pos_feats=self.config.d_model // 2
+        )
+
+        # batch_size, num_queries, d_model
+        query_pos = self.ref_point_head(query_sine_embed)
+        return reference_points_inputs, query_pos
+
+    @merge_with_config_defaults
+    @capture_outputs
+    def forward(
+        self,
+        inputs_embeds: torch.Tensor | None = None,
+        reference_points: torch.Tensor | None = None,
+        spatial_shapes: torch.Tensor | None = None,
+        spatial_shapes_list: torch.Tensor | None = None,
+        level_start_index: torch.Tensor | None = None,
+        valid_ratios: torch.Tensor | None = None,
+        encoder_hidden_states: torch.Tensor | None = None,
+        encoder_attention_mask: torch.Tensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ):
+        intermediate = ()
+        intermediate_reference_points = (reference_points,)
+
+        if inputs_embeds is not None:
+            hidden_states = inputs_embeds
+
+        reference_points_inputs, query_pos = self.get_reference(reference_points, valid_ratios)
+
+        for idx, decoder_layer in enumerate(self.layers):
+            hidden_states = decoder_layer(
+                hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
+                position_embeddings=query_pos,
+                reference_points=reference_points_inputs,
+                spatial_shapes=spatial_shapes,
+                spatial_shapes_list=spatial_shapes_list,
+                level_start_index=level_start_index,
+                **kwargs,
+            )
+            intermediate_hidden_states = self.layernorm(hidden_states)
+            intermediate += (intermediate_hidden_states,)
+
+        intermediate = torch.stack(intermediate)
+        last_hidden_state = intermediate[-1]
+        intermediate_reference_points = torch.stack(intermediate_reference_points)
+
+        return LwDetrDecoderOutput(
+            last_hidden_state=last_hidden_state,
+            intermediate_hidden_states=intermediate,
+            intermediate_reference_points=intermediate_reference_points,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for outputs of the LwDetr backbone-decoder model.
+    """
+)
+@dataclass
+class LwDetrModelOutput(ModelOutput):
+    r"""
+    init_reference_points (`torch.FloatTensor` of shape  `(batch_size, num_queries, 4)`):
+        Initial reference points sent through the Transformer decoder.
+    intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
+        Stacked intermediate hidden states (output of each layer of the decoder).
+    intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+        Stacked intermediate reference points (reference points of each layer of the decoder).
+    enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+        Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
+        picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
+        foreground and background).
+    enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+        Logits of predicted bounding boxes coordinates in the first stage.
+    """
+
+    init_reference_points: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    intermediate_hidden_states: torch.FloatTensor | None = None
+    intermediate_reference_points: torch.FloatTensor | None = None
+    enc_outputs_class: torch.FloatTensor | None = None
+    enc_outputs_coord_logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    cross_attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    The bare LW Detr Model (consisting of a backbone and decoder Transformer) outputting raw
+    hidden-states without any specific head on top.
+    """
+)
+class LwDetrModel(DeformableDetrModel):
+    def __init__(self, config: LwDetrConfig):
+        PreTrainedModel.__init__(self, config)
+
+        # Create backbone + positional encoding
+        self.backbone = LwDetrConvEncoder(config)
+
+        self.group_detr = config.group_detr
+        self.num_queries = config.num_queries
+        hidden_dim = config.d_model
+        self.reference_point_embed = nn.Embedding(self.num_queries * self.group_detr, 4)
+        self.query_feat = nn.Embedding(self.num_queries * self.group_detr, hidden_dim)
+
+        self.decoder = LwDetrDecoder(config)
+
+        self.enc_output = nn.ModuleList([nn.Linear(hidden_dim, hidden_dim) for _ in range(self.group_detr)])
+        self.enc_output_norm = nn.ModuleList([nn.LayerNorm(hidden_dim) for _ in range(self.group_detr)])
+        # Should normally be None and then instantiated in the ForObjectDetection class
+        self.enc_out_bbox_embed = nn.ModuleList(
+            [LwDetrMLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3) for _ in range(self.group_detr)]
+        )
+        self.enc_out_class_embed = nn.ModuleList(
+            [nn.Linear(config.d_model, config.num_labels) for _ in range(self.group_detr)]
+        )
+
+        self.post_init()
+
+    def gen_encoder_output_proposals(self, enc_output, padding_mask, spatial_shapes):
+        """Generate the encoder output proposals from encoded enc_output.
+
+        Args:
+            enc_output (Tensor[batch_size, sequence_length, hidden_size]): Output of the encoder.
+            padding_mask (Tensor[batch_size, sequence_length]): Padding mask for `enc_output`.
+            spatial_shapes (list[tuple[int, int]]): Spatial shapes of the feature maps.
+
+        Returns:
+            `tuple(torch.FloatTensor)`: A tuple of feature map and bbox prediction.
+                - object_query (Tensor[batch_size, sequence_length, hidden_size]): Object query features. Later used to
+                  directly predict a bounding box. (without the need of a decoder)
+                - output_proposals (Tensor[batch_size, sequence_length, 4]): Normalized proposals in [0, 1] space.
+                  Invalid positions (padding or out-of-bounds) are filled with 0.
+                - invalid_mask (Tensor[batch_size, sequence_length, 1]): Boolean mask that is True for invalid positions
+                  (padded pixels or proposals whose coordinates fall outside (0.01, 0.99)).
+        """
+        batch_size = enc_output.shape[0]
+        proposals = []
+        _cur = 0
+        for level, (height, width) in enumerate(spatial_shapes):
+            mask_flatten_ = padding_mask[:, _cur : (_cur + height * width)].view(batch_size, height, width, 1)
+            valid_height = torch.sum(~mask_flatten_[:, :, 0, 0], 1)
+            valid_width = torch.sum(~mask_flatten_[:, 0, :, 0], 1)
+
+            grid_y, grid_x = torch.meshgrid(
+                torch.linspace(
+                    0,
+                    height - 1,
+                    height,
+                    dtype=enc_output.dtype,
+                    device=enc_output.device,
+                ),
+                torch.linspace(
+                    0,
+                    width - 1,
+                    width,
+                    dtype=enc_output.dtype,
+                    device=enc_output.device,
+                ),
+                indexing="ij",
+            )
+            grid = torch.cat([grid_x.unsqueeze(-1), grid_y.unsqueeze(-1)], -1)
+
+            scale = torch.cat([valid_width.unsqueeze(-1), valid_height.unsqueeze(-1)], 1).view(batch_size, 1, 1, 2)
+            grid = (grid.unsqueeze(0).expand(batch_size, -1, -1, -1) + 0.5) / scale
+            width_height = torch.ones_like(grid) * 0.05 * (2.0**level)
+            proposal = torch.cat((grid, width_height), -1).view(batch_size, -1, 4)
+            proposals.append(proposal)
+            _cur += height * width
+        output_proposals = torch.cat(proposals, 1)
+        output_proposals_valid = ((output_proposals > 0.01) & (output_proposals < 0.99)).all(-1, keepdim=True)
+        invalid_mask = padding_mask.unsqueeze(-1) | ~output_proposals_valid
+        output_proposals = output_proposals.masked_fill(invalid_mask, float(0))
+
+        # assign each pixel as an object query
+        object_query = enc_output
+        object_query = object_query.masked_fill(invalid_mask, float(0))
+        return object_query, output_proposals, invalid_mask
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor = None,
+        pixel_mask: torch.LongTensor | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> LwDetrModelOutput:
+        r"""
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, DeformableDetrModel
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> image_processor = AutoImageProcessor.from_pretrained("AnnaZhang/lwdetr_small_60e_coco")
+        >>> model = DeformableDetrModel.from_pretrained("AnnaZhang/lwdetr_small_60e_coco")
+
+        >>> inputs = image_processor(images=image, return_tensors="pt")
+
+        >>> outputs = model(**inputs)
+
+        >>> last_hidden_states = outputs.last_hidden_state
+        >>> list(last_hidden_states.shape)
+        [1, 300, 256]
+        ```"""
+        batch_size, num_channels, height, width = pixel_values.shape
+        device = pixel_values.device
+
+        if pixel_mask is None:
+            pixel_mask = torch.ones(((batch_size, height, width)), dtype=torch.long, device=device)
+
+        # Extract multi-scale feature maps of same resolution `config.d_model` (cf Figure 4 in paper)
+        # First, sent pixel_values + pixel_mask through Backbone to obtain the features
+        # which is a list of tuples
+        features = self.backbone(pixel_values, pixel_mask)
+
+        # Then, apply 1x1 convolution to reduce the channel dimension to d_model (256 by default)
+        sources = []
+        masks = []
+        for level, (source, mask) in enumerate(features):
+            sources.append(source)
+            masks.append(mask)
+            if mask is None:
+                raise ValueError("No attention mask was provided")
+
+        if self.training:
+            reference_points = self.reference_point_embed.weight
+            query_feat = self.query_feat.weight
+        else:
+            # only use one group in inference
+            reference_points = self.reference_point_embed.weight[: self.num_queries]
+            query_feat = self.query_feat.weight[: self.num_queries]
+
+        # Prepare encoder inputs (by flattening)
+        source_flatten = []
+        mask_flatten = []
+        spatial_shapes_list = []
+        for source, mask in zip(sources, masks):
+            batch_size, num_channels, height, width = source.shape
+            spatial_shape = (height, width)
+            spatial_shapes_list.append(spatial_shape)
+            source = source.flatten(2).transpose(1, 2)
+            mask = mask.flatten(1)
+            source_flatten.append(source)
+            mask_flatten.append(mask)
+        source_flatten = torch.cat(source_flatten, 1)
+        mask_flatten = torch.cat(mask_flatten, 1)
+        spatial_shapes = torch.as_tensor(spatial_shapes_list, dtype=torch.long, device=source_flatten.device)
+        level_start_index = torch.cat((spatial_shapes.new_zeros((1,)), spatial_shapes.prod(1).cumsum(0)[:-1]))
+        valid_ratios = torch.stack([self.get_valid_ratio(m, dtype=source_flatten.dtype) for m in masks], 1)
+
+        target = query_feat.unsqueeze(0).expand(batch_size, -1, -1)
+        reference_points = reference_points.unsqueeze(0).expand(batch_size, -1, -1)
+
+        object_query_embedding, output_proposals, invalid_mask = self.gen_encoder_output_proposals(
+            source_flatten, ~mask_flatten, spatial_shapes_list
+        )
+
+        group_detr = self.group_detr if self.training else 1
+        topk = self.num_queries
+        topk_coords_logits = []
+        topk_coords_logits_undetach = []
+        object_query_undetach = []
+
+        for group_id in range(group_detr):
+            group_object_query = self.enc_output[group_id](object_query_embedding)
+            group_object_query = self.enc_output_norm[group_id](group_object_query)
+
+            group_enc_outputs_class = self.enc_out_class_embed[group_id](group_object_query)
+            group_enc_outputs_class = group_enc_outputs_class.masked_fill(invalid_mask, float("-inf"))
+            group_delta_bbox = self.enc_out_bbox_embed[group_id](group_object_query)
+            group_enc_outputs_coord = refine_bboxes(output_proposals, group_delta_bbox)
+
+            group_topk_proposals = torch.topk(group_enc_outputs_class.max(-1)[0], topk, dim=1)[1]
+            group_topk_coords_logits_undetach = torch.gather(
+                group_enc_outputs_coord,
+                1,
+                group_topk_proposals.unsqueeze(-1).repeat(1, 1, 4),
+            )
+            group_topk_coords_logits = group_topk_coords_logits_undetach.detach()
+            group_object_query_undetach = torch.gather(
+                group_object_query, 1, group_topk_proposals.unsqueeze(-1).repeat(1, 1, self.config.d_model)
+            )
+
+            topk_coords_logits.append(group_topk_coords_logits)
+            topk_coords_logits_undetach.append(group_topk_coords_logits_undetach)
+            object_query_undetach.append(group_object_query_undetach)
+
+        topk_coords_logits = torch.cat(topk_coords_logits, 1)
+        topk_coords_logits_undetach = torch.cat(topk_coords_logits_undetach, 1)
+        object_query_undetach = torch.cat(object_query_undetach, 1)
+
+        enc_outputs_class = object_query_undetach
+        enc_outputs_coord_logits = topk_coords_logits_undetach
+
+        reference_points = refine_bboxes(topk_coords_logits, reference_points)
+
+        init_reference_points = reference_points
+        decoder_outputs = self.decoder(
+            inputs_embeds=target,
+            reference_points=reference_points,
+            spatial_shapes=spatial_shapes,
+            spatial_shapes_list=spatial_shapes_list,
+            level_start_index=level_start_index,
+            valid_ratios=valid_ratios,
+            encoder_hidden_states=source_flatten,
+            encoder_attention_mask=mask_flatten,
+            **kwargs,
+        )
+
+        return LwDetrModelOutput(
+            init_reference_points=init_reference_points,
+            last_hidden_state=decoder_outputs.last_hidden_state,
+            intermediate_hidden_states=decoder_outputs.intermediate_hidden_states,
+            intermediate_reference_points=decoder_outputs.intermediate_reference_points,
+            enc_outputs_class=enc_outputs_class,
+            enc_outputs_coord_logits=enc_outputs_coord_logits,
+            hidden_states=decoder_outputs.hidden_states,
+            attentions=decoder_outputs.attentions,
+            cross_attentions=decoder_outputs.cross_attentions,
+        )
+
+    def get_proposal_pos_embed(self, proposals):
+        raise NotImplementedError("get_proposal_pos_embed is not used in LwDetrForObjectDetection")
+
+
+class LwDetrMLPPredictionHead(DeformableDetrMLPPredictionHead):
+    pass
+
+
+@auto_docstring(
+    custom_intro="""
+    Output type of [`LwDetrForObjectDetection`].
+    """
+)
+@dataclass
+class LwDetrObjectDetectionOutput(ModelOutput):
+    r"""
+    loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` are provided)):
+        Total loss as a linear combination of a negative log-likehood (cross-entropy) for class prediction and a
+        bounding box loss. The latter is defined as a linear combination of the L1 loss and the generalized
+        scale-invariant IoU loss.
+    loss_dict (`Dict`, *optional*):
+        A dictionary containing the individual losses. Useful for logging.
+    logits (`torch.FloatTensor` of shape `(batch_size, num_queries, num_classes + 1)`):
+        Classification logits (including no-object) for all queries.
+    pred_boxes (`torch.FloatTensor` of shape `(batch_size, num_queries, 4)`):
+        Normalized boxes coordinates for all queries, represented as (center_x, center_y, width, height). These
+        values are normalized in [0, 1], relative to the size of each individual image in the batch (disregarding
+        possible padding). You can use [`~DeformableDetrProcessor.post_process_object_detection`] to retrieve the
+        unnormalized bounding boxes.
+    auxiliary_outputs (`list[Dict]`, *optional*):
+        Optional, only returned when auxiliary losses are activated (i.e. `config.auxiliary_loss` is set to `True`)
+        and labels are provided. It is a list of dictionaries containing the two above keys (`logits` and
+        `pred_boxes`) for each decoder layer.
+    init_reference_points (`torch.FloatTensor` of shape  `(batch_size, num_queries, 4)`):
+        Initial reference points sent through the Transformer decoder.
+    intermediate_hidden_states (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, hidden_size)`):
+        Stacked intermediate hidden states (output of each layer of the decoder).
+    intermediate_reference_points (`torch.FloatTensor` of shape `(batch_size, config.decoder_layers, num_queries, 4)`):
+        Stacked intermediate reference points (reference points of each layer of the decoder).
+    enc_outputs_class (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.num_labels)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+        Predicted bounding boxes scores where the top `config.two_stage_num_proposals` scoring bounding boxes are
+        picked as region proposals in the first stage. Output of bounding box binary classification (i.e.
+        foreground and background).
+    enc_outputs_coord_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, 4)`, *optional*, returned when `config.with_box_refine=True` and `config.two_stage=True`):
+        Logits of predicted bounding boxes coordinates in the first stage.
+    """
+
+    loss: torch.FloatTensor | None = None
+    loss_dict: dict | None = None
+    logits: torch.FloatTensor | None = None
+    pred_boxes: torch.FloatTensor | None = None
+    auxiliary_outputs: list[dict] | None = None
+    init_reference_points: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    intermediate_hidden_states: torch.FloatTensor | None = None
+    intermediate_reference_points: torch.FloatTensor | None = None
+    enc_outputs_class: Any = None
+    enc_outputs_coord_logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    cross_attentions: tuple[torch.FloatTensor, ...] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    LW DETR Model (consisting of a backbone and decoder Transformer) with object detection heads on
+    top, for tasks such as COCO detection.
+    """
+)
+class LwDetrForObjectDetection(DeformableDetrForObjectDetection):
+    _tied_weights_keys = None
+
+    def __init__(self, config: LwDetrConfig):
+        PreTrainedModel.__init__(self, config)
+        self.model = LwDetrModel(config)
+        self.class_embed = nn.Linear(config.d_model, config.num_labels)
+        self.bbox_embed = LwDetrMLPPredictionHead(config.d_model, config.d_model, 4, num_layers=3)
+
+        self.post_init()
+
+    @can_return_tuple
+    @auto_docstring
+    def forward(
+        self,
+        pixel_values: torch.FloatTensor = None,
+        pixel_mask: torch.LongTensor | None = None,
+        labels: list[dict] | None = None,
+        **kwargs: Unpack[TransformersKwargs],
+    ) -> LwDetrObjectDetectionOutput:
+        r"""
+        labels (`list[Dict]` of len `(batch_size,)`, *optional*):
+            Labels for computing the bipartite matching loss. List of dicts, each dictionary containing at least the
+            following 2 keys: 'class_labels' and 'boxes' (the class labels and bounding boxes of an image in the batch
+            respectively). The class labels themselves should be a `torch.LongTensor` of len `(number of bounding boxes
+            in the image,)` and the boxes a `torch.FloatTensor` of shape `(number of bounding boxes in the image, 4)`.
+
+        Examples:
+
+        ```python
+        >>> from transformers import AutoImageProcessor, LwDetrForObjectDetection
+        >>> from PIL import Image
+        >>> import httpx
+        >>> from io import BytesIO
+
+        >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg"
+        >>> with httpx.stream("GET", url) as response:
+        ...     image = Image.open(BytesIO(response.read()))
+
+        >>> image_processor = AutoImageProcessor.from_pretrained("AnnaZhang/lwdetr_small_60e_coco")
+        >>> model = LwDetrForObjectDetection.from_pretrained("AnnaZhang/lwdetr_small_60e_coco")
+
+        >>> inputs = image_processor(images=image, return_tensors="pt")
+        >>> outputs = model(**inputs)
+
+        >>> # convert outputs (bounding boxes and class logits) to Pascal VOC format (xmin, ymin, xmax, ymax)
+        >>> target_sizes = torch.tensor([image.size[::-1]])
+        >>> results = image_processor.post_process_object_detection(outputs, threshold=0.5, target_sizes=target_sizes)[
+        ...     0
+        ... ]
+        >>> for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
+        ...     box = [round(i, 2) for i in box.tolist()]
+        ...     print(
+        ...         f"Detected {model.config.id2label[label.item()]} with confidence "
+        ...         f"{round(score.item(), 3)} at location {box}"
+        ...     )
+        Detected cat with confidence 0.8 at location [16.5, 52.84, 318.25, 470.78]
+        Detected cat with confidence 0.789 at location [342.19, 24.3, 640.02, 372.25]
+        Detected remote with confidence 0.633 at location [40.79, 72.78, 176.76, 117.25]
+        ```"""
+        outputs = self.model(
+            pixel_values,
+            pixel_mask=pixel_mask,
+            **kwargs,
+        )
+
+        last_hidden_states = outputs.last_hidden_state
+        intermediate_reference_points = outputs.intermediate_reference_points
+        enc_outputs_class_logits = outputs.enc_outputs_class
+        enc_outputs_boxes_logits = outputs.enc_outputs_coord_logits
+
+        logits = self.class_embed(last_hidden_states)
+        pred_boxes_delta = self.bbox_embed(last_hidden_states)
+        pred_boxes = refine_bboxes(intermediate_reference_points[-1], pred_boxes_delta)
+
+        enc_outputs_class_logits_list = enc_outputs_class_logits.split(self.config.num_queries, dim=1)
+        pred_class = []
+        group_detr = self.config.group_detr if self.training else 1
+        for group_index in range(group_detr):
+            group_pred_class = self.model.enc_out_class_embed[group_index](enc_outputs_class_logits_list[group_index])
+            pred_class.append(group_pred_class)
+        enc_outputs_class_logits = torch.cat(pred_class, dim=1)
+
+        loss, loss_dict, auxiliary_outputs = None, None, None
+        if labels is not None:
+            outputs_class, outputs_coord = None, None
+            if self.config.auxiliary_loss:
+                intermediate_hidden_states = outputs.intermediate_hidden_states
+                outputs_coord_delta = self.bbox_embed(intermediate_hidden_states)
+                outputs_coord = refine_bboxes(intermediate_reference_points, outputs_coord_delta)
+                outputs_class = self.class_embed(intermediate_hidden_states)
+
+            loss, loss_dict, auxiliary_outputs = self.loss_function(
+                logits,
+                labels,
+                self.device,
+                pred_boxes,
+                self.config,
+                outputs_class,
+                outputs_coord,
+                enc_outputs_class_logits,
+                enc_outputs_boxes_logits,
+            )
+
+        return LwDetrObjectDetectionOutput(
+            loss=loss,
+            loss_dict=loss_dict,
+            logits=logits,
+            pred_boxes=pred_boxes,
+            auxiliary_outputs=auxiliary_outputs,
+            last_hidden_state=outputs.last_hidden_state,
+            intermediate_hidden_states=outputs.intermediate_hidden_states,
+            intermediate_reference_points=outputs.intermediate_reference_points,
+            init_reference_points=outputs.init_reference_points,
+            enc_outputs_class=enc_outputs_class_logits,
+            enc_outputs_coord_logits=enc_outputs_boxes_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            cross_attentions=outputs.cross_attentions,
+        )
+
+
+__all__ = [
+    "LwDetrConfig",
+    "LwDetrPreTrainedModel",
+    "LwDetrModel",
+    "LwDetrForObjectDetection",
+    "LwDetrViTConfig",
+    "LwDetrViTPreTrainedModel",
+    "LwDetrViTBackbone",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/modernbert/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_modernbert import *
+    from .modeling_modernbert import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/patchtst/__init__.py

@@ -0,0 +1,27 @@
+# Copyright 2024 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import TYPE_CHECKING
+
+from ...utils import _LazyModule
+from ...utils.import_utils import define_import_structure
+
+
+if TYPE_CHECKING:
+    from .configuration_patchtst import *
+    from .modeling_patchtst import *
+else:
+    import sys
+
+    _file = globals()["__file__"]
+    sys.modules[__name__] = _LazyModule(__name__, _file, define_import_structure(_file), module_spec=__spec__)

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/patchtst/configuration_patchtst.py

@@ -0,0 +1,170 @@
+# Copyright 2023 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PatchTST model configuration"""
+
+from huggingface_hub.dataclasses import strict
+
+from transformers.configuration_utils import PreTrainedConfig
+from transformers.utils import auto_docstring
+
+
+@auto_docstring(checkpoint="ibm-granite/granite-timeseries-patchtst")
+@strict
+class PatchTSTConfig(PreTrainedConfig):
+    r"""
+    context_length (`int`, *optional*, defaults to 32):
+        The context length of the input sequence.
+    distribution_output (`str`, *optional*, defaults to `"student_t"`):
+        The distribution emission head for the model when loss is "nll". Could be either "student_t", "normal" or
+        "negative_binomial".
+    loss (`str`, *optional*, defaults to `"mse"`):
+        The loss function for the model corresponding to the `distribution_output` head. For parametric
+        distributions it is the negative log likelihood ("nll") and for point estimates it is the mean squared
+        error "mse".
+    patch_length (`int`, *optional*, defaults to 1):
+        Define the patch length of the patchification process.
+    patch_stride (`int`, *optional*, defaults to 1):
+        Define the stride of the patchification process.
+    num_attention_heads (`int`, *optional*, defaults to 4):
+        Number of attention heads for each attention layer in the Transformer encoder.
+    share_embedding (`bool`, *optional*, defaults to `True`):
+        Sharing the input embedding across all channels.
+    channel_attention (`bool`, *optional*, defaults to `False`):
+        Activate channel attention block in the Transformer to allow channels to attend each other.
+    ffn_dim (`int`, *optional*, defaults to 512):
+        Dimension of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
+    norm_type (`str` , *optional*, defaults to `"batchnorm"`):
+        Normalization at each Transformer layer. Can be `"batchnorm"` or `"layernorm"`.
+    norm_eps (`float`, *optional*, defaults to 1e-05):
+        A value added to the denominator for numerical stability of normalization.
+    positional_dropout (`float`, *optional*, defaults to 0.0):
+        The dropout probability in the positional embedding layer.
+    path_dropout (`float`, *optional*, defaults to 0.0):
+        The dropout path in the residual block.
+    ff_dropout (`float`, *optional*, defaults to 0.0):
+        The dropout probability used between the two layers of the feed-forward networks.
+    bias (`bool`, *optional*, defaults to `True`):
+        Whether to add bias in the feed-forward networks.
+    activation_function (`str`, *optional*, defaults to `"gelu"`):
+        The non-linear activation function (string) in the Transformer.`"gelu"` and `"relu"` are supported.
+    pre_norm (`bool`, *optional*, defaults to `True`):
+        Normalization is applied before self-attention if pre_norm is set to `True`. Otherwise, normalization is
+        applied after residual block.
+    positional_encoding_type (`str`, *optional*, defaults to `"sincos"`):
+        Positional encodings. Options `"random"` and `"sincos"` are supported.
+    use_cls_token (`bool`, *optional*, defaults to `False`):
+        Whether cls token is used.
+    init_std (`float`, *optional*, defaults to 0.02):
+        The standard deviation of the truncated normal weight initialization distribution.
+    share_projection (`bool`, *optional*, defaults to `True`):
+        Sharing the projection layer across different channels in the forecast head.
+    scaling (`Union`, *optional*, defaults to `"std"`):
+        Whether to scale the input targets via "mean" scaler, "std" scaler or no scaler if `None`. If `True`, the
+        scaler is set to "mean".
+    do_mask_input (`bool`, *optional*):
+        Apply masking during the pretraining.
+    mask_type (`str`, *optional*, defaults to `"random"`):
+        Masking type. Only `"random"` and `"forecast"` are currently supported.
+    random_mask_ratio (`float`, *optional*, defaults to 0.5):
+        Masking ratio applied to mask the input data during random pretraining.
+    num_forecast_mask_patches (`int` or `list`, *optional*, defaults to `[2]`):
+        Number of patches to be masked at the end of each batch sample. If it is an integer,
+        all the samples in the batch will have the same number of masked patches. If it is a list,
+        samples in the batch will be randomly masked by numbers defined in the list. This argument is only used
+        for forecast pretraining.
+    channel_consistent_masking (`bool`, *optional*, defaults to `False`):
+        If channel consistent masking is True, all the channels will have the same masking pattern.
+    unmasked_channel_indices (`list`, *optional*):
+        Indices of channels that are not masked during pretraining. Values in the list are number between 1 and
+        `num_input_channels`
+    mask_value (`int`, *optional*, defaults to 0):
+        Values in the masked patches will be filled by `mask_value`.
+    pooling_type (`str`, *optional*, defaults to `"mean"`):
+        Pooling of the embedding. `"mean"`, `"max"` and `None` are supported.
+    head_dropout (`float`, *optional*, defaults to 0.0):
+        The dropout probability for head.
+    prediction_length (`int`, *optional*, defaults to 24):
+        The prediction horizon that the model will output.
+    num_targets (`int`, *optional*, defaults to 1):
+        Number of targets for regression and classification tasks. For classification, it is the number of
+        classes.
+    output_range (`list`, *optional*):
+        Output range for regression task. The range of output values can be set to enforce the model to produce
+        values within a range.
+    num_parallel_samples (`int`, *optional*, defaults to 100):
+        The number of samples is generated in parallel for probabilistic prediction.
+
+    ```python
+    >>> from transformers import PatchTSTConfig, PatchTSTModel
+
+    >>> # Initializing an PatchTST configuration with 12 time steps for prediction
+    >>> configuration = PatchTSTConfig(prediction_length=12)
+
+    >>> # Randomly initializing a model (with random weights) from the configuration
+    >>> model = PatchTSTModel(configuration)
+
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```
+    """
+
+    model_type = "patchtst"
+    attribute_map = {
+        "hidden_size": "d_model",
+        "num_attention_heads": "num_attention_heads",
+        "num_hidden_layers": "num_hidden_layers",
+    }
+
+    num_input_channels: int = 1
+    context_length: int = 32
+    distribution_output: str = "student_t"
+    loss: str | None = "mse"
+    patch_length: int = 1
+    patch_stride: int = 1
+    num_hidden_layers: int = 3
+    d_model: int = 128
+    num_attention_heads: int = 4
+    share_embedding: bool = True
+    channel_attention: bool = False
+    ffn_dim: int = 512
+    norm_type: str = "batchnorm"
+    norm_eps: float = 1e-05
+    attention_dropout: float | int = 0.0
+    positional_dropout: float | int = 0.0
+    path_dropout: float | int = 0.0
+    ff_dropout: float | int = 0.0
+    bias: bool = True
+    activation_function: str = "gelu"
+    pre_norm: bool = True
+    positional_encoding_type: str = "sincos"
+    use_cls_token: bool = False
+    init_std: float = 0.02
+    share_projection: bool = True
+    scaling: str | bool | None = "std"
+    do_mask_input: bool | None = None
+    mask_type: str = "random"
+    random_mask_ratio: float = 0.5
+    num_forecast_mask_patches: list[int] | tuple[int, ...] | int | None = (2,)
+    channel_consistent_masking: bool | None = False
+    unmasked_channel_indices: list[int] | None = None
+    mask_value: int = 0
+    pooling_type: str | None = "mean"
+    head_dropout: float | int = 0.0
+    prediction_length: int = 24
+    num_targets: int = 1
+    output_range: list | None = None
+    num_parallel_samples: int = 100
+
+
+__all__ = ["PatchTSTConfig"]

LTA_openwebtext_dualt/mini_owt_logdirichlet/.venv_qwen35_uv/lib/python3.12/site-packages/transformers/models/patchtst/modeling_patchtst.py

@@ -0,0 +1,1973 @@
+# Copyright 2023 IBM & Hugging Face. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch PatchTST model."""
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+
+from ... import initialization as init
+from ...activations import ACT2CLS
+from ...integrations.deepspeed import is_deepspeed_zero3_enabled
+from ...modeling_flash_attention_utils import FlashAttentionKwargs
+from ...modeling_outputs import BaseModelOutput
+from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
+from ...processing_utils import Unpack
+from ...time_series_utils import NegativeBinomialOutput, NormalOutput, StudentTOutput
+from ...utils import ModelOutput, TransformersKwargs, auto_docstring, logging
+from .configuration_patchtst import PatchTSTConfig
+
+
+logger = logging.get_logger(__name__)
+
+
+# Copied from transformers.models.bert.modeling_bert.eager_attention_forward
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    scaling: float | None = None,
+    dropout: float = 0.0,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    if scaling is None:
+        scaling = query.size(-1) ** -0.5
+
+    # Take the dot product between "query" and "key" to get the raw attention scores.
+    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling
+
+    if attention_mask is not None:
+        attn_weights = attn_weights + attention_mask
+
+    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
+    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
+
+    attn_output = torch.matmul(attn_weights, value)
+    attn_output = attn_output.transpose(1, 2).contiguous()
+
+    return attn_output, attn_weights
+
+
+# Copied from transformers.models.wav2vec2.modeling_wav2vec2.Wav2Vec2Attention with Wav2Vec2->PatchTST
+class PatchTSTAttention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        is_decoder: bool = False,
+        bias: bool = True,
+        is_causal: bool = False,
+        config: PatchTSTConfig | None = None,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.config = config
+
+        if (self.head_dim * num_heads) != self.embed_dim:
+            raise ValueError(
+                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+                f" and `num_heads`: {num_heads})."
+            )
+        self.scaling = self.head_dim**-0.5
+        self.is_decoder = is_decoder
+        self.is_causal = is_causal
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        key_value_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        output_attentions: bool | None = False,
+        # TODO: we need a refactor so that the different attention modules can get their specific kwargs
+        # ATM, we have mixed things encoder, decoder, and encoder-decoder attn
+        **kwargs: Unpack[FlashAttentionKwargs],
+    ) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
+        """Input shape: Batch x Time x Channel"""
+
+        # if key_value_states are provided this layer is used as a cross-attention layer
+        # for the decoder
+        is_cross_attention = key_value_states is not None
+
+        # determine input shapes
+        input_shape = hidden_states.shape[:-1]
+
+        hidden_shape = (*input_shape, -1, self.head_dim)
+
+        # get query proj
+        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        current_states = key_value_states if is_cross_attention else hidden_states
+        kv_shape = (*current_states.shape[:-1], -1, self.head_dim)
+        key_states = self.k_proj(current_states).view(kv_shape).transpose(1, 2)
+        value_states = self.v_proj(current_states).view(kv_shape).transpose(1, 2)
+
+        attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
+            self.config._attn_implementation, eager_attention_forward
+        )
+
+        attn_output, attn_weights = attention_interface(
+            self,
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            dropout=0.0 if not self.training else self.dropout,
+            scaling=self.scaling,
+            output_attentions=output_attentions,
+            **kwargs,
+        )
+
+        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
+        attn_output = self.out_proj(attn_output)
+
+        return attn_output, attn_weights, None
+
+
+class PatchTSTBatchNorm(nn.Module):
+    """
+    Compute batch normalization over the sequence length (time) dimension.
+    """
+
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        self.batchnorm = nn.BatchNorm1d(config.d_model, eps=config.norm_eps)
+
+    def forward(self, inputs: torch.Tensor):
+        """
+        Parameters:
+            inputs (`torch.Tensor` of shape `(batch_size, sequence_length, d_model)`):
+                input for Batch norm calculation
+        Returns:
+            `torch.Tensor` of shape `(batch_size, sequence_length, d_model)`
+        """
+        output = inputs.transpose(1, 2)  # output: (batch_size, d_model, sequence_length)
+        output = self.batchnorm(output)
+        return output.transpose(1, 2)
+
+
+def random_masking(
+    inputs: torch.Tensor,
+    mask_ratio: float,
+    unmasked_channel_indices: list | None = None,
+    channel_consistent_masking: bool = False,
+    mask_value: int = 0,
+):
+    """random_masking: Mask the input considering the control variables.
+
+    Args:
+        inputs (`torch.Tensor` of shape `(batch_size, num_channels, sequence_length, num_features)`):
+            The input tensor to mask.
+        mask_ratio (`float`):
+            Masking ratio applied to mask the input data during random pretraining. It is the number between 0 and 1.
+        unmasked_channel_indices (list, *optional*):
+            Indices of channels that will not be masked.
+        channel_consistent_masking (bool, *optional*, defaults to `False`):
+            When true, masking will be same across all channels of a timeseries. Otherwise, masking positions will vary
+            across channels.
+        mask_value (int, *optional*, defaults to 0):
+            Define the value of masked patches for pretraining.
+
+    Returns:
+        `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as input Tensor and mask tensor of shape [bs x c x
+        n]
+    """
+    if mask_ratio < 0 or mask_ratio >= 1:
+        raise ValueError(f"Mask ratio {mask_ratio} has to be between 0 and 1.")
+
+    batch_size, num_channels, sequence_length, num_features = inputs.shape
+    device = inputs.device
+
+    len_keep = int(sequence_length * (1 - mask_ratio))
+
+    if channel_consistent_masking:
+        noise = torch.rand(batch_size, 1, sequence_length, device=device)  # noise in [0, 1], bs x 1 x  L
+        noise = noise.repeat(1, num_channels, 1)  # bs x num_channels x time
+    else:
+        # noise in [0, 1], bs x num_channels x L
+        noise = torch.rand(batch_size, num_channels, sequence_length, device=device)
+
+    # mask: [bs x num_channels x num_patch]
+    mask = torch.ones(batch_size, num_channels, sequence_length, device=device)
+    mask[:, :, :len_keep] = 0
+
+    # sort noise for each sample
+    ids_shuffle = torch.argsort(noise, dim=-1)  # ascend: small is keep, large is remove
+    ids_restore = torch.argsort(ids_shuffle, dim=-1)  # ids_restore: [bs x num_channels x L]
+
+    mask = torch.gather(mask, dim=-1, index=ids_restore)
+    mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features)  # mask: [bs x num_channels x num_patches x patch_length]
+    if unmasked_channel_indices is not None:
+        mask[:, unmasked_channel_indices, :, :] = 0
+
+    inputs_mask = inputs.masked_fill(mask.bool(), mask_value)
+    return inputs_mask, mask[..., 0]
+
+
+def forecast_masking(
+    inputs: torch.Tensor,
+    num_forecast_mask_patches: list | int,
+    unmasked_channel_indices: list | None = None,
+    mask_value: int = 0,
+):
+    """Forecast masking that masks the last K patches where K is from the num_forecast_mask_patches.
+    If num_forecast_mask_patches is a list, samples in the batch will be randomly masked by numbers defined in the list.
+
+    Parameters:
+        inputs (`torch.Tensor`):
+            Input of shape `(bs, num_channels, num_patch, patch_length)`
+        num_forecast_mask_patches (`list`):
+            Number of patches to be masked at the end of each batch sample. e.g. 4 or [3, 5].
+        unmasked_channel_indices (`list`, *optional*):
+            Indices of channels that are not masked.
+        mask_value (`int`, *optional*, defaults to 0):
+            Values in the masked patches will be filled by `mask_value`.
+
+    Returns:
+        `tuple(torch.Tensor)`: inputs_mask, masked input, same shape as inputs Tensor and Mask tensor of shape `(bs,
+        num_channels , num_patch)` or `(bs, tsg1, tsg2, num_channels, num_patch)`
+    """
+
+    if isinstance(num_forecast_mask_patches, int):
+        num_forecast_mask_patches = [num_forecast_mask_patches]
+    forecast_mask_ratios = [1 for _ in num_forecast_mask_patches]
+
+    batch_size, num_channels, sequence_length, num_features = inputs.shape
+    mask = torch.zeros(batch_size, num_channels, sequence_length, device=inputs.device)
+
+    t_list = []
+    total_length = 0
+    total_ratio = sum(forecast_mask_ratios)
+
+    for patch_length, ratio in zip(num_forecast_mask_patches, forecast_mask_ratios):
+        if patch_length <= 0 or patch_length >= sequence_length:
+            raise ValueError(
+                f"num_forecast_mask_patches {patch_length} should be greater than 0 and less than total patches."
+            )
+        temp_len = int(batch_size * ratio / total_ratio)
+        t_list.append([patch_length, ratio, temp_len])
+        total_length += temp_len
+
+    t_list = sorted(t_list, key=lambda x: x[2])
+
+    if total_length < batch_size:
+        t_list[0][2] = t_list[0][2] + (batch_size - total_length)
+    elif total_length > batch_size:
+        t_list[-1][2] = t_list[-1][2] + (total_length - batch_size)
+
+    batch1 = 0
+    for patch_len, _, temp_len in t_list:
+        batch2 = batch1 + temp_len
+        mask[batch1:batch2, :, -patch_len:] = 1
+        batch1 = batch2
+
+    perm = torch.randperm(mask.shape[0])
+    mask = mask[perm]
+
+    mask = mask.unsqueeze(-1).repeat(1, 1, 1, num_features)  # mask: [bs x num_channels x num_patch x patch_len]
+    if unmasked_channel_indices is not None:
+        mask[:, unmasked_channel_indices, :, :] = 0
+
+    inputs_mask = inputs.masked_fill(mask.bool(), mask_value)
+    return inputs_mask, mask[..., 0]
+
+
+class PatchTSTPatchify(nn.Module):
+    """
+    A class to patchify the time series sequence into different patches
+
+    Returns:
+        `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`
+    """
+
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+
+        self.sequence_length = config.context_length
+        self.patch_length = config.patch_length
+        self.patch_stride = config.patch_stride
+
+        if self.sequence_length <= self.patch_length:
+            raise ValueError(
+                f"Sequence length ({self.sequence_length}) has to be greater than the patch length ({self.patch_length})"
+            )
+
+        # get the number of patches
+        self.num_patches = (max(self.sequence_length, self.patch_length) - self.patch_length) // self.patch_stride + 1
+        new_sequence_length = self.patch_length + self.patch_stride * (self.num_patches - 1)
+        self.sequence_start = self.sequence_length - new_sequence_length
+
+    def forward(self, past_values: torch.Tensor):
+        """
+        Parameters:
+            past_values (`torch.Tensor` of shape `(batch_size, sequence_length, num_channels)`, *required*):
+                Input for patchification
+
+        Returns:
+            `torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`
+        """
+        sequence_length = past_values.shape[-2]
+        if sequence_length != self.sequence_length:
+            raise ValueError(
+                f"Input sequence length ({sequence_length}) doesn't match model configuration ({self.sequence_length})."
+            )
+        # output: [bs x new_sequence_length x num_channels]
+        output = past_values[:, self.sequence_start :, :]
+        # output: [bs x num_patches x num_input_channels x patch_length]
+        output = output.unfold(dimension=-2, size=self.patch_length, step=self.patch_stride)
+        # output: [bs x num_input_channels x num_patches x patch_length]
+        output = output.transpose(-2, -3).contiguous()
+        return output
+
+
+class PatchTSTMasking(nn.Module):
+    """
+    Class to perform random or forecast masking.
+
+    Parameters:
+        config (`PatchTSTConfig`): model config
+    Returns:
+        x_mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`)
+            Masked patched input
+        mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`)
+            Bool tensor indicating True on masked points
+    """
+
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        self.random_mask_ratio = config.random_mask_ratio
+        self.channel_consistent_masking = config.channel_consistent_masking
+        self.mask_type = config.mask_type
+        self.num_forecast_mask_patches = config.num_forecast_mask_patches
+        self.unmasked_channel_indices = config.unmasked_channel_indices
+        self.mask_value = config.mask_value
+        if self.unmasked_channel_indices is not None:
+            self.unmasked_channel_indices = sorted(self.unmasked_channel_indices)
+
+    def forward(self, patch_input: torch.Tensor):
+        """
+        Parameters:
+            patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*):
+                Patch input
+
+        Return:
+            masked_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`)
+                Masked patched input
+            mask (`torch.Tensor` of shape `(batch_size, num_channels, num_patches)`)
+                Bool tensor indicating True on masked points
+
+        """
+        if self.mask_type == "random":
+            masked_input, mask = random_masking(
+                inputs=patch_input,
+                mask_ratio=self.random_mask_ratio,
+                unmasked_channel_indices=self.unmasked_channel_indices,
+                channel_consistent_masking=self.channel_consistent_masking,
+                mask_value=self.mask_value,
+            )
+        elif self.mask_type == "forecast":
+            masked_input, mask = forecast_masking(
+                inputs=patch_input,
+                num_forecast_mask_patches=self.num_forecast_mask_patches,
+                unmasked_channel_indices=self.unmasked_channel_indices,
+                mask_value=self.mask_value,
+            )
+        else:
+            raise ValueError(f"Invalid mask type {self.mask_type}.")
+
+        # mask: [bs x num_input_channels x num_patch]
+        mask = mask.bool()
+        return masked_input, mask
+
+
+class PatchTSTEncoderLayer(nn.Module):
+    """
+    PatchTST encoder layer
+    """
+
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+
+        self.channel_attention = config.channel_attention
+
+        self.self_attn = PatchTSTAttention(
+            embed_dim=config.d_model,
+            num_heads=config.num_attention_heads,
+            dropout=config.attention_dropout,
+            config=config,
+        )
+
+        # Add & Norm of the sublayer 1
+        self.dropout_path1 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()
+        if config.norm_type == "batchnorm":
+            self.norm_sublayer1 = PatchTSTBatchNorm(config)
+        elif config.norm_type == "layernorm":
+            self.norm_sublayer1 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+        else:
+            raise ValueError(f"{config.norm_type} is not a supported norm layer type.")
+
+        # Add & Norm of the sublayer 2
+        if self.channel_attention:
+            self.dropout_path2 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()
+            if config.norm_type == "batchnorm":
+                self.norm_sublayer2 = PatchTSTBatchNorm(config)
+            elif config.norm_type == "layernorm":
+                self.norm_sublayer2 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+            else:
+                raise ValueError(f"{config.norm_type} is not a supported norm layer type.")
+
+        # Position-wise Feed-Forward
+        self.ff = nn.Sequential(
+            nn.Linear(config.d_model, config.ffn_dim, bias=config.bias),
+            ACT2CLS[config.activation_function](),
+            nn.Dropout(config.ff_dropout) if config.ff_dropout > 0 else nn.Identity(),
+            nn.Linear(config.ffn_dim, config.d_model, bias=config.bias),
+        )
+
+        # Add & Norm of sublayer 3
+        self.dropout_path3 = nn.Dropout(config.path_dropout) if config.path_dropout > 0 else nn.Identity()
+        if config.norm_type == "batchnorm":
+            self.norm_sublayer3 = PatchTSTBatchNorm(config)
+        elif config.norm_type == "layernorm":
+            self.norm_sublayer3 = nn.LayerNorm(config.d_model, eps=config.norm_eps)
+        else:
+            raise ValueError(f"{config.norm_type} is not a supported norm layer type.")
+
+        self.pre_norm = config.pre_norm
+
+    def forward(self, hidden_state: torch.Tensor, output_attentions: bool | None = None):
+        """
+        Parameters:
+            hidden_state (`torch.Tensor` of shape `(batch_size, num_channels, sequence_length, d_model)`, *required*):
+                Past values of the time series
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the output attention of all layers
+        Return:
+            `torch.Tensor` of shape `(batch_size, num_channels, sequence_length, d_model)`
+
+        """
+        batch_size, num_input_channels, sequence_length, d_model = hidden_state.shape
+
+        # First sublayer: attention across time
+        # hidden_states: [(bs*num_channels) x sequence_length x d_model]
+        hidden_state = hidden_state.view(batch_size * num_input_channels, sequence_length, d_model)
+
+        if self.pre_norm:
+            ## Norm and Multi-Head attention and Add residual connection
+            attn_output, attn_weights, _ = self.self_attn(
+                hidden_states=self.norm_sublayer1(hidden_state), output_attentions=output_attentions
+            )
+            # Add: residual connection with residual dropout
+            hidden_state = hidden_state + self.dropout_path1(attn_output)
+        else:
+            ## Multi-Head attention and Add residual connection and Norm - Standard Transformer from BERT
+            attn_output, attn_weights, _ = self.self_attn(
+                hidden_states=hidden_state, output_attentions=output_attentions
+            )
+            # hidden_states: [(bs*num_channels) x sequence_length x d_model]
+            hidden_state = self.norm_sublayer1(hidden_state + self.dropout_path1(attn_output))
+
+        # hidden_state: [bs x num_channels x sequence_length x d_model]
+        hidden_state = hidden_state.reshape(batch_size, num_input_channels, sequence_length, d_model)
+
+        # second sublayer: attention across variable at any given time
+        if self.channel_attention:
+            # hidden_state: [bs x sequence_length x num_channels x d_model]
+            hidden_state = hidden_state.transpose(2, 1).contiguous()
+            # hidden_state: [(bs*sequence_length) x num_channels x d_model]
+            hidden_state = hidden_state.view(batch_size * sequence_length, num_input_channels, d_model)
+            if self.pre_norm:
+                ## Norm and Multi-Head attention and Add residual connection
+                attn_output, channel_attn_weights, _ = self.self_attn(
+                    hidden_states=self.norm_sublayer2(hidden_state), output_attentions=output_attentions
+                )
+                # Add: residual connection with residual dropout
+                hidden_state = hidden_state + self.dropout_path2(attn_output)
+            else:
+                ## Multi-Head attention and Add residual connection and Norm
+                attn_output, channel_attn_weights, _ = self.self_attn(
+                    hidden_states=hidden_state, output_attentions=output_attentions
+                )
+                # hidden_states: [(bs*sequence_length) x num_channels x d_model]
+                hidden_state = self.norm_sublayer2(hidden_state + self.dropout_path2(attn_output))
+
+            # Reshape hidden state
+            # hidden_state: [bs x sequence_length x num_channels x d_model]
+            hidden_state = hidden_state.reshape(batch_size, sequence_length, num_input_channels, d_model)
+            # hidden_state: [bs x num_channels x sequence_length x d_model]
+            hidden_state = hidden_state.transpose(1, 2).contiguous()
+
+        # Third sublayer: mixing across hidden
+        # hidden_state: [(batch_size*num_channels) x sequence_length x d_model]
+        hidden_state = hidden_state.view(batch_size * num_input_channels, sequence_length, d_model)
+        if self.pre_norm:
+            ## Norm and Position-wise Feed-Forward and Add residual connection
+            # Add: residual connection with residual dropout
+            hidden_state = hidden_state + self.dropout_path3(self.ff(self.norm_sublayer3(hidden_state)))
+        else:
+            ## Position-wise Feed-Forward and Add residual connection and Norm
+            # Add: residual connection with residual dropout
+            hidden_state = self.norm_sublayer3(hidden_state + self.dropout_path3(self.ff(hidden_state)))
+
+        # [bs x num_channels x sequence_length x d_model]
+        hidden_state = hidden_state.reshape(batch_size, num_input_channels, sequence_length, d_model)
+
+        outputs = (hidden_state,)
+        if output_attentions:
+            outputs += (attn_weights, channel_attn_weights) if self.channel_attention else (attn_weights,)
+
+        return outputs
+
+
+@auto_docstring
+class PatchTSTPreTrainedModel(PreTrainedModel):
+    config: PatchTSTConfig
+    base_model_prefix = "model"
+    main_input_name = "past_values"
+    input_modalities = ("time",)
+    supports_gradient_checkpointing = False
+    _supports_flash_attn = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+
+    @torch.no_grad()
+    def _init_weights(self, module: nn.Module):
+        """
+        Initialize weights
+        """
+        if isinstance(module, PatchTSTPositionalEncoding):
+            # get the number of patches
+            num_patches = (
+                max(self.config.context_length, self.config.patch_length) - self.config.patch_length
+            ) // self.config.patch_stride + 1
+            # initialize cls_token
+            if self.config.use_cls_token:
+                init.normal_(module.cls_token, std=0.02)
+                num_patches += 1
+            # initialize positional encoding
+            position_enc = module._init_pe(self.config, num_patches)
+            if is_deepspeed_zero3_enabled():
+                import deepspeed
+
+                with deepspeed.zero.GatheredParameters(module.position_enc, modifier_rank=None):
+                    if module.position_enc.numel() > 0:
+                        init.copy_(module.position_enc, position_enc)
+            else:
+                init.copy_(module.position_enc, position_enc)
+        elif isinstance(module, (nn.LayerNorm, nn.BatchNorm1d)):
+            init.zeros_(module.bias)
+            init.ones_(module.weight)
+            if getattr(module, "running_mean", None) is not None:
+                init.zeros_(module.running_mean)
+                init.ones_(module.running_var)
+                init.zeros_(module.num_batches_tracked)
+        elif isinstance(module, nn.Linear):
+            init.normal_(module.weight, mean=0.0, std=self.config.init_std)
+            if module.bias is not None:
+                init.zeros_(module.bias)
+
+    def _set_gradient_checkpointing(self, module, value=False):
+        if isinstance(module, (PatchTSTEncoder)):
+            module.gradient_checkpointing = value
+
+
+class PatchTSTEmbedding(nn.Module):
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        self.num_input_channels = config.num_input_channels
+        self.share_embedding = config.share_embedding
+        # Input encoding: projection of feature vectors onto a d-dim vector space
+        if self.share_embedding:
+            self.input_embedding = nn.Linear(config.patch_length, config.d_model)
+        else:
+            self.input_embedding = nn.ModuleList()
+            for _ in range(config.num_input_channels):
+                self.input_embedding.append(nn.Linear(config.patch_length, config.d_model))
+
+    def forward(self, patch_input: torch.Tensor):
+        """
+        Parameters:
+            patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*):
+                Patch input for embedding
+        return:
+            `torch.Tensor` of shape `(batch_size, num_channels, num_patches, d_model)`
+        """
+        # Input encoding
+        num_input_channels = patch_input.shape[1]
+        if num_input_channels != self.num_input_channels:
+            raise ValueError(
+                f"The defined number of input channels ({self.num_input_channels}) in the config "
+                f"has to be the same as the number of channels in the batch input ({num_input_channels})"
+            )
+        if self.share_embedding:
+            embeddings = self.input_embedding(patch_input)  # x: [bs x num_channels  x num_patches x d_model]
+        else:
+            embeddings = [self.input_embedding[i](patch_input[:, i, :, :]) for i in range(num_input_channels)]
+            embeddings = torch.stack(embeddings, dim=1)
+        return embeddings
+
+
+class PatchTSTPositionalEncoding(nn.Module):
+    """
+    Class for positional encoding
+    """
+
+    def __init__(self, config: PatchTSTConfig, num_patches: int):
+        super().__init__()
+        self.use_cls_token = config.use_cls_token
+        self.num_input_channels = config.num_input_channels
+        if config.use_cls_token:
+            # cls_token: [1 x num_input_channels x 1 x d_model]
+            self.cls_token = nn.Parameter(torch.zeros(1, 1, 1, config.d_model))
+            num_patches += 1
+        # positional encoding: [num_patches x d_model]
+        self.position_enc = self._init_pe(config, num_patches)
+        # Positional dropout
+        self.positional_dropout = (
+            nn.Dropout(config.positional_dropout) if config.positional_dropout > 0 else nn.Identity()
+        )
+
+    @staticmethod
+    def _init_pe(config: PatchTSTConfig, num_patches: int) -> nn.Parameter:
+        # Positional encoding
+        if config.positional_encoding_type == "random":
+            position_enc = nn.Parameter(torch.randn(num_patches, config.d_model), requires_grad=True)
+        elif config.positional_encoding_type == "sincos":
+            position_enc = torch.zeros(num_patches, config.d_model)
+            position = torch.arange(0, num_patches).unsqueeze(1)
+            div_term = torch.exp(torch.arange(0, config.d_model, 2) * -(math.log(10000.0) / config.d_model))
+            position_enc[:, 0::2] = torch.sin(position * div_term)
+            position_enc[:, 1::2] = torch.cos(position * div_term)
+            position_enc = position_enc - position_enc.mean()
+            position_enc = position_enc / (position_enc.std() * 10)
+            position_enc = nn.Parameter(position_enc, requires_grad=False)
+        else:
+            raise ValueError(
+                f"{config.positional_encoding_type} is not a valid positional encoder. Available types are 'random' and 'sincos'."
+            )
+        return position_enc
+
+    def forward(self, patch_input: torch.Tensor):
+        if self.use_cls_token:
+            # patch_input: [bs x num_channels x num_patches x d_model]
+            patch_input = self.positional_dropout(patch_input + self.position_enc[1:, :])
+            # append cls token where cls_token: [1 x num_channels x 1 x d_model]
+            cls_token = self.cls_token + self.position_enc[:1, :]
+            # get the same copy of cls_token for all the samples in batch: [bs x num_channels x 1 x d_model]
+            cls_tokens = cls_token.expand(patch_input.shape[0], self.num_input_channels, -1, -1)
+            # hidden_state: [bs x num_channels x (num_patches+1) x d_model]
+            hidden_state = torch.cat((cls_tokens, patch_input), dim=2)
+        else:
+            # hidden_state: [bs x num_channels x num_patches x d_model]
+            hidden_state = self.positional_dropout(patch_input + self.position_enc)
+        return hidden_state
+
+
+class PatchTSTEncoder(PatchTSTPreTrainedModel):
+    """
+    PatchTST Encoder
+    """
+
+    def __init__(self, config: PatchTSTConfig, num_patches: int):
+        super().__init__(config)
+        self.gradient_checkpointing = False
+
+        # Input embedding: projection of feature vectors onto a d-dim vector space
+        self.embedder = PatchTSTEmbedding(config)
+        # Positional encoding
+        self.positional_encoder = PatchTSTPositionalEncoding(config, num_patches)
+        # Encoder
+        self.layers = nn.ModuleList([PatchTSTEncoderLayer(config) for i in range(config.num_hidden_layers)])
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        patch_input: torch.Tensor,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        **kwargs,
+    ) -> BaseModelOutput:
+        """
+        Parameters:
+            patch_input (`torch.Tensor` of shape `(batch_size, num_channels, num_patches, patch_length)`, *required*):
+                Past values of the time series
+            output_hidden_states (bool, optional): Indicates if hidden states should be outputted.
+            output_attentions (bool, optional): Indicates if attentions should be outputted.
+
+        return:
+            `BaseModelOutput`
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        # Input embedding
+        patch_input = self.embedder(patch_input)
+        # Positional encoding
+        hidden_state = self.positional_encoder(patch_input)
+
+        encoder_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        for encoder_layer in self.layers:
+            if output_hidden_states:
+                encoder_states = encoder_states + (hidden_state,)
+
+            layer_outputs = encoder_layer(hidden_state=hidden_state, output_attentions=output_attentions)
+            # get hidden state. hidden_state shape is [bs x num_channels x num_patches x d_model]
+            # or [bs x num_channels x (num_patches+1) x d_model] if use cls_token
+            hidden_state = layer_outputs[0]
+            # append attention matrix at each layer
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+        # return past_values, hidden_states
+        return BaseModelOutput(last_hidden_state=hidden_state, hidden_states=encoder_states, attentions=all_attentions)
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for model's outputs, with potential hidden states.
+    """
+)
+@dataclass
+class PatchTSTModelOutput(ModelOutput):
+    r"""
+    last_hidden_state (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`):
+        Sequence of hidden-states at the output of the last layer of the model.
+    hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`):
+        Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, +
+        one for the output of each layer) of shape `(batch_size, num_channels, height, width)`. Hidden-states of
+        the model at the output of each layer plus the optional initial embedding outputs.
+    mask (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches)`, *optional*):
+        Bool masked tensor indicating which patches are masked
+    loc (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*):
+        Mean of the input data (batch_size, sequence_length, num_channels) over the sequence_length
+    scale (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*):
+        Std of the input data (batch_size, sequence_length, num_channels) over the sequence_length
+    patch_input (`torch.FloatTensor` of shape `(batch_size, num_channels, num_patches, patch_length)`):
+        Patched input to the Transformer
+    """
+
+    last_hidden_state: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+    mask: torch.FloatTensor | None = None
+    loc: torch.FloatTensor | None = None
+    scale: torch.FloatTensor | None = None
+    patch_input: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Output type of [`PatchTSTForPretraining`].
+    """
+)
+@dataclass
+class PatchTSTForPretrainingOutput(ModelOutput):
+    r"""
+    loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+        MSE loss.
+    prediction_output (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`):
+        Prediction outputs of the time series modeling heads.
+    """
+
+    loss: torch.FloatTensor | None = None
+    prediction_output: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Output type of [`PatchTSTForRegression`].
+    """
+)
+@dataclass
+class PatchTSTForRegressionOutput(ModelOutput):
+    r"""
+    loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+        MSE loss.
+    regression_outputs (`torch.FloatTensor` of shape `(batch_size, num_targets)`):
+        Regression outputs of the time series modeling heads.
+    """
+
+    loss: torch.FloatTensor | None = None
+    regression_outputs: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Output type of [`PatchTSTForPrediction`].
+    """
+)
+@dataclass
+class PatchTSTForPredictionOutput(ModelOutput):
+    r"""
+    loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+        MSE loss.
+    prediction_outputs (`torch.FloatTensor` of shape `(batch_size, prediction_length, -1)`):
+        Prediction outputs of the time series modeling heads.
+    attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`):
+        Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length,
+        sequence_length)`.
+
+        Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+        heads.
+    loc: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*)
+        Mean of the input data (batch_size, sequence_length, num_channels) over the sequence_length
+    scale: (`torch.FloatTensor` of shape `(batch_size, 1, num_channels)`, *optional*)
+        Std of the input data (batch_size, sequence_length, num_channels) over the sequence_length
+    """
+
+    loss: torch.FloatTensor | None = None
+    prediction_outputs: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+    loc: torch.FloatTensor | None = None
+    scale: torch.FloatTensor | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Output type of [`PatchTSTForClassification`].
+    """
+)
+@dataclass
+class PatchTSTForClassificationOutput(ModelOutput):
+    r"""
+    loss (*optional*, returned when `labels` is provided, `torch.FloatTensor` of shape `(1,)`):
+        Total loss as the sum of the masked language modeling loss and the next sequence prediction
+        (classification) loss.
+    prediction_logits (`torch.FloatTensor` of shape `(batch_size, num_targets)`):
+        Prediction scores of the PatchTST modeling head (scores before SoftMax).
+    """
+
+    loss: torch.FloatTensor | None = None
+    prediction_logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor] | None = None
+    attentions: tuple[torch.FloatTensor] | None = None
+
+
+@auto_docstring(
+    custom_intro="""
+    Base class for time series model's predictions outputs that contains the sampled values from the chosen
+    distribution.
+    """
+)
+@dataclass
+class SamplePatchTSTOutput(ModelOutput):
+    r"""
+    sequences (`torch.FloatTensor` of shape `(batch_size, num_samples, prediction_length, num_targets)`):
+        Sampled values from the chosen distribution.
+    """
+
+    sequences: torch.FloatTensor | None = None
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.nll
+def nll(input: torch.distributions.Distribution, target: torch.Tensor) -> torch.Tensor:
+    """
+    Computes the negative log likelihood loss from input distribution with respect to target.
+    """
+    return -input.log_prob(target)
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.weighted_average
+def weighted_average(input_tensor: torch.Tensor, weights: torch.Tensor | None = None, dim=None) -> torch.Tensor:
+    """
+    Computes the weighted average of a given tensor across a given `dim`, masking values associated with weight zero,
+    meaning instead of `nan * 0 = nan` you will get `0 * 0 = 0`.
+
+    Args:
+        input_tensor (`torch.FloatTensor`):
+            Input tensor, of which the average must be computed.
+        weights (`torch.FloatTensor`, *optional*):
+            Weights tensor, of the same shape as `input_tensor`.
+        dim (`int`, *optional*):
+            The dim along which to average `input_tensor`.
+
+    Returns:
+        `torch.FloatTensor`: The tensor with values averaged along the specified `dim`.
+    """
+    if weights is not None:
+        weighted_tensor = torch.where(weights != 0, input_tensor * weights, torch.zeros_like(input_tensor))
+        sum_weights = torch.clamp(weights.sum(dim=dim) if dim else weights.sum(), min=1.0)
+        return (weighted_tensor.sum(dim=dim) if dim else weighted_tensor.sum()) / sum_weights
+    else:
+        return input_tensor.mean(dim=dim)
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesStdScaler with TimeSeriesTransformer->PatchTST,TimeSeries->PatchTST
+class PatchTSTStdScaler(nn.Module):
+    """
+    Standardize features by calculating the mean and scaling along the first dimension, and then normalizes it by
+    subtracting from the mean and dividing by the standard deviation.
+    """
+
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+        self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+        self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-5
+
+    def forward(
+        self, data: torch.Tensor, observed_indicator: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                input for Batch norm calculation
+            observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                Calculating the scale on the observed indicator.
+        Returns:
+            tuple of `torch.Tensor` of shapes
+                (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+                `(batch_size, 1, num_input_channels)`)
+        """
+        denominator = observed_indicator.sum(self.dim, keepdim=self.keepdim)
+        denominator = denominator.clamp_min(1.0)
+        loc = (data * observed_indicator).sum(self.dim, keepdim=self.keepdim) / denominator
+
+        variance = (((data - loc) * observed_indicator) ** 2).sum(self.dim, keepdim=self.keepdim) / denominator
+        scale = torch.sqrt(variance + self.minimum_scale)
+        return (data - loc) / scale, loc, scale
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesMeanScaler with TimeSeriesTransformer->PatchTST,TimeSeries->PatchTST
+class PatchTSTMeanScaler(nn.Module):
+    """
+    Computes a scaling factor as the weighted average absolute value along the first dimension, and scales the data
+    accordingly.
+    """
+
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+        self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+        self.minimum_scale = config.minimum_scale if hasattr(config, "minimum_scale") else 1e-10
+        self.default_scale = config.default_scale if hasattr(config, "default_scale") else None
+
+    def forward(
+        self, data: torch.Tensor, observed_indicator: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                input for Batch norm calculation
+            observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                Calculating the scale on the observed indicator.
+        Returns:
+            tuple of `torch.Tensor` of shapes
+                (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+                `(batch_size, 1, num_input_channels)`)
+        """
+        ts_sum = (data * observed_indicator).abs().sum(self.dim, keepdim=True)
+        num_observed = observed_indicator.sum(self.dim, keepdim=True)
+
+        scale = ts_sum / torch.clamp(num_observed, min=1)
+
+        # If `default_scale` is provided, we use it, otherwise we use the scale
+        # of the batch.
+        if self.default_scale is None:
+            batch_sum = ts_sum.sum(dim=0)
+            batch_observations = torch.clamp(num_observed.sum(0), min=1)
+            default_scale = torch.squeeze(batch_sum / batch_observations)
+        else:
+            default_scale = self.default_scale * torch.ones_like(scale)
+
+        # apply default scale where there are no observations
+        scale = torch.where(num_observed > 0, scale, default_scale)
+
+        # ensure the scale is at least `self.minimum_scale`
+        scale = torch.clamp(scale, min=self.minimum_scale)
+        scaled_data = data / scale
+
+        if not self.keepdim:
+            scale = scale.squeeze(dim=self.dim)
+
+        return scaled_data, torch.zeros_like(scale), scale
+
+
+# Copied from transformers.models.time_series_transformer.modeling_time_series_transformer.TimeSeriesNOPScaler with TimeSeriesTransformer->PatchTST,TimeSeries->PatchTST
+class PatchTSTNOPScaler(nn.Module):
+    """
+    Assigns a scaling factor equal to 1 along the first dimension, and therefore applies no scaling to the input data.
+    """
+
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        self.dim = config.scaling_dim if hasattr(config, "scaling_dim") else 1
+        self.keepdim = config.keepdim if hasattr(config, "keepdim") else True
+
+    def forward(
+        self, data: torch.Tensor, observed_indicator: torch.Tensor | None = None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                input for Batch norm calculation
+        Returns:
+            tuple of `torch.Tensor` of shapes
+                (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+                `(batch_size, 1, num_input_channels)`)
+        """
+        scale = torch.ones_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
+        loc = torch.zeros_like(data, requires_grad=False).mean(dim=self.dim, keepdim=self.keepdim)
+        return data, loc, scale
+
+
+class PatchTSTScaler(nn.Module):
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        if config.scaling == "mean" or config.scaling is True:
+            self.scaler = PatchTSTMeanScaler(config)
+        elif config.scaling == "std":
+            self.scaler = PatchTSTStdScaler(config)
+        else:
+            self.scaler = PatchTSTNOPScaler(config)
+
+    def forward(
+        self, data: torch.Tensor, observed_indicator: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Parameters:
+            data (`torch.Tensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                Input for scaler calculation
+            observed_indicator (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                Calculating the scale on the observed indicator.
+        Returns:
+            tuple of `torch.Tensor` of shapes
+                (`(batch_size, sequence_length, num_input_channels)`,`(batch_size, 1, num_input_channels)`,
+                `(batch_size, 1, um_input_channels)`)
+        """
+        data, loc, scale = self.scaler(data, observed_indicator)
+        return data, loc, scale
+
+
+@auto_docstring
+class PatchTSTModel(PatchTSTPreTrainedModel):
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__(config)
+
+        self.scaler = PatchTSTScaler(config)
+        self.patchifier = PatchTSTPatchify(config)
+        self.do_mask_input = config.do_mask_input
+        # get num_patches information from PatchTSTPatchify
+        num_patches = self.patchifier.num_patches
+
+        if self.do_mask_input:
+            self.masking = PatchTSTMasking(config)
+        else:
+            self.masking = nn.Identity()
+        self.encoder = PatchTSTEncoder(config, num_patches=num_patches)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        past_values: torch.Tensor,
+        past_observed_mask: torch.Tensor | None = None,
+        future_values: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | PatchTSTModelOutput:
+        r"""
+        Parameters:
+            past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*):
+                Input sequence to the model
+            past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+                Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+                in `[0, 1]`:
+
+                - 1 for values that are **observed**,
+                - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+            future_values (`torch.BoolTensor` of shape `(batch_size, prediction_length, num_input_channels)`, *optional*):
+                Future target values associated with the `past_values`
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the output attention of all layers
+            return_dict (`bool`, *optional*):
+                Whether or not to return a `ModelOutput` instead of a plain tuple.
+
+        Returns:
+            `PatchTSTModelOutput` or tuple of `torch.Tensor` (if `return_dict`=False or `config.return_dict`=False)
+
+        Examples:
+
+        ```python
+        >>> from huggingface_hub import hf_hub_download
+        >>> import torch
+        >>> from transformers import PatchTSTModel
+
+        >>> file = hf_hub_download(
+        ...     repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset"
+        ... )
+        >>> batch = torch.load(file)
+
+        >>> model = PatchTSTModel.from_pretrained("namctin/patchtst_etth1_pretrain")
+
+        >>> # during training, one provides both past and future values
+        >>> outputs = model(
+        ...     past_values=batch["past_values"],
+        ...     future_values=batch["future_values"],
+        ... )
+
+        >>> last_hidden_state = outputs.last_hidden_state
+        ```"""
+
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+
+        if past_observed_mask is None:
+            past_observed_mask = torch.ones_like(past_values)
+
+        # x: tensor [bs x sequence_length x num_input_channels]
+        scaled_past_values, loc, scale = self.scaler(past_values, past_observed_mask)
+
+        # patched_values: [bs x num_input_channels x num_patches x patch_length] for pretrain
+        patched_values = self.patchifier(scaled_past_values)
+        if self.do_mask_input:
+            masked_values, mask = self.masking(patched_values)
+        else:
+            masked_values, mask = self.masking(patched_values), None
+
+        encoder_output = self.encoder(
+            patch_input=masked_values, output_hidden_states=output_hidden_states, output_attentions=output_attentions
+        )
+
+        if not return_dict:
+            outputs = (encoder_output.last_hidden_state, encoder_output.hidden_states, encoder_output.attentions)
+            outputs = outputs + (mask, loc, scale, patched_values)
+            return tuple(v for v in outputs if v is not None)
+
+        return PatchTSTModelOutput(
+            last_hidden_state=encoder_output.last_hidden_state,
+            hidden_states=encoder_output.hidden_states,
+            attentions=encoder_output.attentions,
+            mask=mask,
+            loc=loc,
+            scale=scale,
+            patch_input=patched_values,
+        )
+
+
+class PatchTSTMaskPretrainHead(nn.Module):
+    """
+    Pretraining head for mask modelling
+    """
+
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        self.dropout = nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity()
+        self.linear = nn.Linear(config.d_model, config.patch_length)
+        self.use_cls_token = config.use_cls_token
+
+    def forward(self, embedding: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters:
+            embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or
+                    `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*):
+                Embedding from the model
+        Returns:
+            `torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or
+                            `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True
+
+        """
+        embedding = self.linear(self.dropout(embedding))  # [bs x num_channels x num_patches x patch_length]
+        if self.use_cls_token:
+            embedding = embedding[:, :, 1:, :]  # remove the first cls token
+        return embedding
+
+
+@auto_docstring(
+    custom_intro="""
+    The PatchTST for pretrain model.
+    """
+)
+class PatchTSTForPretraining(PatchTSTPreTrainedModel):
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__(config)
+
+        config.do_mask_input = True
+        self.model = PatchTSTModel(config=config)
+        self.head = PatchTSTMaskPretrainHead(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        past_values: torch.Tensor,
+        past_observed_mask: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | PatchTSTForPretrainingOutput:
+        r"""
+        Parameters:
+            past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*):
+                Input sequence to the model
+            past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+                Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+                in `[0, 1]`:
+
+                - 1 for values that are **observed**,
+                - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the output attention of all layers
+            return_dict (`bool`, *optional*): Whether or not to return a `ModelOutput` instead of a plain tuple.
+
+        Returns:
+            `PatchTSTForPretrainingOutput` or tuple of `torch.Tensor` (if `return_dict`=False or
+            `config.return_dict`=False)
+
+        Examples:
+
+        ```python
+        >>> from huggingface_hub import hf_hub_download
+        >>> import torch
+        >>> from transformers import PatchTSTConfig, PatchTSTForPretraining
+
+        >>> file = hf_hub_download(
+        ...     repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset"
+        ... )
+        >>> batch = torch.load(file)
+
+        >>> # Config for random mask pretraining
+        >>> config = PatchTSTConfig(
+        ...     num_input_channels=7,
+        ...     context_length=512,
+        ...     patch_length=12,
+        ...     stride=12,
+        ...     mask_type='random',
+        ...     random_mask_ratio=0.4,
+        ...     use_cls_token=True,
+        ... )
+        >>> # Config for forecast mask pretraining
+        >>> config = PatchTSTConfig(
+        ...     num_input_channels=7,
+        ...     context_length=512,
+        ...     patch_length=12,
+        ...     stride=12,
+        ...     mask_type='forecast',
+        ...     num_forecast_mask_patches=5,
+        ...     use_cls_token=True,
+        ... )
+        >>> model = PatchTSTForPretraining(config)
+
+        >>> # during training, one provides both past and future values
+        >>> outputs = model(past_values=batch["past_values"])
+
+        >>> loss = outputs.loss
+        >>> loss.backward()
+        ```"""
+
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        # past_values: [bs x num_channels x num_patches x d_model] or
+        # [bs x num_channels x (num_patches+1) x d_model] if use cls_token
+        model_output = self.model(
+            past_values=past_values,
+            past_observed_mask=past_observed_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+
+        # last_hidden_state: [bs x num_channels x num_patches x patch_length] or
+        # [bs x num_channels x (num_patches+1) x patch_length] if use cls_token
+        x_hat = self.head(model_output.last_hidden_state)
+
+        # calculate masked_loss
+        loss = nn.MSELoss(reduction="none")
+        loss_val = loss(x_hat, model_output.patch_input)
+        masked_loss = (loss_val.mean(dim=-1) * model_output.mask).sum() / (model_output.mask.sum() + 1e-10)
+
+        encoder_states = model_output.hidden_states
+        if not return_dict:
+            outputs = (x_hat,) + model_output[1:-4]
+            outputs = (masked_loss,) + outputs if masked_loss is not None else outputs
+            return outputs
+        return PatchTSTForPretrainingOutput(
+            loss=masked_loss, prediction_output=x_hat, hidden_states=encoder_states, attentions=model_output.attentions
+        )
+
+
+class PatchTSTClassificationHead(nn.Module):
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__()
+        self.use_cls_token = config.use_cls_token
+        self.pooling_type = config.pooling_type
+        self.flatten = nn.Flatten(start_dim=1)
+        self.dropout = nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity()
+        self.linear = nn.Linear(config.num_input_channels * config.d_model, config.num_targets)
+
+    def forward(self, embedding: torch.Tensor):
+        """
+        Parameters:
+            embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or
+                     `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*):
+                Embedding from the model
+        Returns:
+            `torch.Tensor` of shape `(bs, num_targets)`
+
+        """
+        if self.use_cls_token:
+            # use the first output token, pooled_embedding: bs x num_channels x d_model
+            pooled_embedding = embedding[:, :, 0, :]
+        elif self.pooling_type == "mean":
+            # pooled_embedding: [bs x num_channels x d_model]
+            pooled_embedding = embedding.mean(dim=2)
+        elif self.pooling_type == "max":
+            # pooled_embedding: [bs x num_channels x d_model]
+            pooled_embedding = embedding.max(dim=2).values
+        else:
+            raise ValueError(f"pooling operator {self.pooling_type} is not implemented yet")
+        # pooled_embedding: bs x num_channels * d_model
+        pooled_embedding = self.flatten(pooled_embedding)
+        # output: bs x n_classes
+        output = self.linear(self.dropout(pooled_embedding))
+        return output
+
+
+@auto_docstring(
+    custom_intro="""
+    The PatchTST for classification model.
+    """
+)
+class PatchTSTForClassification(PatchTSTPreTrainedModel):
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__(config)
+
+        # Turn off masking
+        if config.do_mask_input:
+            logger.warning("Setting `do_mask_input` parameter to False.")
+            config.do_mask_input = False
+
+        self.model = PatchTSTModel(config)
+        self.head = PatchTSTClassificationHead(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        past_values: torch.Tensor,
+        target_values: torch.Tensor | None = None,
+        past_observed_mask: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | PatchTSTForClassificationOutput:
+        r"""
+        past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*):
+            Input sequence to the model
+        target_values (`torch.Tensor`, *optional*):
+            Labels associates with the `past_values`
+        past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+            Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+            in `[0, 1]`:
+
+            - 1 for values that are **observed**,
+            - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+
+        Examples:
+
+        ```python
+        >>> from transformers import PatchTSTConfig, PatchTSTForClassification
+
+        >>> # classification task with two input channel2 and 3 classes
+        >>> config = PatchTSTConfig(
+        ...     num_input_channels=2,
+        ...     num_targets=3,
+        ...     context_length=512,
+        ...     patch_length=12,
+        ...     stride=12,
+        ...     use_cls_token=True,
+        ... )
+        >>> model = PatchTSTForClassification(config=config)
+
+        >>> # during inference, one only provides past values
+        >>> past_values = torch.randn(20, 512, 2)
+        >>> outputs = model(past_values=past_values)
+        >>> labels = outputs.prediction_logits
+        ```"""
+
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        model_output = self.model(
+            past_values=past_values,
+            past_observed_mask=past_observed_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        y_hat = self.head(model_output.last_hidden_state)
+
+        loss_val = None
+        if target_values is not None:
+            loss = nn.CrossEntropyLoss()
+            loss_val = loss(y_hat, target_values)
+
+        if not return_dict:
+            outputs = (y_hat,) + model_output[1:-3]
+            outputs = (loss_val,) + outputs if loss_val is not None else outputs
+            return outputs
+        return PatchTSTForClassificationOutput(
+            loss=loss_val,
+            prediction_logits=y_hat,
+            hidden_states=model_output.hidden_states,
+            attentions=model_output.attentions,
+        )
+
+
+@auto_docstring(
+    custom_intro="""
+    The PatchTST for regression Model.
+    """
+)
+class PatchTSTPredictionHead(nn.Module):
+    def __init__(self, config: PatchTSTConfig, num_patches: int, distribution_output=None):
+        r"""
+        num_patches (`int`):
+            The number of patches in the input sequence.
+        distribution_output (`DistributionOutput`, *optional*):
+            The distribution output layer for probabilistic forecasting. If None, a linear output layer is used.
+        """
+        super().__init__()
+
+        self.share_projection = config.share_projection
+        self.num_input_channels = config.num_input_channels
+        self.use_cls_token = config.use_cls_token
+        self.pooling_type = config.pooling_type
+        if self.pooling_type or self.use_cls_token:
+            head_dim = config.d_model
+        else:
+            head_dim = config.d_model * num_patches
+
+        if not self.share_projection:
+            # if each channel has its own head
+            self.projections = nn.ModuleList()
+            self.dropouts = nn.ModuleList()
+            self.flattens = nn.ModuleList()
+            for i in range(self.num_input_channels):
+                self.flattens.append(nn.Flatten(start_dim=2))
+                if distribution_output is None:
+                    # use linear head
+                    self.projections.append(nn.Linear(head_dim, config.prediction_length))
+                else:
+                    # use distribution head
+                    self.projections.append(distribution_output.get_parameter_projection(head_dim))
+                self.dropouts.append(nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity())
+        else:
+            # all the channels share the same head
+            self.flatten = nn.Flatten(start_dim=2)
+            if distribution_output is None:
+                # use linear head
+                self.projection = nn.Linear(head_dim, config.prediction_length)
+            else:
+                # use distribution head
+                self.projection = distribution_output.get_parameter_projection(head_dim)
+            self.dropout = nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity()
+
+    def forward(self, embedding: torch.Tensor):
+        """
+        Parameters:
+            embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or
+                     `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*):
+                Embedding from the model
+        Returns:
+            `torch.Tensor` of shape `(bs, forecast_len, num_channels)`
+
+        """
+        if self.use_cls_token:
+            # pooled_embedding: [bs x num_channels x d_model]
+            pooled_embedding = embedding[:, :, 0, :]
+        else:
+            if self.pooling_type == "mean":
+                # pooled_embedding: [bs x num_channels x d_model]
+                pooled_embedding = embedding.mean(dim=2)
+            elif self.pooling_type == "max":
+                # pooled_embedding: [bs x num_channels x d_model]
+                pooled_embedding = embedding.max(dim=2).values
+            else:
+                # pooled_embedding: [bs x num_channels x num_patches x d_model]
+                pooled_embedding = embedding
+
+        if not self.share_projection:
+            output = []
+            for i in range(self.num_input_channels):
+                # pooled_embedding: [bs x (d_model * num_patches)] or [bs x d_model)]
+                pooled_embedding = self.flattens[i](pooled_embedding[:, i, :])
+                pooled_embedding = self.dropouts[i](pooled_embedding)
+                # pooled_embedding: [bs x forecast_len]
+                #  or tuple ([bs x forecast_len], [bs x forecast_len]) if using distribution head
+                pooled_embedding = self.projections[i](pooled_embedding)
+                output.append(pooled_embedding)
+            # output: [bs x num_channels x forecast_len]
+            output = torch.stack(output, dim=1)
+        else:
+            # pooled_embedding: [bs x num_channels x (d_model * num_patches)] or [bs x num_channels x d_model)]
+            pooled_embedding = self.flatten(pooled_embedding)
+            pooled_embedding = self.dropout(pooled_embedding)
+            # output: [bs x num_channels x forecast_len] or
+            # tuple ([bs x num_channels x forecast_len], [bs x num_channels x forecast_len]) if using distribution head
+            output = self.projection(pooled_embedding)
+
+        if isinstance(output, tuple):
+            # output: ([bs x forecast_len x num_channels], [bs x forecast_len x num_channels])
+            output = tuple(z.transpose(2, 1) for z in output)
+        else:
+            output = output.transpose(2, 1)  # [bs x forecast_len x num_channels]
+        return output
+
+
+@auto_docstring(
+    custom_intro="""
+    The PatchTST for prediction model.
+    """
+)
+class PatchTSTForPrediction(PatchTSTPreTrainedModel):
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__(config)
+
+        # Turn off masking
+        if config.do_mask_input:
+            logger.warning("Setting `do_mask_input` parameter to False.")
+            config.do_mask_input = False
+
+        self.model = PatchTSTModel(config)
+
+        if config.loss == "mse":
+            self.distribution_output = None
+        else:
+            if config.distribution_output == "student_t":
+                self.distribution_output = StudentTOutput(dim=config.prediction_length)
+            elif config.distribution_output == "normal":
+                self.distribution_output = NormalOutput(dim=config.prediction_length)
+            elif config.distribution_output == "negative_binomial":
+                self.distribution_output = NegativeBinomialOutput(dim=config.prediction_length)
+            else:
+                raise ValueError(f"Unknown distribution output {config.distribution_output}")
+
+        self.head = PatchTSTPredictionHead(
+            config, self.model.patchifier.num_patches, distribution_output=self.distribution_output
+        )
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        past_values: torch.Tensor,
+        past_observed_mask: torch.Tensor | None = None,
+        future_values: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | PatchTSTForPredictionOutput:
+        r"""
+        Parameters:
+            past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*):
+                Input sequence to the model
+            past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+                Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+                in `[0, 1]`:
+
+                - 1 for values that are **observed**,
+                - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+            future_values (`torch.Tensor` of shape `(bs, forecast_len, num_input_channels)`, *optional*):
+                Future target values associated with the `past_values`
+            output_hidden_states (`bool`, *optional*):
+                Whether or not to return the hidden states of all layers
+            output_attentions (`bool`, *optional*):
+                Whether or not to return the output attention of all layers
+            return_dict (`bool`, *optional*):
+                Whether or not to return a `ModelOutput` instead of a plain tuple.
+
+        Returns:
+            `PatchTSTForPredictionOutput` or tuple of `torch.Tensor` (if `return_dict`=False or
+            `config.return_dict`=False)
+
+        Examples:
+
+        ```python
+        >>> from huggingface_hub import hf_hub_download
+        >>> import torch
+        >>> from transformers import PatchTSTConfig, PatchTSTForPrediction
+
+        >>> file = hf_hub_download(
+        ...     repo_id="hf-internal-testing/etth1-hourly-batch", filename="train-batch.pt", repo_type="dataset"
+        ... )
+        >>> batch = torch.load(file)
+
+        >>> # Prediction task with 7 input channels and prediction length is 96
+        >>> model = PatchTSTForPrediction.from_pretrained("namctin/patchtst_etth1_forecast")
+
+        >>> # during training, one provides both past and future values
+        >>> outputs = model(
+        ...     past_values=batch["past_values"],
+        ...     future_values=batch["future_values"],
+        ... )
+
+        >>> loss = outputs.loss
+        >>> loss.backward()
+
+        >>> # during inference, one only provides past values, the model outputs future values
+        >>> outputs = model(past_values=batch["past_values"])
+        >>> prediction_outputs = outputs.prediction_outputs
+        ```"""
+
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        # get model output
+        model_output = self.model(
+            past_values=past_values,
+            past_observed_mask=past_observed_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        # get output head
+        y_hat = self.head(model_output.last_hidden_state)
+
+        loss_val = None
+
+        if self.distribution_output:
+            y_hat_out = y_hat
+        else:
+            y_hat_out = y_hat * model_output.scale + model_output.loc
+
+        if future_values is not None:
+            if self.distribution_output:
+                distribution = self.distribution_output.distribution(
+                    y_hat, loc=model_output.loc, scale=model_output.scale
+                )
+                loss_val = nll(distribution, future_values)
+                # take average of the loss
+                loss_val = weighted_average(loss_val)
+            else:
+                loss = nn.MSELoss(reduction="mean")
+                loss_val = loss(y_hat_out, future_values)
+
+        loc = model_output.loc
+        scale = model_output.scale
+
+        if not return_dict:
+            outputs = (y_hat_out,) + model_output[1:-1]
+            outputs = (loss_val,) + outputs if loss_val is not None else outputs
+            return outputs
+        return PatchTSTForPredictionOutput(
+            loss=loss_val,
+            prediction_outputs=y_hat_out,
+            hidden_states=model_output.hidden_states,
+            attentions=model_output.attentions,
+            loc=loc,
+            scale=scale,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        past_values: torch.Tensor,
+        past_observed_mask: torch.Tensor | None = None,
+    ) -> SamplePatchTSTOutput:
+        """
+        Generate sequences of sample predictions from a model with a probability distribution head.
+
+        Parameters:
+            past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                Past values of the time series that serves as context in order to predict the future.
+            past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+                Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+                in `[0, 1]`:
+
+                - 1 for values that are **observed**,
+                - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+
+        Return:
+            [`SamplePatchTSTOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of
+            samples, prediction_length, 1)` or `(batch_size, number of samples, prediction_length, num_input_channels)`
+            for multivariate predictions.
+        """
+        # get number of samples
+        num_parallel_samples = self.config.num_parallel_samples
+
+        # get model output
+        outputs = self(
+            past_values=past_values,
+            future_values=None,
+            past_observed_mask=past_observed_mask,
+            output_hidden_states=False,
+        )
+        if self.distribution_output:
+            # get distribution
+            distribution = self.distribution_output.distribution(
+                outputs.prediction_outputs, loc=outputs.loc, scale=outputs.scale
+            )
+            # get samples: list of [bs x forecast_len x num_channels]
+            samples = [distribution.sample() for _ in range(num_parallel_samples)]
+            # samples: [bs x num_samples x forecast_len x num_channels]
+            samples = torch.stack(samples, dim=1)
+        else:
+            samples = outputs.prediction_outputs.unsqueeze(1)
+
+        return SamplePatchTSTOutput(sequences=samples)
+
+
+class PatchTSTRegressionHead(nn.Module):
+    """
+    Regression head
+    """
+
+    def __init__(self, config: PatchTSTConfig, distribution_output=None):
+        super().__init__()
+        self.y_range = config.output_range
+        self.use_cls_token = config.use_cls_token
+        self.pooling_type = config.pooling_type
+        self.distribution_output = distribution_output
+
+        head_dim = config.num_input_channels * config.d_model
+
+        self.flatten = nn.Flatten(start_dim=1)
+        self.dropout = nn.Dropout(config.head_dropout) if config.head_dropout > 0 else nn.Identity()
+
+        if distribution_output is None:
+            self.projection = nn.Linear(head_dim, config.num_targets)
+        else:
+            self.projection = distribution_output.get_parameter_projection(head_dim)
+
+    def forward(self, embedding: torch.Tensor):
+        """
+        Parameters:
+            embedding (`torch.Tensor` of shape `(bs, num_channels, num_patches, d_model)` or
+                    `(bs, num_channels, num_patches+1, d_model)` if `cls_token` is set to True, *required*):
+                Embedding from the model
+        Returns:
+            `torch.Tensor` of shape `(bs, output_dim)`
+
+        """
+        if self.use_cls_token:
+            # use the first output token, pooled_embedding: [bs x num_channels x d_model]
+            pooled_embedding = embedding[:, :, 0, :]
+        elif self.pooling_type == "mean":
+            # pooled_embedding: [bs x num_channels x d_model]
+            pooled_embedding = embedding.mean(dim=2)
+        elif self.pooling_type == "max":
+            # pooled_embedding: [bs x num_channels x d_model]
+            pooled_embedding = embedding.max(dim=2).values
+        else:
+            raise ValueError(f"pooling operator {self.pooling_type} is not implemented yet")
+        # flatten the input
+        # pooled_embedding: bs x (num_channels * d_model)
+        pooled_embedding = self.dropout(self.flatten(pooled_embedding))
+        # projection
+        # output: bs x output_dim or a tuple of this shape for distribution head
+        output = self.projection(pooled_embedding)
+        # apply sigmoid to bound the output if required
+        if (self.distribution_output is None) & (self.y_range is not None):  # linear head
+            output = torch.sigmoid(output) * (self.y_range[1] - self.y_range[0]) + self.y_range[0]
+        return output
+
+
+@auto_docstring(
+    custom_intro="""
+    The PatchTST for regression model.
+    """
+)
+class PatchTSTForRegression(PatchTSTPreTrainedModel):
+    def __init__(self, config: PatchTSTConfig):
+        super().__init__(config)
+
+        # Turn off masking
+        if config.do_mask_input:
+            logger.warning("Setting `do_mask_input` parameter to False.")
+            config.do_mask_input = False
+
+        self.model = PatchTSTModel(config)
+        if config.loss == "mse":
+            self.distribution_output = None
+        else:
+            if config.distribution_output == "student_t":
+                self.distribution_output = StudentTOutput(dim=config.num_targets)
+            elif config.distribution_output == "normal":
+                self.distribution_output = NormalOutput(dim=config.num_targets)
+            elif config.distribution_output == "negative_binomial":
+                self.distribution_output = NegativeBinomialOutput(dim=config.num_targets)
+            else:
+                raise ValueError(f"Unknown distribution output {config.distribution_output}")
+
+        self.head = PatchTSTRegressionHead(config, self.distribution_output)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @auto_docstring
+    def forward(
+        self,
+        past_values: torch.Tensor,
+        target_values: torch.Tensor | None = None,
+        past_observed_mask: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> tuple | PatchTSTForRegressionOutput:
+        r"""
+        past_values (`torch.Tensor` of shape `(bs, sequence_length, num_input_channels)`, *required*):
+            Input sequence to the model
+        target_values (`torch.Tensor` of shape `(bs, num_input_channels)`):
+            Target values associates with the `past_values`
+        past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+            Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+            in `[0, 1]`:
+
+            - 1 for values that are **observed**,
+            - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+            Whether or not to return a `ModelOutput` instead of a plain tuple.
+
+        Examples:
+
+        ```python
+        >>> from transformers import PatchTSTConfig, PatchTSTForRegression
+
+        >>> # Regression task with 6 input channels and regress 2 targets
+        >>> model = PatchTSTForRegression.from_pretrained("namctin/patchtst_etth1_regression")
+
+        >>> # during inference, one only provides past values, the model outputs future values
+        >>> past_values = torch.randn(20, 512, 6)
+        >>> outputs = model(past_values=past_values)
+        >>> regression_outputs = outputs.regression_outputs
+        ```"""
+
+        return_dict = return_dict if return_dict is not None else self.config.return_dict
+
+        model_output = self.model(
+            past_values=past_values,
+            past_observed_mask=past_observed_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        # get output head. y_hat is of shape [bs x num_targets] or tuple of this shape
+        y_hat = self.head(model_output.last_hidden_state)
+
+        loss = None
+        if target_values is not None:
+            if self.distribution_output:
+                distribution = self.distribution_output.distribution(y_hat)
+                # y_hat should be a 2-tuple, each with dimension [bs, num_targets]
+                y_hat = tuple(item.view(-1, self.config.num_targets) for item in y_hat)
+                loss = nll(distribution, target_values)
+                # take average of the loss
+                loss = weighted_average(loss)
+            else:
+                loss = nn.MSELoss(reduction="mean")
+                loss = loss(y_hat, target_values)
+
+        if not return_dict:
+            # hidden_states, attentions, mask
+            outputs = (y_hat,) + model_output[1:-3]
+            outputs = (loss,) + outputs if loss is not None else outputs
+            return outputs
+        return PatchTSTForRegressionOutput(
+            loss=loss,
+            regression_outputs=y_hat,
+            hidden_states=model_output.hidden_states,
+            attentions=model_output.attentions,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        past_values: torch.Tensor,
+        past_observed_mask: torch.Tensor | None = None,
+    ) -> SamplePatchTSTOutput:
+        """
+        Generate sequences of sample predictions from a model with a probability distribution head.
+
+        Parameters:
+            past_values (`torch.FloatTensor` of shape `(batch_size, sequence_length, num_input_channels)`):
+                Past values of the time series that serves as context in order to predict the future.
+            past_observed_mask (`torch.BoolTensor` of shape `(batch_size, sequence_length, num_input_channels)`, *optional*):
+                Boolean mask to indicate which `past_values` were observed and which were missing. Mask values selected
+                in `[0, 1]`:
+
+                - 1 for values that are **observed**,
+                - 0 for values that are **missing** (i.e. NaNs that were replaced by zeros).
+
+        Return:
+            [`SamplePatchTSTOutput`] where the outputs `sequences` tensor will have shape `(batch_size, number of
+            samples, num_targets)`.
+        """
+        # get number of samples
+        num_parallel_samples = self.config.num_parallel_samples
+
+        # get model output
+        outputs = self(
+            past_values=past_values,
+            target_values=None,
+            past_observed_mask=past_observed_mask,
+            output_hidden_states=False,
+        )
+
+        # get distribution
+        distribution = self.distribution_output.distribution(outputs.regression_outputs)
+        # get samples: list of [bs x num_targets]
+        samples = [distribution.sample() for _ in range(num_parallel_samples)]
+        # samples: [bs x num_samples x num_targets]
+        samples = torch.stack(samples, dim=1).view(-1, num_parallel_samples, self.config.num_targets)
+        return SamplePatchTSTOutput(sequences=samples)
+
+
+__all__ = [
+    "PatchTSTModel",
+    "PatchTSTPreTrainedModel",
+    "PatchTSTForPrediction",
+    "PatchTSTForPretraining",
+    "PatchTSTForRegression",
+    "PatchTSTForClassification",
+]

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_001000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1aa341eb9ed99506627f80f4b70887e720d25512825a349a73cf28fe114cd46b
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_036000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5b2d72944bb49d0bcb8bbb5f01497fe800afeaf9cc5fbd3e27d0424c7d08e861
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_038000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ad049284a7e7301eb78966cc2be57487df86efe92e711a03e7fd48400173973
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_100000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87fdeaead6e09d7836628bbdd5a5bc626b6edc0b3c54efbd4b382a50c5cbe84f
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_135000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:73471d3765f03038acf5c0c134ecdeda95b7b67556b92c8dd8ed100dae8e890d
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_173000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28fbbcc3b528281bdca696e2e9fbd0018dc9b148a534da21d52258b4114136a1
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_197000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b2eb645b31651c7e54eb2ff7e06e3c5a8533f0fb976678fee816da1107eed6d4
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_203000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:565b196d1cf6ea7889eec1b25ed3c60a58e486b3b9b8b92a680bcf758b0eb158
+size 515519058

LTA_openwebtext_dualt/mini_owt_logdirichlet/runs/owt_t5_elftokenized_full_len1024_C1_to_1024_pow1_d768_l12_h12_gbs512_8gpu_50ep_lr3e4_elfopt_t5embed_unfixed_stateprobadd_selfcond_ce_fast_20260531_230026/step_236000.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a1c6f2dd3e478d9b8eedfc6e83464fc3b96fa7c58f0b4e013f9487237d42904
+size 515519058