LightPFN 1.0.0__py3-none-any.whl

lightpfn/__init__.py

@@ -0,0 +1,25 @@
+"""LightPFN: compact tabular inference, pretrained only on synthetic data."""
+
+__version__ = "1.0.0"
+__all__ = ["Config", "LightPFN", "LightPFNClassifier", "load_model", "load_pretrained", "save_model", "__version__"]
+
+
+def __getattr__(name):
+    if name in ("Config", "LightPFN"):
+        from lightpfn.model import lightpfn
+        value = getattr(lightpfn, name)
+    elif name == "LightPFNClassifier":
+        try:
+            from lightpfn.sklearn import LightPFNClassifier
+        except ModuleNotFoundError as exc:
+            if exc.name != "sklearn":
+                raise
+            raise ImportError("LightPFNClassifier requires scikit-learn: `pip install scikit-learn`.") from exc
+        value = LightPFNClassifier
+    elif name in ("load_model", "load_pretrained", "save_model"):
+        from lightpfn import checkpoint
+        value = getattr(checkpoint, name)
+    else:
+        raise AttributeError(name)
+    globals()[name] = value
+    return value

lightpfn/checkpoint.py

@@ -0,0 +1,107 @@
+"""Safe local checkpoints and revision-pinned Hugging Face downloads.
+
+The release format is model.safetensors plus config.json. Legacy tensor/dict
+checkpoints remain supported using PyTorch's restricted weights-only loader.
+There is deliberately no fallback to unrestricted pickle.
+"""
+
+import json
+from collections.abc import Mapping
+from dataclasses import asdict
+from importlib.resources import files
+from pathlib import Path
+
+import torch
+
+from lightpfn.model.lightpfn import Config, LightPFN
+
+
+def _safetensors():
+    try:
+        from safetensors.torch import load_file, save_file
+    except ImportError as exc:
+        raise ImportError("Safetensors weights require `pip install safetensors`.") from exc
+    return load_file, save_file
+
+
+def _config(values):
+    if not isinstance(values, Mapping):
+        raise ValueError("Checkpoint config must be a mapping of architecture parameters.")
+    values = dict(values)
+    if "group_offsets" in values:
+        values["group_offsets"] = tuple(values["group_offsets"])
+    return Config(**values)
+
+
+def load_model(path, device="cpu"):
+    """Load a safetensors directory/file or a weights-only compatible .pt file.
+
+    Safetensors files need config.json in the same directory. Full training
+    checkpoints containing arbitrary Python objects are intentionally rejected.
+    """
+    path = Path(path)
+    if path.is_dir():
+        path = path / "model.safetensors"
+    if path.suffix == ".safetensors":
+        load_file, _ = _safetensors()
+        config = json.loads(path.with_name("config.json").read_text(encoding="utf-8"))
+        state = load_file(str(path), device="cpu")
+    else:
+        checkpoint = torch.load(path, map_location="cpu", weights_only=True)
+        if not isinstance(checkpoint, Mapping) or "config" not in checkpoint:
+            raise ValueError("Checkpoint must contain 'config' and 'ema' or 'model' weights.")
+        config = checkpoint["config"]
+        state = checkpoint.get("ema", checkpoint.get("model"))
+    if not isinstance(state, Mapping) or not state or not all(
+        isinstance(key, str) and isinstance(value, torch.Tensor) for key, value in state.items()
+    ):
+        raise ValueError("Checkpoint weights must be a nonempty tensor state dictionary.")
+    model = LightPFN(_config(config))
+    model.load_state_dict(state, strict=True)
+    return model.to(device).eval()
+
+
+def save_model(model, directory):
+    """Export unfused inference weights and JSON config, excluding optimizer state."""
+    if getattr(model, "is_folded", False):
+        raise ValueError("Export the original model; folded weights use a different layout.")
+    _, save_file = _safetensors()
+    directory = Path(directory)
+    directory.mkdir(parents=True, exist_ok=True)
+    state = {key: value.detach().cpu().contiguous().clone() for key, value in model.state_dict().items()}
+    save_file(state, str(directory / "model.safetensors"), metadata={"format": "pt"})
+    (directory / "config.json").write_text(json.dumps(asdict(model.cfg), indent=2) + "\n", encoding="utf-8")
+    return directory
+
+
+def pretrained_spec():
+    """Return the packaged model repository and immutable revision."""
+    return json.loads(files("lightpfn").joinpath("pretrained.json").read_text(encoding="utf-8"))
+
+
+def load_pretrained(*, repo_id=None, revision=None, device="cpu", cache_dir=None, local_files_only=False):
+    """Load the released model, or an explicit repo and revision using cached HF authentication.
+
+    A custom repository requires a revision. No remote Python code is executed.
+    For private repositories authenticate with `hf auth login` or HF_TOKEN.
+    """
+    try:
+        from huggingface_hub import hf_hub_download
+    except ImportError as exc:
+        raise ImportError("Hugging Face downloads require `pip install huggingface-hub`.") from exc
+    spec = pretrained_spec()
+    repo_id = spec["repo_id"] if repo_id is None else repo_id
+    if revision is None:
+        if repo_id != spec["repo_id"]:
+            raise ValueError("Provide revision when using a custom Hugging Face repository.")
+        revision = spec["revision"]
+    if not revision:
+        raise ValueError("No pretrained revision configured; provide a checkpoint or explicit revision.")
+    options = dict(repo_id=repo_id, revision=revision, cache_dir=cache_dir, local_files_only=local_files_only)
+    config_path = Path(hf_hub_download(filename="config.json", **options))
+    # Use the resolved commit from the snapshot path for both files, even when a
+    # caller explicitly chose a mutable branch or tag.
+    resolved = config_path.parent.name
+    options["revision"] = resolved
+    weights_path = Path(hf_hub_download(filename="model.safetensors", **options))
+    return load_model(weights_path, device=device)

lightpfn/device.py

@@ -0,0 +1,51 @@
+"""Inference device: "auto" takes torch's GPU backend when there is one (CUDA on NVIDIA, also ROCm builds of
+torch, which use the same "cuda" device), otherwise a GPU with a Vulkan driver (AMD, Intel or NVIDIA,
+through wgpu), otherwise the CPU. Each can be asked for explicitly: "cuda[:i]", "vulkan[:i]" (index into
+lightpfn.vulkan.adapters()), "cpu". The environment variable LIGHTPFN_DEVICE replaces "auto" with any of
+these, e.g. LIGHTPFN_DEVICE=cpu to keep a GPU free.
+"""
+
+import os
+
+import torch
+
+DEVICE_ENV = "LIGHTPFN_DEVICE"
+
+
+def resolve_device(device="auto"):
+    """The device string LightPFNClassifier runs on: "cuda[:i]", "vulkan[:i]", "cpu" (or "mps")."""
+    device = "auto" if device is None else str(device).strip().lower()
+    if device == "auto":
+        device = os.environ.get(DEVICE_ENV, "").strip().lower() or "auto"
+    if device == "auto":
+        if torch.cuda.is_available():
+            return "cuda"
+        from lightpfn import vulkan
+
+        return "vulkan" if vulkan.is_available() else "cpu"
+    kind, _, index = device.partition(":")
+    if ":" in device and not index:
+        raise ValueError(f"device {device!r}: missing index after ':'")
+    if index and not index.isdigit():
+        raise ValueError(f"device {device!r}: the index after ':' must be a number")
+    if kind == "cpu":
+        return "cpu"
+    if kind == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("device 'cuda' requested, but this torch build sees no CUDA/ROCm GPU")
+        return device
+    if kind == "vulkan":
+        from lightpfn import vulkan
+
+        if not vulkan.is_available(int(index) if index else None):
+            raise RuntimeError(f"device {device!r} requested, but no Vulkan adapter was found: it needs a GPU "
+                               "driver with Vulkan and `pip install wgpu` (adapters: lightpfn.vulkan.adapters())")
+        return device
+    if kind == "mps":
+        return device
+    raise ValueError(f"unknown device {device!r}: use 'auto', 'cpu', 'cuda[:i]' or 'vulkan[:i]'")
+
+
+def kind(device):
+    """"cpu", "cuda", "vulkan" or "mps" of a resolved device string."""
+    return str(device).partition(":")[0]

lightpfn/model/__init__.py

@@ -0,0 +1 @@
+"""LightPFN architecture; the inference package does not import synthetic priors."""

lightpfn/model/catstats.py

@@ -0,0 +1,117 @@
+"""Category statistics of the context rows for the categorical adapter (Config.cat_adapter).
+
+Every categorical cell receives the class distribution of its category among the context rows,
+as a deviation from the class prior, plus count features. Query rows use every context row of
+their category. A context row uses only the context rows that come before it in a random order
+of the rows (CatBoost's ordered target statistics), so its own label never enters its features,
+and no label is ever correlated with its own statistics beyond the true dependence (leave-one-out
+statistics are: the model collapsed on them in our tests). The count of rows the
+statistics rest on is a feature, so the model can weigh a context row's noisier estimate.
+
+Missing values form one category of their own. An unseen category gets zero features.
+Everything is vectorized over tasks and columns; only the categorical columns are processed.
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+
+N_COUNT = 3  # log1p(rows the statistics use), log1p(other rows of the category), their fraction of the rows
+
+
+def _keys(x):
+    """Category keys: the values, with missing values as one key that sorts last."""
+    x = x.float()
+    return torch.where(torch.isnan(x), torch.full_like(x, math.inf), x)
+
+
+def _features(counts, n_used, n_freq, prior, n, smooth):
+    """counts (..., C) class counts of the rows used, n_used/n_freq (...) row counts, prior (B, C)
+    broadcastable to counts; returns (..., C + N_COUNT)."""
+    delta = (counts - n_used[..., None] * prior) / (n_used[..., None] + smooth)
+    cnt = torch.stack([torch.log1p(n_used) / 4, torch.log1p(n_freq) / 4, n_freq / n], -1)
+    return torch.cat([delta, cnt], -1)
+
+
+def context_stats(x, y, cat, n_classes, smooth=1.0, perm=None, generator=None):
+    """x: (B, n, m) context values (NaN = missing), y: (B, n) classes in 0..n_classes-1, cat: (B, m) bool.
+    perm: (B, n) the row order of the ordered statistics. Default: random per task from the global RNG, or,
+    with a generator, one order drawn from it for every task (the same whatever the batch).
+
+    Returns (feats, table): feats (B, n, m, n_classes + N_COUNT) float32, zero on numerical columns, and
+    the table query_stats needs; (None, None) when no column is categorical."""
+    B, n, m = x.shape
+    C = int(n_classes)
+    dev = x.device
+    mc = int(cat.sum(1).max()) if cat.numel() else 0
+    if mc == 0:
+        return None, None
+    # the categorical columns of each task first, in order; the tail of a task with fewer is padding
+    cols = torch.argsort((~cat).to(torch.int8), dim=1, stable=True)[:, :mc]  # (B, mc)
+    valid = torch.gather(cat, 1, cols)
+    keys = _keys(torch.gather(x, 2, cols[:, None].expand(B, n, mc)))
+    if perm is None:
+        if generator is None:
+            perm = torch.argsort(torch.rand(B, n, device=dev), 1)
+        else:
+            perm = torch.argsort(torch.rand(n, generator=generator)).expand(B, n)
+    perm = perm.to(dev)
+    pe = perm[..., None].expand(B, n, mc)
+    kp = torch.gather(keys, 1, pe)
+    vs, order = torch.sort(kp, dim=1, stable=True)  # stable: rows of a category keep the random order
+    rows = torch.gather(pe, 1, order)  # original row of every sorted position
+    ys = torch.gather(y.long()[..., None].expand(B, n, mc), 1, rows)
+    oh = F.one_hot(ys, C).float()  # (B, n, mc, C)
+    start = torch.ones_like(vs, dtype=torch.bool)
+    start[:, 1:] = vs[:, 1:] != vs[:, :-1]
+    end = torch.ones_like(start)
+    end[:, :-1] = start[:, 1:]
+    ar = torch.arange(n, device=dev)[None, :, None]
+    first = torch.where(start, ar, 0).cummax(1).values  # first sorted position of the category
+    last = (n - 1 - torch.where(end.flip(1), ar, 0).cummax(1).values).flip(1)  # last one
+    incl = oh.cumsum(1)
+    excl = incl - oh
+    del oh
+    base = torch.gather(excl, 1, first[..., None].expand_as(excl))
+    prefix = excl.sub_(base)  # rows of the category before this one in the random order
+    # Query counts are constant within a category. Cache one entry per category rather than n
+    # repeated entries per column (hundreds of MB per estimator at 60k rows / 100 columns).
+    n_keys = int((start.sum(1) * valid).max())
+    category = start.long().cumsum(1) - 1
+    bi, ri, ci = torch.where(end & valid[:, None])
+    ki = category[bi, ri, ci]
+    table_keys = torch.full((B, mc, n_keys), math.inf, device=dev)
+    table_total = torch.zeros(B, mc, n_keys, C, device=dev)
+    table_keys[bi, ci, ki] = vs[bi, ri, ci]
+    table_total[bi, ci, ki] = incl[bi, ri, ci] - base[bi, ri, ci]
+    del incl, excl, base
+    prior = F.one_hot(y.long(), C).float().mean(1)[:, None, None]  # (B, 1, 1, C)
+    n_used = (ar - first).float()
+    feats = _features(prefix, n_used, (last - first).float(), prior, n, smooth) * valid[:, None, :, None]
+    out = torch.zeros(B, n, mc, feats.shape[-1], device=dev)
+    out.scatter_(1, rows[..., None].expand_as(feats), feats)
+    full = torch.zeros(B, n, m, feats.shape[-1], device=dev)
+    full.scatter_(2, cols[:, None, :, None].expand_as(out), out)  # cols is a permutation prefix: no repeats
+    table = dict(cols=cols, valid=valid, keys=table_keys,
+                 total=table_total, prior=prior, n=n, smooth=smooth)
+    return full, table
+
+
+def query_stats(table, x):
+    """Features (B, M, m, C + N_COUNT) of query rows x: (B, M, m) from the table of context_stats."""
+    if table is None:
+        return None
+    B, M, m = x.shape
+    cols, valid, keys, total = table["cols"], table["valid"], table["keys"], table["total"]
+    mc, n_keys, C = keys.shape[1], keys.shape[2], total.shape[-1]
+    q = _keys(torch.gather(x, 2, cols[:, None].expand(B, M, mc))).transpose(1, 2).contiguous()  # (B, mc, M)
+    pos = torch.searchsorted(keys, q).clamp(max=n_keys - 1)
+    hit = torch.gather(keys, 2, pos) == q
+    counts = torch.gather(total, 2, pos[..., None].expand(B, mc, M, C)) * hit[..., None]
+    n_used = counts.sum(-1)
+    feats = _features(counts, n_used, n_used, table["prior"], table["n"], table["smooth"])
+    feats = (feats * valid[:, :, None, None]).transpose(1, 2)  # (B, M, mc, F)
+    full = torch.zeros(B, M, m, feats.shape[-1], device=x.device)
+    full.scatter_(2, cols[:, None, :, None].expand_as(feats), feats)
+    return full

lightpfn/model/layers.py

@@ -0,0 +1,190 @@
+"""Building blocks shared by the three stages of the model.
+
+Tensors are (batch, sequence, dim) unless a name says otherwise. Attention runs through
+F.scaled_dot_product_attention, so the same code uses flash kernels on GPU and fused CPU kernels.
+"""
+
+import math
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class MLP(nn.Module):
+    """Two-layer GELU feed-forward block; the output projection starts at zero so every residual
+    block starts as the identity (as in TabPFN-3)."""
+
+    def __init__(self, dim, hidden):
+        super().__init__()
+        self.fc1 = nn.Linear(dim, hidden, bias=False)
+        self.fc2 = nn.Linear(hidden, dim, bias=False)
+        nn.init.zeros_(self.fc2.weight)
+
+    def forward(self, x):
+        return self.fc2(F.gelu(self.fc1(x)))
+
+
+class SoftmaxScaling(nn.Module):
+    """Scales attention queries as a function of the number of keys n (TabPFN-3 SoftmaxScalingMLP):
+    q * base(log n) * (1 + tanh(mod(q))). Keeps attention sharp when the context grows far beyond
+    the lengths seen in training. base starts at 1 so the layer starts as a no-op."""
+
+    def __init__(self, n_heads, head_dim, hidden=64):
+        super().__init__()
+        self.n_heads, self.head_dim = n_heads, head_dim
+        self.base = nn.Sequential(nn.Linear(1, hidden), nn.GELU(), nn.Linear(hidden, n_heads * head_dim))
+        self.mod = nn.Sequential(nn.Linear(head_dim, hidden), nn.GELU(), nn.Linear(hidden, head_dim))
+        nn.init.zeros_(self.base[2].weight)
+        nn.init.ones_(self.base[2].bias)
+        nn.init.zeros_(self.mod[2].weight)
+        nn.init.zeros_(self.mod[2].bias)
+
+    def forward(self, q, n):
+        """q: (B, H, L, D) queries, n: number of keys."""
+        logn = torch.full((1, 1), math.log(max(n, 2)), device=q.device, dtype=q.dtype)
+        base = self.base(logn).view(1, self.n_heads, 1, self.head_dim)
+        return q * base * (1 + torch.tanh(self.mod(q)))
+
+
+# CUDA attention kernels put the batch on a grid dimension of at most 65535 blocks: a larger batch fails
+# with "invalid argument" (the row stages of a 2-task group of 50k-row tables have 100k rows in the batch).
+SDPA_MAX_BATCH = 32768
+
+
+def sdpa(q, k, v, mask=None):
+    """F.scaled_dot_product_attention in chunks of at most SDPA_MAX_BATCH along the batch dim (a size-1
+    batch dim of k, v or mask broadcasts)."""
+    B = q.shape[0]
+    if B <= SDPA_MAX_BATCH:
+        return F.scaled_dot_product_attention(q, k, v, attn_mask=mask)
+    out = []
+    for i in range(0, B, SDPA_MAX_BATCH):
+        part = [t if t is None or t.shape[0] == 1 else t[i:i + SDPA_MAX_BATCH] for t in (k, v, mask)]
+        out.append(F.scaled_dot_product_attention(q[i:i + SDPA_MAX_BATCH], *part[:2], attn_mask=part[2]))
+    return torch.cat(out)
+
+
+def rope(x, pos, base=10000.0, interleaved=False):
+    """Rotary position embedding (rotate-half form) on the last dim of x: (..., L, D) with
+    positions pos: (L,). interleaved=True rotates the pairs (2i, 2i + 1) instead of (i, i + D/2),
+    with one complex multiply: the same function for weights passed through interleave_rotary()."""
+    d = x.shape[-1]
+    inv_freq = base ** (-torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d)
+    ang = pos.to(torch.float32)[:, None] * inv_freq[None]  # (L, D/2)
+    if interleaved and x.dtype in (torch.float32, torch.float64):
+        rot = torch.complex(ang.cos(), ang.sin()).to(torch.complex64 if x.dtype == torch.float32 else torch.complex128)
+        return torch.view_as_real(torch.view_as_complex(x.unflatten(-1, (d // 2, 2))) * rot).flatten(-2)
+    cos, sin = ang.cos().to(x.dtype), ang.sin().to(x.dtype)
+    if interleaved:  # reduced precision (autocast): the same arithmetic as below on adjacent pairs
+        x1, x2 = x[..., 0::2], x[..., 1::2]
+        return torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1).flatten(-2)
+    x1, x2 = x[..., : d // 2], x[..., d // 2 :]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+
+class Attention(nn.Module):
+    """Multi-head attention with separate query and key/value inputs.
+
+    kv_heads_query < n_heads lets a set of queries use only the first key/value heads (multi-query
+    attention for test rows, as in TabPFN-3.5): keys/values of the other heads never need to be
+    cached for prediction."""
+
+    def __init__(self, dim, n_heads, scaling=False):
+        super().__init__()
+        assert dim % n_heads == 0
+        self.n_heads, self.head_dim = n_heads, dim // n_heads
+        self.q = nn.Linear(dim, dim, bias=False)
+        self.kv = nn.Linear(dim, 2 * dim, bias=False)
+        self.out = nn.Linear(dim, dim, bias=False)
+        nn.init.zeros_(self.out.weight)
+        self.scaling = SoftmaxScaling(n_heads, self.head_dim) if scaling else None
+
+    def split(self, x):
+        B, L, _ = x.shape
+        return x.view(B, L, self.n_heads, self.head_dim).transpose(1, 2)  # (B, H, L, D)
+
+    def keys_values(self, x):
+        k, v = self.kv(x).chunk(2, dim=-1)
+        return self.split(k), self.split(v)
+
+    def attend(self, x, k, v, mask=None, q_rope=None, kv_heads=None):
+        """x: (B, Lq, dim) queries; k, v: (B, Hkv, Lk, D). kv_heads=h: all query heads use the
+        first h key/value heads."""
+        q = self.split(self.q(x))
+        if q_rope is not None:
+            q = q_rope(q)
+        if self.scaling is not None:
+            q = self.scaling(q, k.shape[2])
+        if kv_heads is not None:
+            k = k[:, :kv_heads].repeat_interleave(self.n_heads // kv_heads, dim=1)
+            v = v[:, :kv_heads].repeat_interleave(self.n_heads // kv_heads, dim=1)
+        o = sdpa(q, k, v, mask)
+        B, H, L, D = o.shape
+        return self.out(o.transpose(1, 2).reshape(B, L, H * D))
+
+
+class Block(nn.Module):
+    """Pre-norm residual block: attention of x on a key/value sequence, then an MLP."""
+
+    def __init__(self, dim, n_heads, ff_factor=2, scaling=False, ff_hidden=None):
+        super().__init__()
+        self.norm_q = nn.RMSNorm(dim)
+        self.norm_kv = nn.RMSNorm(dim)
+        self.norm_ff = nn.RMSNorm(dim)
+        self.attn = Attention(dim, n_heads, scaling)
+        self.mlp = MLP(dim, dim * ff_factor if ff_hidden is None else ff_hidden)
+
+    def keys_values(self, kv_input):
+        return self.attn.keys_values(self.norm_kv(kv_input))
+
+    def forward(self, x, k, v, mask=None, q_rope=None, kv_heads=None):
+        x = x + self.attn.attend(self.norm_q(x), k, v, mask, q_rope, kv_heads)
+        return x + self.mlp(self.norm_ff(x))
+
+
+class FoldedRMSNorm(nn.Module):
+    """RMSNorm whose weight was folded into the linear layer that reads it (inference only): one
+    reduction and one multiply instead of the unfused CPU kernels. Same eps as nn.RMSNorm (eps=None:
+    that of the accumulation dtype, float32 for reduced precision), reduction in at least float32."""
+
+    def __init__(self, dim, eps=None):
+        super().__init__()
+        self.dim, self.eps = dim, eps
+
+    def forward(self, x):
+        acc = torch.promote_types(x.dtype, torch.float32)
+        eps = torch.finfo(acc).eps if self.eps is None else self.eps
+        ms = torch.linalg.vector_norm(x, dim=-1, keepdim=True, dtype=acc).square_().div_(self.dim)
+        return x * ms.add_(eps).rsqrt_().to(x.dtype)
+
+
+@torch.no_grad()
+def fold_norm(norm, *linears):
+    """Moves the weight of an RMSNorm into the linear layers fed by it; returns the weightless norm."""
+    for lin in linears:
+        lin.weight.mul_(norm.weight)
+    return FoldedRMSNorm(norm.weight.numel(), norm.eps)
+
+
+@torch.no_grad()
+def interleave_rotary(block):
+    """Reorders the query and key dims of each head of a block from (i, i + D/2) pairs to adjacent
+    (2i, 2i + 1) pairs: q.k is unchanged (same permutation on both), and rope(..., interleaved=True)
+    rotates the same pairs."""
+    a = block.attn
+    H, D = a.n_heads, a.head_dim
+    pairs = torch.stack([torch.arange(D // 2), torch.arange(D // 2) + D // 2], -1).flatten()
+    rows = (torch.arange(H)[:, None] * D + pairs).flatten().to(a.q.weight.device)
+    a.q.weight.copy_(a.q.weight[rows])
+    a.kv.weight[: H * D].copy_(a.kv.weight[rows])
+
+
+class OrthogonalEmbedding(nn.Embedding):
+    """Label embedding with orthonormal initial rows (TabPFN-3 TrainableOrthogonalEmbedding)."""
+
+    def __init__(self, n, dim):
+        super().__init__(n, dim)
+        with torch.no_grad():
+            q, _ = torch.linalg.qr(torch.randn(dim, min(n, dim)))
+            self.weight[: q.shape[1]] = q.T