dgenerate-ultralytics-headless 8.3.143__py3-none-any.whl → 8.3.145__py3-none-any.whl

ultralytics/models/yolo/yoloe/train.py

@@ -2,6 +2,8 @@
 
 import itertools
 from copy import copy, deepcopy
+from pathlib import Path
+from typing import Dict, List, Optional, Union
 
 import torch
 
@@ -17,9 +19,22 @@ from .val import YOLOEDetectValidator
 
 
 class YOLOETrainer(DetectionTrainer):
-    """A base trainer for YOLOE training."""
+    """
+    A trainer class for YOLOE object detection models.
 
-    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
+    This class extends DetectionTrainer to provide specialized training functionality for YOLOE models,
+    including custom model initialization, validation, and dataset building with multi-modal support.
+
+    Attributes:
+        loss_names (tuple): Names of loss components used during training.
+
+    Methods:
+        get_model: Initialize and return a YOLOEModel with specified configuration.
+        get_validator: Return a YOLOEDetectValidator for model validation.
+        build_dataset: Build YOLO dataset with multi-modal support for training.
+    """
+
+    def __init__(self, cfg=DEFAULT_CFG, overrides: Optional[Dict] = None, _callbacks=None):
         """
         Initialize the YOLOE Trainer with specified configurations.
 
@@ -36,14 +51,14 @@ class YOLOETrainer(DetectionTrainer):
         overrides["overlap_mask"] = False
         super().__init__(cfg, overrides, _callbacks)
 
-    def get_model(self, cfg=None, weights=None, verbose=True):
+    def get_model(self, cfg=None, weights=None, verbose: bool = True):
         """
         Return a YOLOEModel initialized with the specified configuration and weights.
 
         Args:
-            cfg (dict | str | None): Model configuration. Can be a dictionary containing a 'yaml_file' key,
+            cfg (dict | str, optional): Model configuration. Can be a dictionary containing a 'yaml_file' key,
                 a direct path to a YAML file, or None to use default configuration.
-            weights (str | Path | None): Path to pretrained weights file to load into the model.
+            weights (str | Path, optional): Path to pretrained weights file to load into the model.
             verbose (bool): Whether to display model information during initialization.
 
         Returns:
@@ -68,20 +83,20 @@ class YOLOETrainer(DetectionTrainer):
         return model
 
     def get_validator(self):
-        """Returns a DetectionValidator for YOLO model validation."""
+        """Return a YOLOEDetectValidator for YOLOE model validation."""
         self.loss_names = "box", "cls", "dfl"
         return YOLOEDetectValidator(
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
-    def build_dataset(self, img_path, mode="train", batch=None):
+    def build_dataset(self, img_path: str, mode: str = "train", batch: Optional[int] = None):
         """
         Build YOLO Dataset.
 
         Args:
             img_path (str): Path to the folder containing images.
-            mode (str): `train` mode or `val` mode, users are able to customize different augmentations for each mode.
-            batch (int, optional): Size of batches, this is for `rect`.
+            mode (str): 'train' mode or 'val' mode, users are able to customize different augmentations for each mode.
+            batch (int, optional): Size of batches, this is for rectangular training.
 
         Returns:
             (Dataset): YOLO dataset configured for training or validation.
@@ -93,9 +108,17 @@ class YOLOETrainer(DetectionTrainer):
 
 
 class YOLOEPETrainer(DetectionTrainer):
-    """Fine-tune YOLOE model in linear probing way."""
+    """
+    Fine-tune YOLOE model using linear probing approach.
 
-    def get_model(self, cfg=None, weights=None, verbose=True):
+    This trainer freezes most model layers and only trains specific projection layers for efficient
+    fine-tuning on new datasets while preserving pretrained features.
+
+    Methods:
+        get_model: Initialize YOLOEModel with frozen layers except projection layers.
+    """
+
+    def get_model(self, cfg=None, weights=None, verbose: bool = True):
         """
         Return YOLOEModel initialized with specified config and weights.
 
@@ -139,9 +162,19 @@ class YOLOEPETrainer(DetectionTrainer):
 
 
 class YOLOETrainerFromScratch(YOLOETrainer, WorldTrainerFromScratch):
-    """Train YOLOE models from scratch."""
+    """
+    Train YOLOE models from scratch with text embedding support.
+
+    This trainer combines YOLOE training capabilities with world training features, enabling
+    training from scratch with text embeddings and grounding datasets.
 
-    def build_dataset(self, img_path, mode="train", batch=None):
+    Methods:
+        build_dataset: Build datasets for training with grounding support.
+        preprocess_batch: Process batches with text features.
+        generate_text_embeddings: Generate and cache text embeddings for training.
+    """
+
+    def build_dataset(self, img_path: Union[List[str], str], mode: str = "train", batch: Optional[int] = None):
         """
         Build YOLO Dataset for training or validation.
 
@@ -168,7 +201,7 @@ class YOLOETrainerFromScratch(YOLOETrainer, WorldTrainerFromScratch):
         batch["txt_feats"] = txt_feats
         return batch
 
-    def generate_text_embeddings(self, texts, batch, cache_dir):
+    def generate_text_embeddings(self, texts: List[str], batch: int, cache_dir: Path):
         """
         Generate text embeddings for a list of text samples.
 
@@ -196,21 +229,31 @@ class YOLOETrainerFromScratch(YOLOETrainer, WorldTrainerFromScratch):
 
 
 class YOLOEPEFreeTrainer(YOLOEPETrainer, YOLOETrainerFromScratch):
-    """Train prompt-free YOLOE model."""
+    """
+    Train prompt-free YOLOE model.
+
+    This trainer combines linear probing capabilities with from-scratch training for prompt-free
+    YOLOE models that don't require text prompts during inference.
+
+    Methods:
+        get_validator: Return standard DetectionValidator for validation.
+        preprocess_batch: Preprocess batches without text features.
+        set_text_embeddings: Set text embeddings for datasets (no-op for prompt-free).
+    """
 
     def get_validator(self):
-        """Returns a DetectionValidator for YOLO model validation."""
+        """Return a DetectionValidator for YOLO model validation."""
         self.loss_names = "box", "cls", "dfl"
         return DetectionValidator(
             self.test_loader, save_dir=self.save_dir, args=copy(self.args), _callbacks=self.callbacks
         )
 
     def preprocess_batch(self, batch):
-        """Preprocesses a batch of images for YOLOE training, adjusting formatting and dimensions as needed."""
+        """Preprocess a batch of images for YOLOE training, adjusting formatting and dimensions as needed."""
         batch = DetectionTrainer.preprocess_batch(self, batch)
         return batch
 
-    def set_text_embeddings(self, datasets, batch):
+    def set_text_embeddings(self, datasets, batch: int):
         """
         Set text embeddings for datasets to accelerate training by caching category names.
 
@@ -231,9 +274,18 @@ class YOLOEPEFreeTrainer(YOLOEPETrainer, YOLOETrainerFromScratch):
 
 
 class YOLOEVPTrainer(YOLOETrainerFromScratch):
-    """Train YOLOE model with visual prompts."""
+    """
+    Train YOLOE model with visual prompts.
+
+    This trainer extends YOLOETrainerFromScratch to support visual prompt-based training,
+    where visual cues are provided alongside images to guide the detection process.
+
+    Methods:
+        build_dataset: Build dataset with visual prompt loading transforms.
+        preprocess_batch: Preprocess batches with visual prompts.
+    """
 
-    def build_dataset(self, img_path, mode="train", batch=None):
+    def build_dataset(self, img_path: Union[List[str], str], mode: str = "train", batch: Optional[int] = None):
         """
         Build YOLO Dataset for training or validation with visual prompts.
 
@@ -263,7 +315,7 @@ class YOLOEVPTrainer(YOLOETrainerFromScratch):
             self.train_loader.dataset.transforms.append(LoadVisualPrompt())
 
     def preprocess_batch(self, batch):
-        """Preprocesses a batch of images for YOLOE training, moving visual prompts to the appropriate device."""
+        """Preprocess a batch of images for YOLOE training, moving visual prompts to the appropriate device."""
         batch = super().preprocess_batch(batch)
         batch["visuals"] = batch["visuals"].to(self.device)
         return batch

ultralytics/models/yolo/yoloe/train_seg.py

@@ -14,8 +14,8 @@ class YOLOESegTrainer(YOLOETrainer, SegmentationTrainer):
     """
     Trainer class for YOLOE segmentation models.
 
-    This class combines YOLOETrainer and SegmentationTrainer to provide training functionality
-    specifically for YOLOE segmentation models.
+    This class combines YOLOETrainer and SegmentationTrainer to provide training functionality specifically for YOLOE
+    segmentation models, enabling both object detection and instance segmentation capabilities.
 
     Attributes:
         cfg (dict): Configuration dictionary with training parameters.
@@ -28,7 +28,7 @@ class YOLOESegTrainer(YOLOETrainer, SegmentationTrainer):
         Return YOLOESegModel initialized with specified config and weights.
 
         Args:
-            cfg (dict | str): Model configuration dictionary or YAML file path.
+            cfg (dict | str, optional): Model configuration dictionary or YAML file path.
             weights (str, optional): Path to pretrained weights file.
             verbose (bool): Whether to display model information.
 
@@ -66,7 +66,10 @@ class YOLOEPESegTrainer(SegmentationTrainer):
     Fine-tune YOLOESeg model in linear probing way.
 
     This trainer specializes in fine-tuning YOLOESeg models using a linear probing approach, which involves freezing
-    most of the model and only training specific layers.
+    most of the model and only training specific layers for efficient adaptation to new tasks.
+
+    Attributes:
+        data (dict): Dataset configuration containing channels, class names, and number of classes.
     """
 
     def get_model(self, cfg=None, weights=None, verbose=True):
@@ -74,7 +77,7 @@ class YOLOEPESegTrainer(SegmentationTrainer):
         Return YOLOESegModel initialized with specified config and weights for linear probing.
 
         Args:
-            cfg (dict | str): Model configuration dictionary or YAML file path.
+            cfg (dict | str, optional): Model configuration dictionary or YAML file path.
             weights (str, optional): Path to pretrained weights file.
             verbose (bool): Whether to display model information.
 
@@ -113,12 +116,12 @@ class YOLOEPESegTrainer(SegmentationTrainer):
 
 
 class YOLOESegTrainerFromScratch(YOLOETrainerFromScratch, YOLOESegTrainer):
-    """Trainer for YOLOE segmentation from scratch."""
+    """Trainer for YOLOE segmentation models trained from scratch without pretrained weights."""
 
     pass
 
 
 class YOLOESegVPTrainer(YOLOEVPTrainer, YOLOESegTrainerFromScratch):
-    """Trainer for YOLOE segmentation with VP."""
+    """Trainer for YOLOE segmentation models with Vision Prompt (VP) capabilities."""
 
     pass

ultralytics/models/yolo/yoloe/val.py

@@ -1,6 +1,7 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 
 from copy import deepcopy
+from typing import Any, Dict, Optional, Union
 
 import torch
 from torch.nn import functional as F
@@ -18,26 +19,40 @@ from ultralytics.utils.torch_utils import select_device, smart_inference_mode
 
 class YOLOEDetectValidator(DetectionValidator):
     """
-    A mixin class for YOLOE model validation that handles both text and visual prompt embeddings.
+    A validator class for YOLOE detection models that handles both text and visual prompt embeddings.
 
-    This mixin provides functionality to validate YOLOE models using either text or visual prompt embeddings.
-    It includes methods for extracting visual prompt embeddings from samples, preprocessing batches, and
-    running validation with different prompt types.
+    This class extends DetectionValidator to provide specialized validation functionality for YOLOE models.
+    It supports validation using either text prompts or visual prompt embeddings extracted from training samples,
+    enabling flexible evaluation strategies for prompt-based object detection.
 
     Attributes:
         device (torch.device): The device on which validation is performed.
         args (namespace): Configuration arguments for validation.
         dataloader (DataLoader): DataLoader for validation data.
+
+    Methods:
+        get_visual_pe: Extract visual prompt embeddings from training samples.
+        preprocess: Preprocess batch data ensuring visuals are on the same device as images.
+        get_vpe_dataloader: Create a dataloader for LVIS training visual prompt samples.
+        __call__: Run validation using either text or visual prompt embeddings.
+
+    Examples:
+        Validate with text prompts
+        >>> validator = YOLOEDetectValidator()
+        >>> stats = validator(model=model, load_vp=False)
+
+        Validate with visual prompts
+        >>> stats = validator(model=model, refer_data="path/to/data.yaml", load_vp=True)
     """
 
     @smart_inference_mode()
-    def get_visual_pe(self, dataloader, model):
+    def get_visual_pe(self, dataloader: torch.utils.data.DataLoader, model: YOLOEModel) -> torch.Tensor:
         """
         Extract visual prompt embeddings from training samples.
 
-        This function processes a dataloader to compute visual prompt embeddings for each class
-        using a YOLOE model. It normalizes the embeddings and handles cases where no samples
-        exist for a class.
+        This method processes a dataloader to compute visual prompt embeddings for each class using a YOLOE model.
+        It normalizes the embeddings and handles cases where no samples exist for a class by setting their
+        embeddings to zero.
 
         Args:
             dataloader (torch.utils.data.DataLoader): The dataloader providing training samples.
@@ -53,6 +68,7 @@ class YOLOEDetectValidator(DetectionValidator):
 
         desc = "Get visual prompt embeddings from samples"
 
+        # Count samples per class
         for batch in dataloader:
             cls = batch["cls"].squeeze(-1).to(torch.int).unique()
             count = torch.bincount(cls, minlength=len(names))
@@ -60,6 +76,7 @@ class YOLOEDetectValidator(DetectionValidator):
 
         cls_visual_num = cls_visual_num.to(self.device)
 
+        # Extract visual prompt embeddings
         pbar = TQDM(dataloader, total=len(dataloader), desc=desc)
         for batch in pbar:
             batch = self.preprocess(batch)
@@ -73,30 +90,31 @@ class YOLOEDetectValidator(DetectionValidator):
                 for c in cls:
                     visual_pe[c] += preds[i][pad_cls == c].sum(0) / cls_visual_num[c]
 
+        # Normalize embeddings for classes with samples, set others to zero
         visual_pe[cls_visual_num != 0] = F.normalize(visual_pe[cls_visual_num != 0], dim=-1, p=2)
         visual_pe[cls_visual_num == 0] = 0
         return visual_pe.unsqueeze(0)
 
-    def preprocess(self, batch):
+    def preprocess(self, batch: Dict[str, Any]) -> Dict[str, Any]:
         """Preprocess batch data, ensuring visuals are on the same device as images."""
         batch = super().preprocess(batch)
         if "visuals" in batch:
             batch["visuals"] = batch["visuals"].to(batch["img"].device)
         return batch
 
-    def get_vpe_dataloader(self, data):
+    def get_vpe_dataloader(self, data: Dict[str, Any]) -> torch.utils.data.DataLoader:
         """
         Create a dataloader for LVIS training visual prompt samples.
 
-        This function prepares a dataloader for visual prompt embeddings (VPE) using the LVIS dataset.
-        It applies necessary transformations and configurations to the dataset and returns a dataloader
+        This method prepares a dataloader for visual prompt embeddings (VPE) using the specified dataset.
+        It applies necessary transformations including LoadVisualPrompt and configurations to the dataset
         for validation purposes.
 
         Args:
             data (dict): Dataset configuration dictionary containing paths and settings.
 
         Returns:
-            (torch.utils.data.DataLoader): The dataLoader for visual prompt samples.
+            (torch.utils.data.DataLoader): The dataloader for visual prompt samples.
         """
         dataset = build_yolo_dataset(
             self.args,
@@ -120,17 +138,23 @@ class YOLOEDetectValidator(DetectionValidator):
         )
 
     @smart_inference_mode()
-    def __call__(self, trainer=None, model=None, refer_data=None, load_vp=False):
+    def __call__(
+        self,
+        trainer: Optional[Any] = None,
+        model: Optional[Union[YOLOEModel, str]] = None,
+        refer_data: Optional[str] = None,
+        load_vp: bool = False,
+    ) -> Dict[str, Any]:
         """
         Run validation on the model using either text or visual prompt embeddings.
 
-        This method validates the model using either text prompts or visual prompts, depending
-        on the `load_vp` flag. It supports validation during training (using a trainer object)
-        or standalone validation with a provided model.
+        This method validates the model using either text prompts or visual prompts, depending on the load_vp flag.
+        It supports validation during training (using a trainer object) or standalone validation with a provided
+        model. For visual prompts, reference data can be specified to extract embeddings from a different dataset.
 
         Args:
             trainer (object, optional): Trainer object containing the model and device.
-            model (YOLOEModel, optional): Model to validate. Required if `trainer` is not provided.
+            model (YOLOEModel | str, optional): Model to validate. Required if trainer is not provided.
             refer_data (str, optional): Path to reference data for visual prompts.
             load_vp (bool): Whether to load visual prompts. If False, text prompts are used.
 

ultralytics/nn/autobackend.py

@@ -6,7 +6,7 @@ import platform
 import zipfile
 from collections import OrderedDict, namedtuple
 from pathlib import Path
-from typing import List, Optional, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import cv2
 import numpy as np
@@ -19,8 +19,19 @@ from ultralytics.utils.checks import check_requirements, check_suffix, check_ver
 from ultralytics.utils.downloads import attempt_download_asset, is_url
 
 
-def check_class_names(names):
-    """Check class names and convert to dict format if needed."""
+def check_class_names(names: Union[List, Dict]) -> Dict[int, str]:
+    """
+    Check class names and convert to dict format if needed.
+
+    Args:
+        names (list | dict): Class names as list or dict format.
+
+    Returns:
+        (dict): Class names in dict format with integer keys and string values.
+
+    Raises:
+        KeyError: If class indices are invalid for the dataset size.
+    """
     if isinstance(names, list):  # names is a list
         names = dict(enumerate(names))  # convert to dict
     if isinstance(names, dict):
@@ -38,8 +49,16 @@ def check_class_names(names):
     return names
 
 
-def default_class_names(data=None):
-    """Applies default class names to an input YAML file or returns numerical class names."""
+def default_class_names(data: Optional[Union[str, Path]] = None) -> Dict[int, str]:
+    """
+    Apply default class names to an input YAML file or return numerical class names.
+
+    Args:
+        data (str | Path, optional): Path to YAML file containing class names.
+
+    Returns:
+        (dict): Dictionary mapping class indices to class names.
+    """
     if data:
         try:
             return YAML.load(check_yaml(data))["names"]
@@ -50,7 +69,7 @@ def default_class_names(data=None):
 
 class AutoBackend(nn.Module):
     """
-    Handles dynamic backend selection for running inference using Ultralytics YOLO models.
+    Handle dynamic backend selection for running inference using Ultralytics YOLO models.
 
     The AutoBackend class is designed to provide an abstraction layer for various inference engines. It supports a wide
     range of formats, each with specific naming conventions as outlined below:
@@ -82,6 +101,24 @@ class AutoBackend(nn.Module):
         names (dict): A dictionary of class names that the model can detect.
         stride (int): The model stride, typically 32 for YOLO models.
         fp16 (bool): Whether the model uses half-precision (FP16) inference.
+        nhwc (bool): Whether the model expects NHWC input format instead of NCHW.
+        pt (bool): Whether the model is a PyTorch model.
+        jit (bool): Whether the model is a TorchScript model.
+        onnx (bool): Whether the model is an ONNX model.
+        xml (bool): Whether the model is an OpenVINO model.
+        engine (bool): Whether the model is a TensorRT engine.
+        coreml (bool): Whether the model is a CoreML model.
+        saved_model (bool): Whether the model is a TensorFlow SavedModel.
+        pb (bool): Whether the model is a TensorFlow GraphDef.
+        tflite (bool): Whether the model is a TensorFlow Lite model.
+        edgetpu (bool): Whether the model is a TensorFlow Edge TPU model.
+        tfjs (bool): Whether the model is a TensorFlow.js model.
+        paddle (bool): Whether the model is a PaddlePaddle model.
+        mnn (bool): Whether the model is an MNN model.
+        ncnn (bool): Whether the model is an NCNN model.
+        imx (bool): Whether the model is an IMX model.
+        rknn (bool): Whether the model is an RKNN model.
+        triton (bool): Whether the model is a Triton Inference Server model.
 
     Methods:
         forward: Run inference on an input image.
@@ -113,7 +150,7 @@ class AutoBackend(nn.Module):
             weights (str | List[str] | torch.nn.Module): Path to the model weights file or a module instance.
             device (torch.device): Device to run the model on.
             dnn (bool): Use OpenCV DNN module for ONNX inference.
-            data (str | Path | optional): Path to the additional data.yaml file containing class names.
+            data (str | Path, optional): Path to the additional data.yaml file containing class names.
             fp16 (bool): Enable half-precision inference. Supported only on specific backends.
             batch (int): Batch-size to assume for inference.
             fuse (bool): Fuse Conv2D + BatchNorm layers for optimization.
@@ -567,15 +604,22 @@ class AutoBackend(nn.Module):
 
         self.__dict__.update(locals())  # assign all variables to self
 
-    def forward(self, im, augment=False, visualize=False, embed=None, **kwargs):
+    def forward(
+        self,
+        im: torch.Tensor,
+        augment: bool = False,
+        visualize: bool = False,
+        embed: Optional[List] = None,
+        **kwargs: Any,
+    ) -> Union[torch.Tensor, List[torch.Tensor]]:
         """
-        Runs inference on the YOLOv8 MultiBackend model.
+        Run inference on an AutoBackend model.
 
         Args:
             im (torch.Tensor): The image tensor to perform inference on.
             augment (bool): Whether to perform data augmentation during inference.
             visualize (bool): Whether to visualize the output predictions.
-            embed (list | None): A list of feature vectors/embeddings to return.
+            embed (list, optional): A list of feature vectors/embeddings to return.
             **kwargs (Any): Additional keyword arguments for model configuration.
 
         Returns:
@@ -632,7 +676,7 @@ class AutoBackend(nn.Module):
                 results = [None] * n  # preallocate list with None to match the number of images
 
                 def callback(request, userdata):
-                    """Places result in preallocated list using userdata index."""
+                    """Place result in preallocated list using userdata index."""
                     results[userdata] = request.results
 
                 # Create AsyncInferQueue, set the callback and start asynchronous inference for each input image
@@ -780,7 +824,7 @@ class AutoBackend(nn.Module):
         else:
             return self.from_numpy(y)
 
-    def from_numpy(self, x):
+    def from_numpy(self, x: np.ndarray) -> torch.Tensor:
         """
         Convert a numpy array to a tensor.
 
@@ -792,7 +836,7 @@ class AutoBackend(nn.Module):
         """
         return torch.tensor(x).to(self.device) if isinstance(x, np.ndarray) else x
 
-    def warmup(self, imgsz=(1, 3, 640, 640)):
+    def warmup(self, imgsz: Tuple[int, int, int, int] = (1, 3, 640, 640)) -> None:
         """
         Warm up the model by running one forward pass with a dummy input.
 
@@ -808,9 +852,9 @@ class AutoBackend(nn.Module):
                 self.forward(im)  # warmup
 
     @staticmethod
-    def _model_type(p="path/to/model.pt"):
+    def _model_type(p: str = "path/to/model.pt") -> List[bool]:
         """
-        Takes a path to a model file and returns the model type.
+        Take a path to a model file and return the model type.
 
         Args:
             p (str): Path to the model file.

ultralytics/nn/modules/__init__.py

@@ -1,12 +1,12 @@
 # Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
 """
-Ultralytics modules.
+Ultralytics neural network modules.
 
-This module provides access to various neural network components used in Ultralytics models, including convolution blocks,
-attention mechanisms, transformer components, and detection/segmentation heads.
+This module provides access to various neural network components used in Ultralytics models, including convolution
+blocks, attention mechanisms, transformer components, and detection/segmentation heads.
 
 Examples:
-    Visualize a module with Netron.
+    Visualize a module with Netron
     >>> from ultralytics.nn.modules import *
     >>> import torch
     >>> import os

ultralytics/nn/modules/activation.py

@@ -10,7 +10,7 @@ class AGLU(nn.Module):
     Unified activation function module from AGLU.
 
     This class implements a parameterized activation function with learnable parameters lambda and kappa, based on the
-    AGLU (Adaptive Gated Linear Unit) approach (https://github.com/kostas1515/AGLU).
+    AGLU (Adaptive Gated Linear Unit) approach.
 
     Attributes:
         act (nn.Softplus): Softplus activation function with negative beta.
@@ -27,6 +27,9 @@ class AGLU(nn.Module):
         >>> output = m(input)
         >>> print(output.shape)
         torch.Size([2])
+
+    References:
+        https://github.com/kostas1515/AGLU
     """
 
     def __init__(self, device=None, dtype=None) -> None: