dgenerate-ultralytics-headless 8.3.214__py3-none-any.whl → 8.3.248__py3-none-any.whl

ultralytics/nn/autobackend.py

@@ -19,11 +19,11 @@ from PIL import Image
 from ultralytics.utils import ARM64, IS_JETSON, LINUX, LOGGER, PYTHON_VERSION, ROOT, YAML, is_jetson
 from ultralytics.utils.checks import check_requirements, check_suffix, check_version, check_yaml, is_rockchip
 from ultralytics.utils.downloads import attempt_download_asset, is_url
+from ultralytics.utils.nms import non_max_suppression
 
 
 def check_class_names(names: list | dict) -> dict[int, str]:
-    """
-    Check class names and convert to dict format if needed.
+    """Check class names and convert to dict format if needed.
 
     Args:
         names (list | dict): Class names as list or dict format.
@@ -52,8 +52,7 @@ def check_class_names(names: list | dict) -> dict[int, str]:
 
 
 def default_class_names(data: str | Path | None = None) -> dict[int, str]:
-    """
-    Apply default class names to an input YAML file or return numerical class names.
+    """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.
@@ -70,8 +69,7 @@ def default_class_names(data: str | Path | None = None) -> dict[int, str]:
 
 
 class AutoBackend(nn.Module):
-    """
-    Handle 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:
@@ -95,6 +93,9 @@ class AutoBackend(nn.Module):
             | NCNN                  | *_ncnn_model/     |
             | IMX                   | *_imx_model/      |
             | RKNN                  | *_rknn_model/     |
+            | Triton Inference      | triton://model    |
+            | ExecuTorch            | *.pte             |
+            | Axelera               | *_axelera_model/  |
 
     Attributes:
         model (torch.nn.Module): The loaded YOLO model.
@@ -121,10 +122,12 @@ class AutoBackend(nn.Module):
         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.
+        pte (bool): Whether the model is a PyTorch ExecuTorch model.
+        axelera (bool): Whether the model is an Axelera model.
 
     Methods:
         forward: Run inference on an input image.
-        from_numpy: Convert numpy array to tensor.
+        from_numpy: Convert NumPy arrays to tensors on the model device.
         warmup: Warm up the model with a dummy input.
         _model_type: Determine the model type from file path.
 
@@ -144,8 +147,7 @@ class AutoBackend(nn.Module):
         fuse: bool = True,
         verbose: bool = True,
     ):
-        """
-        Initialize the AutoBackend for inference.
+        """Initialize the AutoBackend for inference.
 
         Args:
             model (str | torch.nn.Module): Path to the model weights file or a module instance.
@@ -175,10 +177,12 @@ class AutoBackend(nn.Module):
             ncnn,
             imx,
             rknn,
+            pte,
+            axelera,
             triton,
         ) = self._model_type("" if nn_module else model)
         fp16 &= pt or jit or onnx or xml or engine or nn_module or triton  # FP16
-        nhwc = coreml or saved_model or pb or tflite or edgetpu or rknn  # BHWC formats (vs torch BCWH)
+        nhwc = coreml or saved_model or pb or tflite or edgetpu or rknn  # BHWC formats (vs torch BCHW)
         stride, ch = 32, 3  # default stride and channels
         end2end, dynamic = False, False
         metadata, task = None, None
@@ -241,25 +245,28 @@ class AutoBackend(nn.Module):
             check_requirements(("onnx", "onnxruntime-gpu" if cuda else "onnxruntime"))
             import onnxruntime
 
-            providers = ["CPUExecutionProvider"]
-            if cuda:
-                if "CUDAExecutionProvider" in onnxruntime.get_available_providers():
-                    providers.insert(0, "CUDAExecutionProvider")
-                else:  # Only log warning if CUDA was requested but unavailable
-                    LOGGER.warning("Failed to start ONNX Runtime with CUDA. Using CPU...")
-                    device = torch.device("cpu")
-                    cuda = False
-            LOGGER.info(f"Using ONNX Runtime {onnxruntime.__version__} {providers[0]}")
+            # Select execution provider: CUDA > CoreML (mps) > CPU
+            available = onnxruntime.get_available_providers()
+            if cuda and "CUDAExecutionProvider" in available:
+                providers = [("CUDAExecutionProvider", {"device_id": device.index}), "CPUExecutionProvider"]
+            elif device.type == "mps" and "CoreMLExecutionProvider" in available:
+                providers = ["CoreMLExecutionProvider", "CPUExecutionProvider"]
+            else:
+                providers = ["CPUExecutionProvider"]
+                if cuda:
+                    LOGGER.warning("CUDA requested but CUDAExecutionProvider not available. Using CPU...")
+                    device, cuda = torch.device("cpu"), False
+            LOGGER.info(
+                f"Using ONNX Runtime {onnxruntime.__version__} with {providers[0] if isinstance(providers[0], str) else providers[0][0]}"
+            )
             if onnx:
                 session = onnxruntime.InferenceSession(w, providers=providers)
             else:
-                check_requirements(
-                    ("model-compression-toolkit>=2.4.1", "sony-custom-layers[torch]>=0.3.0", "onnxruntime-extensions")
-                )
+                check_requirements(("model-compression-toolkit>=2.4.1", "edge-mdt-cl<1.1.0", "onnxruntime-extensions"))
                 w = next(Path(w).glob("*.onnx"))
                 LOGGER.info(f"Loading {w} for ONNX IMX inference...")
                 import mct_quantizers as mctq
-                from sony_custom_layers.pytorch.nms import nms_ort  # noqa
+                from edgemdt_cl.pytorch.nms import nms_ort  # noqa - register custom NMS ops
 
                 session_options = mctq.get_ort_session_options()
                 session_options.enable_mem_reuse = False  # fix the shape mismatch from onnxruntime
@@ -269,7 +276,10 @@ class AutoBackend(nn.Module):
             metadata = session.get_modelmeta().custom_metadata_map
             dynamic = isinstance(session.get_outputs()[0].shape[0], str)
             fp16 = "float16" in session.get_inputs()[0].type
-            if not dynamic:
+
+            # Setup IO binding for optimized inference (CUDA only, not supported for CoreML)
+            use_io_binding = not dynamic and cuda
+            if use_io_binding:
                 io = session.io_binding()
                 bindings = []
                 for output in session.get_outputs():
@@ -332,11 +342,11 @@ class AutoBackend(nn.Module):
                 check_requirements("numpy==1.23.5")
 
             try:  # https://developer.nvidia.com/nvidia-tensorrt-download
-                import tensorrt as trt  # noqa
+                import tensorrt as trt
             except ImportError:
                 if LINUX:
                     check_requirements("tensorrt>7.0.0,!=10.1.0")
-                import tensorrt as trt  # noqa
+                import tensorrt as trt
             check_version(trt.__version__, ">=7.0.0", hard=True)
             check_version(trt.__version__, "!=10.1.0", msg="https://github.com/ultralytics/ultralytics/pull/14239")
             if device.type == "cpu":
@@ -369,39 +379,42 @@ class AutoBackend(nn.Module):
             is_trt10 = not hasattr(model, "num_bindings")
             num = range(model.num_io_tensors) if is_trt10 else range(model.num_bindings)
             for i in num:
+                # Get tensor info using TRT10+ or legacy API
                 if is_trt10:
                     name = model.get_tensor_name(i)
                     dtype = trt.nptype(model.get_tensor_dtype(name))
                     is_input = model.get_tensor_mode(name) == trt.TensorIOMode.INPUT
-                    if is_input:
-                        if -1 in tuple(model.get_tensor_shape(name)):
-                            dynamic = True
-                            context.set_input_shape(name, tuple(model.get_tensor_profile_shape(name, 0)[1]))
-                        if dtype == np.float16:
-                            fp16 = True
-                    else:
-                        output_names.append(name)
-                    shape = tuple(context.get_tensor_shape(name))
-                else:  # TensorRT < 10.0
+                    shape = tuple(model.get_tensor_shape(name))
+                    profile_shape = tuple(model.get_tensor_profile_shape(name, 0)[2]) if is_input else None
+                else:
                     name = model.get_binding_name(i)
                     dtype = trt.nptype(model.get_binding_dtype(i))
                     is_input = model.binding_is_input(i)
-                    if model.binding_is_input(i):
-                        if -1 in tuple(model.get_binding_shape(i)):  # dynamic
-                            dynamic = True
-                            context.set_binding_shape(i, tuple(model.get_profile_shape(0, i)[1]))
-                        if dtype == np.float16:
-                            fp16 = True
-                    else:
-                        output_names.append(name)
-                    shape = tuple(context.get_binding_shape(i))
+                    shape = tuple(model.get_binding_shape(i))
+                    profile_shape = tuple(model.get_profile_shape(0, i)[1]) if is_input else None
+
+                # Process input/output tensors
+                if is_input:
+                    if -1 in shape:
+                        dynamic = True
+                        if is_trt10:
+                            context.set_input_shape(name, profile_shape)
+                        else:
+                            context.set_binding_shape(i, profile_shape)
+                    if dtype == np.float16:
+                        fp16 = True
+                else:
+                    output_names.append(name)
+                shape = tuple(context.get_tensor_shape(name)) if is_trt10 else tuple(context.get_binding_shape(i))
                 im = torch.from_numpy(np.empty(shape, dtype=dtype)).to(device)
                 bindings[name] = Binding(name, dtype, shape, im, int(im.data_ptr()))
             binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
 
         # CoreML
         elif coreml:
-            check_requirements("coremltools>=8.0")
+            check_requirements(
+                ["coremltools>=9.0", "numpy>=1.14.5,<=2.3.5"]
+            )  # latest numpy 2.4.0rc1 breaks coremltools exports
             LOGGER.info(f"Loading {w} for CoreML inference...")
             import coremltools as ct
 
@@ -414,8 +427,7 @@ class AutoBackend(nn.Module):
             LOGGER.info(f"Loading {w} for TensorFlow SavedModel inference...")
             import tensorflow as tf
 
-            keras = False  # assume TF1 saved_model
-            model = tf.keras.models.load_model(w) if keras else tf.saved_model.load(w)
+            model = tf.saved_model.load(w)
             metadata = Path(w) / "metadata.yaml"
 
         # TF GraphDef
@@ -423,7 +435,7 @@ class AutoBackend(nn.Module):
             LOGGER.info(f"Loading {w} for TensorFlow GraphDef inference...")
             import tensorflow as tf
 
-            from ultralytics.engine.exporter import gd_outputs
+            from ultralytics.utils.export.tensorflow import gd_outputs
 
             def wrap_frozen_graph(gd, inputs, outputs):
                 """Wrap frozen graphs for deployment."""
@@ -491,7 +503,7 @@ class AutoBackend(nn.Module):
                 if ARM64
                 else "paddlepaddle>=3.0.0"
             )
-            import paddle.inference as pdi  # noqa
+            import paddle.inference as pdi
 
             w = Path(w)
             model_file, params_file = None, None
@@ -569,6 +581,51 @@ class AutoBackend(nn.Module):
             rknn_model.init_runtime()
             metadata = w.parent / "metadata.yaml"
 
+        # Axelera
+        elif axelera:
+            import os
+
+            if not os.environ.get("AXELERA_RUNTIME_DIR"):
+                LOGGER.warning(
+                    "Axelera runtime environment is not activated."
+                    "\nPlease run: source /opt/axelera/sdk/latest/axelera_activate.sh"
+                    "\n\nIf this fails, verify driver installation: https://docs.ultralytics.com/integrations/axelera/#axelera-driver-installation"
+                )
+            try:
+                from axelera.runtime import op
+            except ImportError:
+                check_requirements(
+                    "axelera_runtime2==0.1.2",
+                    cmds="--extra-index-url https://software.axelera.ai/artifactory/axelera-runtime-pypi",
+                )
+            from axelera.runtime import op
+
+            w = Path(w)
+            if (found := next(w.rglob("*.axm"), None)) is None:
+                raise FileNotFoundError(f"No .axm file found in: {w}")
+
+            ax_model = op.load(str(found))
+            metadata = found.parent / "metadata.yaml"
+
+        # ExecuTorch
+        elif pte:
+            LOGGER.info(f"Loading {w} for ExecuTorch inference...")
+            # TorchAO release compatibility table bug https://github.com/pytorch/ao/issues/2919
+            check_requirements("setuptools<71.0.0")  # Setuptools bug: https://github.com/pypa/setuptools/issues/4483
+            check_requirements(("executorch==1.0.1", "flatbuffers"))
+            from executorch.runtime import Runtime
+
+            w = Path(w)
+            if w.is_dir():
+                model_file = next(w.rglob("*.pte"))
+                metadata = w / "metadata.yaml"
+            else:
+                model_file = w
+                metadata = w.parent / "metadata.yaml"
+
+            program = Runtime.get().load_program(str(model_file))
+            model = program.load_method("forward")
+
         # Any other format (unsupported)
         else:
             from ultralytics.engine.exporter import export_formats
@@ -585,14 +642,15 @@ class AutoBackend(nn.Module):
             for k, v in metadata.items():
                 if k in {"stride", "batch", "channels"}:
                     metadata[k] = int(v)
-                elif k in {"imgsz", "names", "kpt_shape", "args"} and isinstance(v, str):
-                    metadata[k] = eval(v)
+                elif k in {"imgsz", "names", "kpt_shape", "kpt_names", "args"} and isinstance(v, str):
+                    metadata[k] = ast.literal_eval(v)
             stride = metadata["stride"]
             task = metadata["task"]
             batch = metadata["batch"]
             imgsz = metadata["imgsz"]
             names = metadata["names"]
             kpt_shape = metadata.get("kpt_shape")
+            kpt_names = metadata.get("kpt_names")
             end2end = metadata.get("args", {}).get("nms", False)
             dynamic = metadata.get("args", {}).get("dynamic", dynamic)
             ch = metadata.get("channels", 3)
@@ -614,8 +672,7 @@ class AutoBackend(nn.Module):
         embed: list | None = None,
         **kwargs: Any,
     ) -> torch.Tensor | list[torch.Tensor]:
-        """
-        Run inference on an AutoBackend model.
+        """Run inference on an AutoBackend model.
 
         Args:
             im (torch.Tensor): The image tensor to perform inference on.
@@ -627,7 +684,7 @@ class AutoBackend(nn.Module):
         Returns:
             (torch.Tensor | list[torch.Tensor]): The raw output tensor(s) from the model.
         """
-        b, ch, h, w = im.shape  # batch, channel, height, width
+        _b, _ch, h, w = im.shape  # batch, channel, height, width
         if self.fp16 and im.dtype != torch.float16:
             im = im.half()  # to FP16
         if self.nhwc:
@@ -649,10 +706,7 @@ class AutoBackend(nn.Module):
 
         # ONNX Runtime
         elif self.onnx or self.imx:
-            if self.dynamic:
-                im = im.cpu().numpy()  # torch to numpy
-                y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})
-            else:
+            if self.use_io_binding:
                 if not self.cuda:
                     im = im.cpu()
                 self.io.bind_input(
@@ -665,13 +719,21 @@ class AutoBackend(nn.Module):
                 )
                 self.session.run_with_iobinding(self.io)
                 y = self.bindings
+            else:
+                im = im.cpu().numpy()  # torch to numpy
+                y = self.session.run(self.output_names, {self.session.get_inputs()[0].name: im})
             if self.imx:
                 if self.task == "detect":
                     # boxes, conf, cls
                     y = np.concatenate([y[0], y[1][:, :, None], y[2][:, :, None]], axis=-1)
                 elif self.task == "pose":
                     # boxes, conf, kpts
-                    y = np.concatenate([y[0], y[1][:, :, None], y[2][:, :, None], y[3]], axis=-1)
+                    y = np.concatenate([y[0], y[1][:, :, None], y[2][:, :, None], y[3]], axis=-1, dtype=y[0].dtype)
+                elif self.task == "segment":
+                    y = (
+                        np.concatenate([y[0], y[1][:, :, None], y[2][:, :, None], y[3]], axis=-1, dtype=y[0].dtype),
+                        y[4],
+                    )
 
         # OpenVINO
         elif self.xml:
@@ -771,11 +833,19 @@ class AutoBackend(nn.Module):
             im = im if isinstance(im, (list, tuple)) else [im]
             y = self.rknn_model.inference(inputs=im)
 
+        # Axelera
+        elif self.axelera:
+            y = self.ax_model(im.cpu())
+
+        # ExecuTorch
+        elif self.pte:
+            y = self.model.execute([im])
+
         # TensorFlow (SavedModel, GraphDef, Lite, Edge TPU)
         else:
             im = im.cpu().numpy()
             if self.saved_model:  # SavedModel
-                y = self.model(im, training=False) if self.keras else self.model.serving_default(im)
+                y = self.model.serving_default(im)
                 if not isinstance(y, list):
                     y = [y]
             elif self.pb:  # GraphDef
@@ -820,8 +890,6 @@ class AutoBackend(nn.Module):
                     y[1] = np.transpose(y[1], (0, 3, 1, 2))  # should be y = (1, 116, 8400), (1, 32, 160, 160)
             y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
 
-        # for x in y:
-        #     print(type(x), len(x)) if isinstance(x, (list, tuple)) else print(type(x), x.shape)  # debug shapes
         if isinstance(y, (list, tuple)):
             if len(self.names) == 999 and (self.task == "segment" or len(y) == 2):  # segments and names not defined
                 nc = y[0].shape[1] - y[1].shape[1] - 4  # y = (1, 32, 160, 160), (1, 116, 8400)
@@ -830,35 +898,35 @@ class AutoBackend(nn.Module):
         else:
             return self.from_numpy(y)
 
-    def from_numpy(self, x: np.ndarray) -> torch.Tensor:
-        """
-        Convert a numpy array to a tensor.
+    def from_numpy(self, x: np.ndarray | torch.Tensor) -> torch.Tensor:
+        """Convert a NumPy array to a torch tensor on the model device.
 
         Args:
-            x (np.ndarray): The array to be converted.
+            x (np.ndarray | torch.Tensor): Input array or tensor.
 
         Returns:
-            (torch.Tensor): The converted tensor
+            (torch.Tensor): Tensor on `self.device`.
         """
         return torch.tensor(x).to(self.device) if isinstance(x, np.ndarray) else x
 
     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.
+        """Warm up the model by running one forward pass with a dummy input.
 
         Args:
-            imgsz (tuple): The shape of the dummy input tensor in the format (batch_size, channels, height, width)
+            imgsz (tuple[int, int, int, int]): Dummy input shape in (batch, channels, height, width) format.
         """
         warmup_types = self.pt, self.jit, self.onnx, self.engine, self.saved_model, self.pb, self.triton, self.nn_module
         if any(warmup_types) and (self.device.type != "cpu" or self.triton):
             im = torch.empty(*imgsz, dtype=torch.half if self.fp16 else torch.float, device=self.device)  # input
             for _ in range(2 if self.jit else 1):
-                self.forward(im)  # warmup
+                self.forward(im)  # warmup model
+                warmup_boxes = torch.rand(1, 84, 16, device=self.device)  # 16 boxes works best empirically
+                warmup_boxes[:, :4] *= imgsz[-1]
+                non_max_suppression(warmup_boxes)  # warmup NMS
 
     @staticmethod
     def _model_type(p: str = "path/to/model.pt") -> list[bool]:
-        """
-        Take a path to a model file and return the model type.
+        """Take a path to a model file and return the model type.
 
         Args:
             p (str): Path to the model file.
@@ -867,8 +935,8 @@ class AutoBackend(nn.Module):
             (list[bool]): List of booleans indicating the model type.
 
         Examples:
-            >>> model = AutoBackend(model="path/to/model.onnx")
-            >>> model_type = model._model_type()  # returns "onnx"
+            >>> types = AutoBackend._model_type("path/to/model.onnx")
+            >>> assert types[2]  # onnx
         """
         from ultralytics.engine.exporter import export_formats
 
@@ -887,4 +955,4 @@ class AutoBackend(nn.Module):
             url = urlsplit(p)
             triton = bool(url.netloc) and bool(url.path) and url.scheme in {"http", "grpc"}
 
-        return types + [triton]
+        return [*types, triton]

ultralytics/nn/modules/__init__.py

@@ -103,80 +103,80 @@ from .transformer import (
 )
 
 __all__ = (
-    "Conv",
-    "Conv2",
-    "LightConv",
-    "RepConv",
-    "DWConv",
-    "DWConvTranspose2d",
-    "ConvTranspose",
-    "Focus",
-    "GhostConv",
-    "ChannelAttention",
-    "SpatialAttention",
+    "AIFI",
+    "C1",
+    "C2",
+    "C2PSA",
+    "C3",
+    "C3TR",
     "CBAM",
-    "Concat",
-    "TransformerLayer",
-    "TransformerBlock",
-    "MLPBlock",
-    "LayerNorm2d",
+    "CIB",
     "DFL",
-    "HGBlock",
-    "HGStem",
+    "ELAN1",
+    "MLP",
+    "OBB",
+    "PSA",
     "SPP",
+    "SPPELAN",
     "SPPF",
-    "C1",
-    "C2",
-    "C3",
+    "A2C2f",
+    "AConv",
+    "ADown",
+    "Attention",
+    "BNContrastiveHead",
+    "Bottleneck",
+    "BottleneckCSP",
     "C2f",
-    "C3k2",
-    "SCDown",
-    "C2fPSA",
-    "C2PSA",
     "C2fAttn",
-    "C3x",
-    "C3TR",
+    "C2fCIB",
+    "C2fPSA",
     "C3Ghost",
-    "GhostBottleneck",
-    "Bottleneck",
-    "BottleneckCSP",
-    "Proto",
-    "Detect",
-    "Segment",
-    "Pose",
+    "C3k2",
+    "C3x",
+    "CBFuse",
+    "CBLinear",
+    "ChannelAttention",
     "Classify",
-    "TransformerEncoderLayer",
-    "RepC3",
-    "RTDETRDecoder",
-    "AIFI",
+    "Concat",
+    "ContrastiveHead",
+    "Conv",
+    "Conv2",
+    "ConvTranspose",
+    "DWConv",
+    "DWConvTranspose2d",
     "DeformableTransformerDecoder",
     "DeformableTransformerDecoderLayer",
+    "Detect",
+    "Focus",
+    "GhostBottleneck",
+    "GhostConv",
+    "HGBlock",
+    "HGStem",
+    "ImagePoolingAttn",
+    "Index",
+    "LRPCHead",
+    "LayerNorm2d",
+    "LightConv",
+    "MLPBlock",
     "MSDeformAttn",
-    "MLP",
+    "MaxSigmoidAttnBlock",
+    "Pose",
+    "Proto",
+    "RTDETRDecoder",
+    "RepC3",
+    "RepConv",
+    "RepNCSPELAN4",
+    "RepVGGDW",
     "ResNetLayer",
-    "OBB",
+    "SCDown",
+    "Segment",
+    "SpatialAttention",
+    "TorchVision",
+    "TransformerBlock",
+    "TransformerEncoderLayer",
+    "TransformerLayer",
     "WorldDetect",
     "YOLOEDetect",
     "YOLOESegment",
     "v10Detect",
-    "LRPCHead",
-    "ImagePoolingAttn",
-    "MaxSigmoidAttnBlock",
-    "ContrastiveHead",
-    "BNContrastiveHead",
-    "RepNCSPELAN4",
-    "ADown",
-    "SPPELAN",
-    "CBFuse",
-    "CBLinear",
-    "AConv",
-    "ELAN1",
-    "RepVGGDW",
-    "CIB",
-    "C2fCIB",
-    "Attention",
-    "PSA",
-    "TorchVision",
-    "Index",
-    "A2C2f",
 )

ultralytics/nn/modules/activation.py

@@ -6,8 +6,7 @@ import torch.nn as nn
 
 
 class AGLU(nn.Module):
-    """
-    Unified activation function module from AGLU.
+    """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.
@@ -40,11 +39,10 @@ class AGLU(nn.Module):
         self.kappa = nn.Parameter(nn.init.uniform_(torch.empty(1, device=device, dtype=dtype)))  # kappa parameter
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """
-        Apply the Adaptive Gated Linear Unit (AGLU) activation function.
+        """Apply the Adaptive Gated Linear Unit (AGLU) activation function.
 
-        This forward method implements the AGLU activation function with learnable parameters lambda and kappa.
-        The function applies a transformation that adaptively combines linear and non-linear components.
+        This forward method implements the AGLU activation function with learnable parameters lambda and kappa. The
+        function applies a transformation that adaptively combines linear and non-linear components.
 
         Args:
             x (torch.Tensor): Input tensor to apply the activation function to.