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 PyPI downloads License Python GitHub status DOI

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.

CNN Example (MNIST)

from pureml.base import NN
from pureml.layers import Conv2D, MaxPool2D, Affine, BatchNorm2d, Dropout, Dropout2d, output_shape_2d
from pureml.activations import relu
from pureml import Tensor, no_grad
from pureml.training_utils import DataLoader
from pureml.datasets.MNIST import MnistDataset
from pureml.losses import CCE
from pureml.optimizers import SGD
from pureml.evaluation import accuracy


class ConvNet(NN):

    def __init__(self, in_features: int, img_res: tuple[int, int]):
        # Conv+Pool+BN Layer 1:
        H, W = img_res
        conv_kwargs = dict(kernel_size=2, stride=1, padding=0, dilation=1)
        self.conv1 = Conv2D(in_channels=in_features, out_channels=32, **conv_kwargs)
        H, W = output_shape_2d(H_in=H, W_in=W, **conv_kwargs)
        conv_kwargs["stride"] = 2
        conv_kwargs["kernel_size"] = 3
        self.pool1 = MaxPool2D(**conv_kwargs)
        H, W = output_shape_2d(H_in=H, W_in=W, **conv_kwargs)
        self.bn1 = BatchNorm2d(32)
        # Conv+Pool+BN Layer 2:
        conv_kwargs = dict(kernel_size=2, stride=1, padding=0, dilation=1)
        self.conv2 = Conv2D(in_channels=32, out_channels=64, **conv_kwargs)
        H, W = output_shape_2d(H_in=H, W_in=W, **conv_kwargs)
        conv_kwargs["stride"] = 2
        conv_kwargs["kernel_size"] = 3
        self.pool2 = MaxPool2D(**conv_kwargs)
        H, W = output_shape_2d(H_in=H, W_in=W, **conv_kwargs)
        self.bn2 = BatchNorm2d(64)
        # Conv+BN Layer 3:
        conv_kwargs = dict(kernel_size=2, stride=1, padding=0, dilation=1)
        self.conv3 = Conv2D(in_channels=64, out_channels=128, **conv_kwargs)
        H, W = output_shape_2d(H_in=H, W_in=W, **conv_kwargs)
        self.bn3 = BatchNorm2d(128)
        self.dropout2d = Dropout2d(p=0.1)
        # FC Layer 4:
        self.fc1 = Affine(128*H*W, 320, bias=False)
        # Dropout layer (fully connected block):
        self.dropout = Dropout(p=0.25)
        # FC Layer 5:
        self.fc2 = Affine(320, 10, bias=False)
    
    def predict(self, X: Tensor):
        X = X.unsqueeze(1)
        X = self.conv1(X)
        X = relu(X)
        X = self.pool1(X)
        X = self.bn1(X)
        # ---------------
        X = self.conv2(X)
        X = relu(X)
        X = self.pool2(X)
        X = self.bn2(X)
        # ---------------
        X = self.conv3(X)
        X = relu(X)
        X = self.bn3(X)
        X = self.dropout2d(X)
        # ---------------
        X = X.flatten()
        X = self.fc1(X)
        X = self.dropout(X)
        X = self.fc2(X)
        # ---------------
        if self.training:
            return X
        return X.argmax(axis=-1)
    

if __name__ == "__main__":
    model = ConvNet(1, (28, 28))
    optim = SGD(model.parameters, lr=0.01, beta=0.9, weight_decay=0.0001)
    with MnistDataset("train") as train:
        for epoch in range(1, 4):
            model.train()
            for X, Y in DataLoader(train, batch_size=64, shuffle=True):
                logits = model(X)
                loss = CCE(Y, logits, from_logits=True)
                optim.zero_grad()
                loss.backward()
                optim.step()
            model.eval()
            with no_grad():
                with MnistDataset("test") as test:
                    print(f"Accuracy after epoch {epoch} is {accuracy(model, test, batch_size=64)}")
    model.save_state("convnet_mnist_demo")
    print("Saved checkpoint to convnet_mnist_demo.pureml.zip")

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

Squeeze and Unsqueeze (Tensor)

Use these when you need to explicitly control singleton dimensions (size == 1) without breaking autograd.

When to use:

Examples:

import numpy as np
from pureml import Tensor

# Remove singleton dims
logits = Tensor(np.random.randn(32, 10, 1), requires_grad=True)
logits2d = logits.squeeze(-1)      # (32, 10)

# Add channel dim
images = Tensor(np.random.randn(32, 28, 28), requires_grad=True)
images_nchw = images.unsqueeze(1)  # (32, 1, 28, 28)

# Gradients flow through both ops
out = (images_nchw * images_nchw).reshape(32, -1)
out.backward()

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)

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.

Unfold helpers

Output-shape helpers

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)

Contributing and Protocols

For development workflow, architecture contracts, protocol rules, and release/branching guidance, see CONTRIBUTING.md in the repository root.