PureML

/PUREML-DOCS
Docs Home Main Site

PureML Documentation

LOGO

Tiny but powerful 100% NumPy-based deep learning framework with explicit autodiff and lightweight utilities.

PyPI License Python GitHub status

Installation

PureML targets Python 3.11+ with NumPy 2.x and zarr. PyPI: https://pypi.org/project/ym-pure-ml/

pip install ym-pure-ml

Quickstart (MNIST)

from pureml import Tensor
from pureml.activations import relu
from pureml.layers import Affine
from pureml.base import NN
from pureml.datasets import MnistDataset
from pureml.optimizers import Adam
from pureml.losses import CCE
from pureml.training_utils import DataLoader
from pureml.evaluation import accuracy
import time

class MNIST_BEATER(NN):

    def __init__(self) -> None:
        self.L1 = Affine(28*28, 256)
        self.L2 = Affine(256, 10)

    def predict(self, x: Tensor) -> Tensor:
        x = x.flatten(sample_ndim=2) # passing 2 because imgs in MNIST are 2D
        x = relu(self.L1(x))
        x = self.L2(x)
        if self.training:
            return x
        return x.argmax(axis=x.ndim-1) # argmax over the feature dim

with MnistDataset("train") as train, MnistDataset("test") as test:
    model = MNIST_BEATER().train()
    opt = Adam(model.parameters, lr=1e-3, weight_decay=1e-2)
    start_time = time.perf_counter()
    for _ in range(5):
        for X, Y in DataLoader(train, batch_size=128, shuffle=True):
            opt.zero_grad()
            logits = model(X)
            loss = CCE(Y, logits, from_logits=True)
            loss.backward()
            opt.step()
    end_time = time.perf_counter()
    model.eval()
    acc = accuracy(model, test, batch_size=1024)
print("Time taken: ", end_time - start_time, " sec.")
print(f"Test accuracy: {acc * 100}")

MnistDataset("train") yields (Tensor image, one_hot Tensor label), normalized to [0, 1]. In eval mode MNIST_BEATER.predict returns class indices; in train mode it returns logits for loss computation.

Tensors and Autograd

Creating tensors

from pureml import Tensor, is_grad_enabled, no_grad
from pureml.general_math import sum as tensor_sum

x = Tensor([[1, 2], [3, 4]], requires_grad=True)
y = Tensor(5)                      # requires_grad=False by default

Gradient flow

out = tensor_sum(x * Tensor(2))
out.backward()        # seeds ones_like when grad is omitted
print(x.grad)         # accumulated gradient
x.zero_grad()         # sets grad to None

Built-in ops

Defining custom ops

TensorValuedFunction(forward_fn, grad_fn) builds a differentiable node. Both functions may accept a keyword-only context dict to reuse cached intermediates. If grad_fn is omitted, calling it raises GradientNotDefined.

Activations (pureml.activations)

Each returns a Tensor and has a Jacobian-free backward pass.

Losses (pureml.losses)

All losses return scalar tensors (mean over all elements/samples).

Layers (pureml.layers)

Common interface: .parameters (trainables), .named_buffers() (non-trainable state), .train()/.eval() toggle training and call on_mode_change.

General Math (pureml.general_math)

All return Tensors with gradient support and honor the specified axis (negative axes allowed). Defaults: mean reduces over all elements when axis=None; variance/std/sum default to axis=-1.

Data utilities (pureml.training_utils)

Dataset abstractions

DataLoader

from pureml.training_utils import DataLoader
loader = DataLoader(dataset, batch_size=32, shuffle=True, drop_last=False)
for batch in loader:
    ...

Collation helpers

Encoding helpers

Models and state management (pureml.base)

Mode propagation matters for layers like Dropout/BatchNorm and for MNIST_BEATER.predict, which changes its return type depending on mode.

Evaluation (pureml.evaluation)

Optimizers and Schedulers (pureml.optimizers)

Usage pattern:

opt = SGD(model.parameters, lr=0.1, beta=0.9, weight_decay=1e-4, decoupled_wd=True)
...
loss.backward()
opt.step()
opt.zero_grad()

Common features:

Optimizers:

Learning-rate schedulers (operate in-place on attached optimizer’s lr):

Persistence backend (pureml.util.ArrayStorage)

ArrayStorage wraps a Zarr v3 root group for storing multiple appendable arrays with metadata. Backends:

Key methods:

Other utilities in pureml.util:

Logging (pureml.logging_util)

configure_logging(logs_dir, file_level=logging.DEBUG, console_level=logging.WARNING) sets up the root logger once with:

Returns immediately if root already has handlers.

Built-in datasets and models (pureml.datasets, pureml.models)

MnistDataset

from pureml.datasets import MnistDataset
train = MnistDataset("train")   # (Tensor image, one_hot label)
test  = MnistDataset("test")    # (Tensor image, class index)

MNIST_BEATER (neural network)

KNN (classical)