types, review

src/lora_lite/__init__.py

@@ -1,3 +1,11 @@
+import os as _os
+
+# Optional runtime shape/dtype checking via jaxtyping + beartype.
+# Set BEARTYPE=1 for smoke tests / debugging; off by default for zero overhead.
+if _os.environ.get("BEARTYPE"):
+    from beartype.claw import beartype_this_package as _bt
+    _bt()
+
 from .config import LoraLiteConfig
 from .adapter import attach, detach, save, load
 from .variant import REGISTRY, register, ParamSpec, Variant

src/lora_lite/config.py

@@ -1,17 +1,19 @@
 from dataclasses import dataclass, field, asdict
-from typing import Any
+from typing import Any, Literal
 import torch
 
+Role = Literal["reader", "writer", "inner"]
+
 
 @dataclass
 class LoraLiteConfig:
     variant: str = "lora"
     r: int = 8
-    alpha: float = 16.0
+    alpha: float | int = 16.0
     dtype: torch.dtype = torch.bfloat16
 
     # targeting
-    target_roles: tuple[str, ...] = ("reader", "writer")
+    target_roles: tuple[Role, ...] = ("reader", "writer")
     target_names: tuple[str, ...] = ()
     exclude_names: tuple[str, ...] = ("lm_head", "embed_tokens")
     layers: tuple[int, ...] | None = None
@@ -26,12 +28,9 @@ class LoraLiteConfig:
 
     @classmethod
     def from_dict(cls, d: dict) -> "LoraLiteConfig":
+        # to_dict always serializes dtype as str; torch.save preserves tuples.
+        # If you build the dict by hand, pass the right types -- fail loud otherwise.
         d = dict(d)
-        if isinstance(d.get("dtype"), str):
-            d["dtype"] = getattr(torch, d["dtype"])
-        if isinstance(d.get("layers"), list):
-            d["layers"] = tuple(d["layers"])
-        for k in ("target_roles", "target_names", "exclude_names"):
-            if isinstance(d.get(k), list):
-                d[k] = tuple(d[k])
+        d["dtype"] = getattr(torch, d["dtype"])
         return cls(**d)
+

src/lora_lite/variants/antipasto.py

@@ -34,12 +34,12 @@ WHICH BASIS IS ROTATED:
 
 REQUIRES even rank divisible by `block_size` (default 4). r=8, bs=4 -> 2 blocks.
 """
-from __future__ import annotations
 import math
 
 import torch
 from einops import einsum
-from torch import nn
+from jaxtyping import Float
+from torch import nn, Tensor as T
 
 from ..variant import register, ParamSpec
 
@@ -96,7 +96,7 @@ class AntiPaSTO:
         }
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         if type(layer) is not nn.Linear:
             raise TypeError(
                 "AntiPaSTO mutates layer.weight into W_res (like PiSSA), so v1 "
@@ -116,7 +116,11 @@ class AntiPaSTO:
             layer.weight.data.copy_(W_res)
 
     @staticmethod
-    def forward(layer: nn.Linear, x, y):
+    def forward(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+        y: Float[T, '*B o'],
+    ) -> Float[T, '*B o']:
         cfg = layer._lora_cfg
         bs = int(cfg.variant_kwargs.get("block_size", 4))
         max_angle = float(cfg.variant_kwargs.get("max_rotation_angle", 0.5))

src/lora_lite/variants/delora.py

@@ -33,7 +33,8 @@ Reference implementations:
 """
 import torch
 from einops import einsum
-from torch import nn
+from jaxtyping import Float
+from torch import nn, Tensor as T
 
 from ..variant import register, ParamSpec
 
@@ -57,7 +58,7 @@ class DeLoRA:
         }
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         # Reading layer.weight only works for plain Linear; for bnb layers this
         # dequantizes via .float() round-trip if available, or fails cleanly.
         with torch.no_grad():
@@ -67,7 +68,11 @@ class DeLoRA:
         return
 
     @staticmethod
-    def forward(layer: nn.Linear, x, y):
+    def forward(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+        y: Float[T, '*B o'],
+    ) -> Float[T, '*B o']:
         cfg = layer._lora_cfg
         A = layer.lora_A                                   # (r, d_in)
         B = layer.lora_B                                   # (d_out, r)

src/lora_lite/variants/dora.py

@@ -19,9 +19,9 @@ Reference implementations (for review/cross-check):
     (offline: docs/refs/peft_lora_dora.py)
 """
 import torch
-import torch.nn.functional as F
 from einops import einsum
-from torch import nn
+from jaxtyping import Float
+from torch import nn, Tensor as T
 
 from ..variant import register, ParamSpec
 
@@ -40,7 +40,7 @@ class DoRA:
         }
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         if type(layer) is not nn.Linear:
             raise TypeError(
                 "DoRA needs ||W||_c, so v1 only supports plain nn.Linear. "
@@ -52,7 +52,11 @@ class DoRA:
             layer.lora_m.data.copy_(col_norm)
 
     @staticmethod
-    def forward(layer: nn.Linear, x, y):
+    def forward(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+        y: Float[T, '*B o'],
+    ) -> Float[T, '*B o']:
         cfg = layer._lora_cfg
         scale = cfg.alpha / cfg.r
         # V = W + scale * B @ A

src/lora_lite/variants/eva.py

@@ -31,14 +31,17 @@ References:
     https://github.com/huggingface/peft/blob/main/examples/eva_finetuning/eva_finetuning.py
     (offline: docs/refs/peft_eva_finetuning.py)
 """
-from __future__ import annotations
-
 import torch
 from einops import einsum
-from torch import nn
+from jaxtyping import Float
+from torch import nn, Tensor as T
+from typing import Iterable
 
 from ..variant import register, ParamSpec
 
+CalibrationBatch = dict | tuple | list | T
+CalibrationData = Iterable[CalibrationBatch]
+
 
 @register
 class EVA:
@@ -55,12 +58,12 @@ class EVA:
         }
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         # No-op; group_init does the data-driven SVD across all targets at once.
         return
 
     @staticmethod
-    def group_init(model: nn.Module, targets, cfg, calibration_data) -> None:
+    def group_init(model: nn.Module, targets, cfg, calibration_data: CalibrationData | None) -> None:
         # adapter.load() passes _skip_group_init=True so this is only called on
         # the live attach path where calibration_data is required.
         if calibration_data is None:
@@ -72,7 +75,7 @@ class EVA:
             )
         # Collect input activations per target via forward hooks.
         layers = {name: layer for name, layer, _ in targets}
-        captured: dict[str, list[torch.Tensor]] = {n: [] for n in layers}
+        captured: dict[str, list[T]] = {n: [] for n in layers}
 
         def make_hook(name):
             def _h(module, args, kwargs):
@@ -115,7 +118,11 @@ class EVA:
             layer.lora_A.copy_(A)
 
     @staticmethod
-    def forward(layer: nn.Linear, x, y):
+    def forward(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+        y: Float[T, '*B o'],
+    ) -> Float[T, '*B o']:
         cfg = layer._lora_cfg
         scale = cfg.alpha / cfg.r
         h = einsum(x, layer.lora_A, "... i, r i -> ... r")

src/lora_lite/variants/hra.py

@@ -30,7 +30,8 @@ Reference implementations (for review/cross-check):
 """
 import torch
 from einops import einsum
-from torch import nn
+from jaxtyping import Float
+from torch import nn, Tensor as T
 
 from ..variant import register, ParamSpec
 
@@ -53,7 +54,7 @@ class HRA:
         }
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         # Symmetric init per peft (docs/refs/peft_hra_layer.py:101-108):
         #   half = kaiming(r//2, d_in); U = repeat_interleave(half, 2, dim=0)
         # Adjacent pairs (H_2k H_2k+1) cancel since H^2 = I, so R = I exactly,
@@ -66,7 +67,10 @@ class HRA:
         return
 
     @staticmethod
-    def forward_input(layer: nn.Linear, x: torch.Tensor) -> torch.Tensor:
+    def forward_input(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+    ) -> Float[T, '*B i']:
         """Apply Rx where R = prod_i H_i, H_i = I - 2 u_i u_i^T / ||u_i||^2."""
         U = layer.lora_U                                     # (r, d_in)
         Rx = x

src/lora_lite/variants/ia3.py

@@ -25,7 +25,8 @@ Reference implementation:
     https://github.com/huggingface/peft/blob/main/src/peft/tuners/ia3/layer.py
 """
 import torch
-from torch import nn
+from jaxtyping import Float
+from torch import nn, Tensor as T
 
 from ..variant import register, ParamSpec
 
@@ -39,11 +40,15 @@ class IA3:
         return {"lora_g": ParamSpec((d_out,), init="ones", trainable=True)}
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         return
 
     @staticmethod
-    def forward(layer: nn.Linear, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    def forward(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+        y: Float[T, '*B o'],
+    ) -> Float[T, '*B o']:
         return y * layer.lora_g
 
 
@@ -56,9 +61,12 @@ class IA3FF:
         return {"lora_g": ParamSpec((d_in,), init="ones", trainable=True)}
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         return
 
     @staticmethod
-    def forward_input(layer: nn.Linear, x: torch.Tensor) -> torch.Tensor:
+    def forward_input(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+    ) -> Float[T, '*B i']:
         return x * layer.lora_g

src/lora_lite/variants/lora.py

@@ -10,7 +10,8 @@ Reference implementations (for review/cross-check):
     (see docs/refs/peft_lora_layer.py for offline copy)
 """
 from einops import einsum
-from torch import nn
+from jaxtyping import Float
+from torch import nn, Tensor as T
 import torch
 
 from ..variant import register, ParamSpec
@@ -28,12 +29,16 @@ class LoRA:
         }
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         # B is zeros => delta=0 at t=0; identity invariant holds.
         return
 
     @staticmethod
-    def forward(layer: nn.Linear, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    def forward(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+        y: Float[T, '*B o'],
+    ) -> Float[T, '*B o']:
         cfg = layer._lora_cfg
         scale = cfg.alpha / cfg.r
         h = einsum(x, layer.lora_A, "... i, r i -> ... r")

src/lora_lite/variants/pissa.py

@@ -22,7 +22,8 @@ Reference implementations (for review/cross-check):
 """
 import torch
 from einops import einsum
-from torch import nn
+from jaxtyping import Float
+from torch import nn, Tensor as T
 
 from ..variant import register, ParamSpec
 
@@ -39,7 +40,7 @@ class PiSSA:
         }
 
     @staticmethod
-    def init(layer: nn.Linear, cfg) -> None:
+    def init(layer: nn.Module, cfg) -> None:
         if type(layer) is not nn.Linear:
             raise TypeError(
                 "PiSSA mutates layer.weight into W_res, so v1 only supports plain nn.Linear. "
@@ -63,7 +64,11 @@ class PiSSA:
         layer.weight.data.copy_((W - scale * BA).to(layer.weight.dtype))
 
     @staticmethod
-    def forward(layer: nn.Linear, x, y):
+    def forward(
+        layer: nn.Module,
+        x: Float[T, '*B i'],
+        y: Float[T, '*B o'],
+    ) -> Float[T, '*B o']:
         cfg = layer._lora_cfg
         scale = cfg.alpha / cfg.r
         h = einsum(x, layer.lora_A, "... i, r i -> ... r")