kasn 0.1.0__py3-none-any.whl

kasn/__init__.py

@@ -0,0 +1,49 @@
+"""Kolmogorov-Arnold Sieve Networks (KASN).
+
+Sieve extremum estimation with deep Kolmogorov-Arnold networks, learned
+B-spline edges and Group Lasso regularisation.
+
+KASNRegressor
+    Squared, Huber, quantile and Poisson criteria.
+KASNClassifier
+    Binary logistic criterion.
+tune_kasn
+    Prune-aware greedy hyperparameter search on a validation fold.
+make_kasn_regression
+    Monte Carlo designs of Section 4.
+load
+    Restore any estimator written by ``.save()``.
+"""
+
+from typing import Optional
+
+from .base import _read_payload
+from .classifier import KASNClassifier
+from .datasets import make_kasn_regression, target_function
+from .losses import make_loss
+from .network import KASNLayer, KASNNetwork
+from .preprocessing import EmpiricalCDFTransformer
+from .regressor import KASNRegressor
+from .splines import BSplineBasis
+from .theory import architecture, check_conditions
+from .tuning import TuningResult, tune_kasn
+
+__version__ = "0.1.0"
+
+
+def load(path, device: Optional[str] = None):
+    """Load a ``KASNRegressor`` or ``KASNClassifier`` written by ``.save()``."""
+    payload = _read_payload(path)
+    classes = {"KASNRegressor": KASNRegressor,
+               "KASNClassifier": KASNClassifier}
+    if payload["class"] not in classes:
+        raise ValueError(f"unknown estimator class {payload['class']!r}")
+    return classes[payload["class"]]._from_payload(payload, device)
+
+
+__all__ = [
+    "KASNRegressor", "KASNClassifier", "load", "tune_kasn", "TuningResult",
+    "make_kasn_regression", "target_function", "make_loss",
+    "EmpiricalCDFTransformer", "BSplineBasis", "KASNLayer", "KASNNetwork",
+    "architecture", "check_conditions", "__version__",
+]

kasn/_utils.py

@@ -0,0 +1,55 @@
+"""Small internal helpers: devices, seeding and serialisation hygiene."""
+
+from __future__ import annotations
+
+import contextlib
+from typing import Any, Iterator, Optional, Union
+
+import numpy as np
+import torch
+
+
+def resolve_device(device: Union[str, torch.device, None]) -> torch.device:
+    """``'auto'`` (or ``None``) picks CUDA when available, otherwise CPU."""
+    if device is None or device == "auto":
+        return torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    return torch.device(device)
+
+
+@contextlib.contextmanager
+def seeded(random_state: Optional[int],
+           device: torch.device) -> Iterator[None]:
+    """Seed torch inside a forked RNG so the caller's global state is untouched.
+
+    Every candidate in a hyperparameter search therefore starts from the same
+    initialisation, and repeated fits do not consume each other's draws.
+    """
+    if random_state is None:
+        yield
+        return
+    devices = []
+    if device.type == "cuda":
+        devices = [device.index if device.index is not None
+                   else torch.cuda.current_device()]
+    with torch.random.fork_rng(devices=devices):
+        torch.manual_seed(int(random_state))
+        yield
+
+
+def plain(obj: Any) -> Any:
+    """Recursively convert numpy scalars and arrays to built-in Python types.
+
+    ``torch.load(weights_only=True)`` accepts only built-in containers and
+    tensors, so everything saved alongside the state dict passes through here.
+    """
+    if isinstance(obj, dict):
+        return {str(k): plain(v) for k, v in obj.items()}
+    if isinstance(obj, (list, tuple)):
+        return type(obj)(plain(v) for v in obj)
+    if isinstance(obj, np.generic):
+        return obj.item()
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    if isinstance(obj, torch.device):
+        return str(obj)
+    return obj