newch 0.2.0__py3-none-win_amd64.whl

newch/__init__.py

@@ -0,0 +1,29 @@
+"""newch: a small machine-learning library for low-spec machines.
+
+PyTorch-style tensors and autograd, Keras-style ``compile``/``fit``, backed
+by a portable C engine. No Python dependencies (numpy is used if present).
+
+    import newch as nc
+
+    model = nc.Sequential([nc.net.Linear(2, 32), nc.net.ReLU(), nc.net.Linear(32, 3)])
+    model.compile(optimizer="adam", loss="cross_entropy", metrics=["accuracy"])
+    model.fit(x, y, epochs=20, batch_size=32)
+"""
+
+from . import callbacks, data, net, optim
+from .net import Model, Module, Sequential, quantize_
+from .tensor import (Context, NcError, Tensor, avg_pool2d, batch_norm, conv2d, cross_entropy, default_context,
+                     dropout, exp, float32, full, int8, int32, linear, load_tensors, log, log_softmax, manual_seed,
+                     matmul, max_pool2d, mem_stats, mse_loss, no_grad, ones, parameter, persistent, rand, randn,
+                     release, relu, save_tensors, set_default_context, sigmoid, softmax, sqrt, step_end, tanh, tensor,
+                     zeros)
+
+__version__ = "0.2.0"
+
+__all__ = [
+    "Context", "Model", "Module", "NcError", "Sequential", "Tensor", "avg_pool2d", "batch_norm", "callbacks",
+    "conv2d", "cross_entropy", "data", "default_context", "dropout", "exp", "float32", "full", "int8", "int32",
+    "linear", "load_tensors", "log", "log_softmax", "manual_seed", "matmul", "max_pool2d", "mem_stats", "mse_loss",
+    "net", "no_grad", "ones", "optim", "parameter", "persistent", "quantize_", "rand", "randn", "release", "relu",
+    "save_tensors", "set_default_context", "sigmoid", "softmax", "sqrt", "step_end", "tanh", "tensor", "zeros",
+]

newch/_lib.py

@@ -0,0 +1,168 @@
+"""ctypes bindings to the C engine (nc.h). Internal: use the public API."""
+
+from __future__ import annotations
+
+import ctypes
+import os
+import sys
+from ctypes import POINTER, Structure, c_char_p, c_float, c_int, c_int64, c_size_t, c_uint64, c_void_p
+from pathlib import Path
+
+ALL_AXES = -128
+F32, I32, I8 = 0, 1, 3
+DTYPE_NAMES = {F32: "float32", I32: "int32", I8: "int8"}
+DTYPE_CODES = {"float32": F32, "int32": I32, "int8": I8}
+DTYPE_SIZE = {"float32": 4, "int32": 4, "int8": 1}
+ARRAY_CODE = {"float32": "f", "int32": "i", "int8": "b"}
+
+
+class Config(Structure):
+    _fields_ = [("arena_chunk_bytes", c_size_t), ("seed", c_uint64), ("force_reference", c_int), ("threads", c_int)]
+
+
+class CpuInfo(Structure):
+    _fields_ = [
+        ("arch", c_char_p),
+        ("has_neon", c_int),
+        ("has_dotprod", c_int),
+        ("has_fp16", c_int),
+        ("has_avx2", c_int),
+        ("has_fma", c_int),
+        ("has_rvv", c_int),
+        ("n_cores", c_int),
+        ("n_big_cores", c_int),
+    ]
+
+
+class MemStats(Structure):
+    _fields_ = [
+        ("persistent_bytes", c_size_t),
+        ("arena_bytes_in_use", c_size_t),
+        ("arena_bytes_reserved", c_size_t),
+        ("arena_peak_bytes", c_size_t),
+    ]
+
+
+class OptimOpts(Structure):
+    _fields_ = [(n, c_float) for n in ("lr", "momentum", "weight_decay", "beta1", "beta2", "eps")]
+
+
+def _lib_names() -> list[str]:
+    if sys.platform == "win32":
+        return ["nc.dll"]
+    if sys.platform == "darwin":
+        return ["libnc.dylib"]
+    return ["libnc.so"]
+
+
+def _find_library() -> str:
+    env = os.environ.get("NEWCH_LIB")
+    if env:
+        return env
+    here = Path(__file__).resolve().parent
+    for d in (here, here.parent / "build" / "host"):
+        for name in _lib_names():
+            if (d / name).exists():
+                return str(d / name)
+    raise OSError(
+        "newch: native library not found. Build it with `python tools/build.py` "
+        "or set NEWCH_LIB to the path of nc.dll / libnc.so."
+    )
+
+
+lib = ctypes.CDLL(_find_library())
+
+P = c_void_p
+I64P = POINTER(c_int64)
+
+_SIGNATURES = {
+    # context
+    "nc_init": (P, [POINTER(Config)]),
+    "nc_free_ctx": (None, [P]),
+    "nc_manual_seed": (None, [P, c_uint64]),
+    "nc_linear_fn": (P, [P, P, P]),
+    "nc_step_end": (None, [P]),
+    "nc_arena_trim": (None, [P]),
+    "nc_num_threads": (c_int, [P]),
+    "nc_persistent_begin": (None, [P]),
+    "nc_persistent_end": (None, [P]),
+    "nc_set_grad_enabled": (c_int, [P, c_int]),
+    "nc_ok": (c_int, [P]),
+    "nc_status_code": (c_int, [P]),
+    "nc_error": (c_char_p, [P]),
+    "nc_clear_error": (None, [P]),
+    "nc_get_mem_stats": (MemStats, [P]),
+    "nc_reset_peak": (None, [P]),
+    "nc_cpu": (POINTER(CpuInfo), [P]),
+    "nc_kernel_backend": (c_char_p, [P]),
+    # tensors
+    "nc_empty": (P, [P, c_int, c_int, I64P]),
+    "nc_zeros": (P, [P, c_int, c_int, I64P]),
+    "nc_full": (P, [P, c_float, c_int, I64P]),
+    "nc_randn": (P, [P, c_float, c_int, I64P]),
+    "nc_rand_uniform": (P, [P, c_float, c_float, c_int, I64P]),
+    "nc_from_data": (P, [P, c_int, c_void_p, c_int, I64P]),
+    "nc_scalar": (P, [P, c_float]),
+    "nc_parameter": (P, [P, P, c_int, I64P]),
+    "nc_ndim": (c_int, [P]),
+    "nc_shape": (I64P, [P]),
+    "nc_numel": (c_int64, [P]),
+    "nc_get_dtype": (c_int, [P]),
+    "nc_is_contiguous": (c_int, [P]),
+    "nc_requires_grad": (c_int, [P]),
+    "nc_set_requires_grad": (None, [P, c_int]),
+    "nc_is_persistent": (c_int, [P]),
+    "nc_release": (None, [P]),
+    "nc_grad": (P, [P]),
+    "nc_item": (c_float, [P]),
+    "nc_data_f32": (c_void_p, [P]),
+    "nc_copy_to": (c_int, [P, c_void_p, c_size_t]),
+    "nc_copy_from": (c_int, [P, c_void_p, c_size_t]),
+    # ops
+    **{f"nc_{n}": (P, [P, P]) for n in ("add", "sub", "mul", "div", "matmul", "cross_entropy", "mse_loss",
+                                         "index_select")},
+    **{f"nc_{n}": (P, [P]) for n in ("neg", "exp", "log", "sqrt", "relu", "sigmoid", "tanh", "contiguous",
+                                      "log_softmax", "softmax", "detach")},
+    **{f"nc_{n}": (P, [P, c_float]) for n in ("scale", "add_scalar", "pow_scalar")},
+    **{f"nc_{n}": (P, [P, c_int, c_int]) for n in ("sum", "mean", "max", "transpose")},
+    "nc_argmax": (P, [P, c_int]),
+    "nc_conv2d": (P, [P, P, P, c_int, c_int, c_int, c_int]),
+    "nc_max_pool2d": (P, [P, c_int, c_int, c_int, c_int, c_int, c_int]),
+    "nc_avg_pool2d": (P, [P, c_int, c_int, c_int, c_int, c_int, c_int]),
+    "nc_dropout": (P, [P, c_float]),
+    "nc_batch_norm": (P, [P, P, P, P, P, c_int, c_float, c_float]),
+    "nc_quantize_rows": (P, [P, POINTER(c_void_p)]),
+    "nc_linear_int8": (P, [P, P, P, P]),
+    "nc_save": (c_int, [P, c_char_p, c_int, POINTER(c_char_p), POINTER(c_void_p)]),
+    "nc_load": (c_int, [P, c_char_p, c_int, POINTER(c_char_p), POINTER(c_void_p)]),
+    "nc_reshape": (P, [P, c_int, I64P]),
+    # autograd
+    "nc_backward": (c_int, [P]),
+    "nc_zero_grad": (None, [P]),
+    # optimizers
+    "nc_sgd": (P, [P, POINTER(c_void_p), c_int, OptimOpts]),
+    "nc_adam": (P, [P, POINTER(c_void_p), c_int, OptimOpts]),
+    "nc_optim_zero_grad": (None, [P]),
+    "nc_optim_step": (None, [P]),
+    "nc_optim_set_lr": (None, [P, c_float]),
+}
+
+for _name, (_res, _args) in _SIGNATURES.items():
+    _fn = getattr(lib, _name)
+    _fn.restype = _res
+    _fn.argtypes = _args
+
+
+def shape_arg(shape) -> tuple[int, ctypes.Array]:
+    shape = tuple(int(s) for s in shape)
+    return len(shape), (c_int64 * max(len(shape), 1))(*shape)
+
+
+def ptr_array(ptrs) -> ctypes.Array:
+    ptrs = list(ptrs)
+    return (c_void_p * max(len(ptrs), 1))(*ptrs)
+
+
+def name_array(names) -> ctypes.Array:
+    names = [n.encode("utf-8") for n in names]
+    return (c_char_p * max(len(names), 1))(*names)

newch/callbacks.py

@@ -0,0 +1,91 @@
+"""Keras-style training callbacks."""
+
+from __future__ import annotations
+
+import ctypes
+import math
+
+
+class Callback:
+    model = None
+
+    def set_model(self, model) -> None:
+        self.model = model
+
+    def on_train_begin(self, logs: dict) -> None: ...
+    def on_epoch_end(self, epoch: int, logs: dict) -> None: ...
+    def on_train_end(self, logs: dict) -> None: ...
+
+
+class History(Callback):
+    """Returned by fit(): ``history.history["loss"]`` is one value per epoch."""
+
+    def on_train_begin(self, logs):
+        self.epoch: list[int] = []
+        self.history: dict[str, list[float]] = {}
+
+    def on_epoch_end(self, epoch, logs):
+        self.epoch.append(epoch)
+        for k, v in logs.items():
+            self.history.setdefault(k, []).append(v)
+
+
+class EarlyStopping(Callback):
+    """Stops when ``monitor`` has not improved by ``min_delta`` for ``patience`` epochs."""
+
+    def __init__(self, monitor: str = "val_loss", min_delta: float = 0.0, patience: int = 0,
+                 mode: str = "auto", restore_best_weights: bool = False):
+        self.monitor, self.min_delta, self.patience = monitor, abs(min_delta), patience
+        if mode == "auto":
+            mode = "max" if "acc" in monitor else "min"
+        if mode not in ("min", "max"):
+            raise ValueError("mode must be 'auto', 'min' or 'max'")
+        self.sign = 1.0 if mode == "min" else -1.0
+        self.restore_best_weights = restore_best_weights
+        self.stopped_epoch = None
+
+    def on_train_begin(self, logs):
+        self.best = math.inf
+        self.wait = 0
+        self.best_weights = None
+        self.stopped_epoch = None
+
+    def on_epoch_end(self, epoch, logs):
+        if self.monitor not in logs:
+            raise KeyError(f"EarlyStopping: '{self.monitor}' not in logs (available: {sorted(logs)})")
+        value = self.sign * logs[self.monitor]
+        if value < self.best - self.min_delta:
+            self.best, self.wait = value, 0
+            if self.restore_best_weights:
+                self.best_weights = [p._flat() for p in self.model.parameters()]
+            return
+        self.wait += 1
+        if self.wait >= self.patience:
+            self.model.stop_training = True
+            self.stopped_epoch = epoch
+
+    def on_train_end(self, logs):
+        if self.restore_best_weights and self.best_weights is not None:
+            from ._lib import lib
+
+            for p, saved in zip(self.model.parameters(), self.best_weights):
+                ctypes.memmove(lib.nc_data_f32(p.ptr), saved.buffer_info()[0], len(saved) * 4)
+
+
+class LambdaCallback(Callback):
+    """Keyword-only, with Keras' argument names: on_epoch_end(epoch, logs), ..."""
+
+    def __init__(self, *, on_epoch_end=None, on_train_begin=None, on_train_end=None):
+        self._e, self._b, self._t = on_epoch_end, on_train_begin, on_train_end
+
+    def on_train_begin(self, logs):
+        if self._b:
+            self._b(logs)
+
+    def on_epoch_end(self, epoch, logs):
+        if self._e:
+            self._e(epoch, logs)
+
+    def on_train_end(self, logs):
+        if self._t:
+            self._t(logs)

newch/data.py

@@ -0,0 +1,227 @@
+"""Datasets and a streaming DataLoader (torch.utils.data style).
+
+For data larger than RAM, store samples in a raw binary file and open it
+with FileArray: rows are read through mmap one batch at a time, so memory use
+stays around one batch regardless of the dataset size.
+
+    FileArray.create("x.bin", rows, row_shape=(28, 28))     # write once
+    ds = TensorDataset(FileArray("x.bin", (60000, 28, 28)), labels)
+    for xb, yb in DataLoader(ds, batch_size=64, shuffle=True): ...
+"""
+
+from __future__ import annotations
+
+import array
+import mmap
+import random
+from pathlib import Path
+
+from . import _lib as L
+from .tensor import (Context, Tensor, _flatten, _is_numpy, default_context, float32, int8, int32, tensor,
+                     tensor_from_bytes)
+
+
+class Dataset:
+    """Map-style dataset: implement ``__len__`` and ``__getitem__``."""
+
+    def __len__(self) -> int:
+        raise NotImplementedError
+
+    def __getitem__(self, index):
+        raise NotImplementedError
+
+
+# --------------------------------------------------------------------------
+# Row sources: array-likes that can hand out a batch of rows as raw bytes
+# --------------------------------------------------------------------------
+
+
+class FileArray:
+    """A read-only array stored row-major in a raw binary file (like numpy.memmap).
+
+    shape includes the number of rows; dtype is "float32", "int32" or "int8".
+    """
+
+    def __init__(self, path, shape, dtype: str = float32, offset: int = 0):
+        if dtype not in L.DTYPE_CODES:
+            raise ValueError(f"unsupported dtype {dtype!r}")
+        self.path, self.shape, self.dtype, self.offset = Path(path), tuple(int(s) for s in shape), dtype, offset
+        self.row_shape = self.shape[1:]
+        n = 1
+        for d in self.row_shape:
+            n *= d
+        self.row_bytes = n * L.DTYPE_SIZE[dtype]
+        need = offset + self.shape[0] * self.row_bytes
+        size = self.path.stat().st_size
+        if size < need:
+            raise ValueError(f"{path}: {size} bytes, but shape {self.shape} {dtype} needs {need}")
+        self._f = open(self.path, "rb")
+        self._mm = mmap.mmap(self._f.fileno(), 0, access=mmap.ACCESS_READ) if need else None
+
+    @classmethod
+    def create(cls, path, rows, row_shape, dtype: str = float32) -> "FileArray":
+        """Writes an iterable of rows (sequences / numpy rows) without holding them all in memory."""
+        code = L.ARRAY_CODE[dtype]
+        count = 0
+        with open(path, "wb") as f:
+            for row in rows:
+                if _is_numpy(row):
+                    import numpy as np
+
+                    f.write(np.ascontiguousarray(row, dtype={float32: np.float32, int32: np.int32,
+                                                             int8: np.int8}[dtype]).tobytes())
+                else:
+                    flat: list = []
+                    _flatten(row, [], flat) if isinstance(row, (list, tuple)) else flat.append(row)
+                    f.write(array.array(code, flat if dtype == float32 else map(int, flat)).tobytes())
+                count += 1
+        return cls(path, (count, *row_shape), dtype)
+
+    def __len__(self) -> int:
+        return self.shape[0]
+
+    def rows(self, indices) -> bytes:
+        rb, off, mm = self.row_bytes, self.offset, self._mm
+        if not indices:
+            return b""
+        first = indices[0]
+        if indices == list(range(first, first + len(indices))):  # contiguous: one slice
+            return mm[off + first * rb: off + (first + len(indices)) * rb]
+        return b"".join(mm[off + i * rb: off + (i + 1) * rb] for i in indices)
+
+    def __getitem__(self, i: int):
+        return array.array(L.ARRAY_CODE[self.dtype], self.rows([i % len(self)])).tolist()
+
+    def close(self) -> None:
+        if self._mm is not None:
+            self._mm.close()
+            self._mm = None
+        self._f.close()
+
+
+class _MemoryRows:
+    """In-memory row source built once from lists, array.array, numpy or a Tensor."""
+
+    def __init__(self, data):
+        if isinstance(data, Tensor):
+            self.dtype, shape = data.dtype, data.shape
+            flat = data._flat()
+            self.buf = flat.tobytes()
+        elif _is_numpy(data):
+            import numpy as np
+
+            kind = int32 if np.issubdtype(data.dtype, np.integer) or data.dtype == np.bool_ else float32
+            arr = np.ascontiguousarray(data, dtype=np.int32 if kind == int32 else np.float32)
+            self.dtype, shape, self.buf = kind, arr.shape, arr.tobytes()
+        else:
+            shape_l: list = []
+            flat_l: list = []
+            _flatten(list(data), shape_l, flat_l)
+            self.dtype = float32 if any(isinstance(v, float) for v in flat_l) else int32
+            self.buf = array.array(L.ARRAY_CODE[self.dtype],
+                                   flat_l if self.dtype == float32 else map(int, flat_l)).tobytes()
+            shape = tuple(shape_l)
+        if not shape:
+            raise ValueError("dataset arrays need at least one dimension (the sample axis)")
+        self.shape = tuple(shape)
+        self.row_shape = self.shape[1:]
+        n = 1
+        for d in self.row_shape:
+            n *= d
+        self.row_bytes = n * L.DTYPE_SIZE[self.dtype]
+        self._mv = memoryview(self.buf)
+
+    def __len__(self) -> int:
+        return self.shape[0]
+
+    def rows(self, indices) -> bytes:
+        rb, mv = self.row_bytes, self._mv
+        if not indices:
+            return b""
+        first = indices[0]
+        if indices == list(range(first, first + len(indices))):
+            return bytes(mv[first * rb:(first + len(indices)) * rb])
+        return b"".join(mv[i * rb:(i + 1) * rb] for i in indices)
+
+    def __getitem__(self, i: int):
+        return array.array(L.ARRAY_CODE[self.dtype], self.rows([i % len(self)])).tolist()
+
+
+def _row_source(data):
+    return data if isinstance(data, FileArray) else _MemoryRows(data)
+
+
+class TensorDataset(Dataset):
+    """Samples are rows of each array (torch.utils.data.TensorDataset).
+
+    Arrays may be nested lists, array.array, numpy arrays, Tensors or FileArray
+    (disk-backed). Integer data becomes int32, floating data float32.
+    """
+
+    def __init__(self, *arrays):
+        if not arrays:
+            raise ValueError("TensorDataset needs at least one array")
+        self.sources = [_row_source(a) for a in arrays]
+        n = len(self.sources[0])
+        if any(len(s) != n for s in self.sources):
+            raise ValueError("all arrays must have the same number of rows")
+
+    def __len__(self) -> int:
+        return len(self.sources[0])
+
+    def __getitem__(self, index):
+        return tuple(s[index] for s in self.sources)
+
+    def get_batch(self, indices, ctx: Context) -> list[Tensor]:
+        return [tensor_from_bytes(s.rows(indices), (len(indices), *s.row_shape), s.dtype, ctx)
+                for s in self.sources]
+
+
+class DataLoader:
+    """Iterates a dataset in batches (torch.utils.data.DataLoader).
+
+    Yields a list with one temporary Tensor per field, valid until the next
+    ``step_end()``. TensorDataset batches are gathered directly as bytes; other
+    datasets are collated from ``__getitem__``.
+    """
+
+    def __init__(self, dataset, batch_size: int = 1, shuffle: bool = False, drop_last: bool = False,
+                 generator: random.Random | None = None, ctx: Context | None = None):
+        if batch_size < 1:
+            raise ValueError("batch_size must be >= 1")
+        self.dataset, self.batch_size, self.shuffle, self.drop_last = dataset, batch_size, shuffle, drop_last
+        self.generator = generator or random.Random()
+        self._ctx = ctx
+
+    def __len__(self) -> int:
+        n = len(self.dataset)
+        return n // self.batch_size if self.drop_last else -(-n // self.batch_size)
+
+    def __iter__(self):
+        ctx = self._ctx or default_context()
+        order = list(range(len(self.dataset)))
+        if self.shuffle:
+            self.generator.shuffle(order)
+        bs = self.batch_size
+        stop = len(order) - (len(order) % bs if self.drop_last else 0)
+        for s in range(0, stop, bs):
+            idx = order[s:s + bs]
+            if hasattr(self.dataset, "get_batch"):
+                yield self.dataset.get_batch(idx, ctx)
+            else:
+                yield _collate([self.dataset[i] for i in idx], ctx)
+
+
+def _stack(items, ctx: Context) -> Tensor:
+    if _is_numpy(items[0]):
+        import numpy as np
+
+        return tensor(np.stack(items), ctx=ctx)
+    return tensor(list(items), ctx=ctx)
+
+
+def _collate(samples, ctx: Context):
+    """Tuple samples -> list of batched fields (as torch's default_collate); else one Tensor."""
+    if isinstance(samples[0], tuple):
+        return [_stack([s[k] for s in samples], ctx) for k in range(len(samples[0]))]
+    return _stack(samples, ctx)