Skip to main content

A pure-NumPy deep learning framework with reverse-mode automatic differentiation

Project description

BareTensor

Python NumPy PyTorch Parity License PyPI version

Autograd engine and deep learning framework in pure Python/NumPy — verified against PyTorch's C++ backend to ≤ 1e⁻⁴.


Installation

pip install baretensor

Requires Python 3.10+ and NumPy ≥ 1.26. Zero other dependencies.


Quick Start

from baretensor import Tensor, Linear, Sequential, SGD

# Autograd
x = Tensor([[1.0, 2.0]], requires_grad=True)
w = Tensor([[0.5], [-0.3]], requires_grad=True)
y = x @ w
y.backward()
print(w.grad)  # exact analytical gradient

# Build models
model = Sequential(
    Linear(784, 256),
    Linear(256, 10),
)

What's Here

Autograd Engine

  • Dynamic DAG, topological sort, reverse-mode differentiation
  • Analytical Jacobians — Softmax Cross-Entropy, LayerNorm, Batched MatMul, GELU, MSE
  • Gradient un-broadcasting across batch axes
  • Full suite of activations: relu, sigmoid, tanh, gelu, softmax
  • Element-wise math: exp, log, +, -, *, /, @, negation

NN Modules

Module Description
Linear Fully-connected: y = xW + b
Conv2d im2col-based 2D convolution with stride/padding
MaxPool2d Max pooling with stride/padding
Dropout Inverted dropout regularization
BatchNorm1d Batch normalization with running statistics
LayerNorm Layer normalization with learnable γ, β
RMSNorm LLaMA/Mistral-style RMS normalization (γ only)
Embedding Token → dense vector lookup with scatter-add grad
MultiHeadAttention Multi-head scaled dot-product attention
TransformerEncoderBlock Attention + FFN + residual + LayerNorm
Sequential Chain modules in order

Functional Ops

Function Description
scaled_dot_product_attention(Q, K, V, mask) Attention with optional causal masking
layer_norm(x, gamma, beta, eps) Layer normalization
cross_entropy_loss(logits, targets) Fused softmax + NLL with analytical Jacobian
mse_loss(y_pred, y_true) Fused MSE with analytical Jacobian
rope(x, positions, base) Rotary Position Embeddings (RoPE)
cat(tensors, axis) Concatenate along axis

Optimizers

Optimizer Features
SGD Vanilla stochastic gradient descent
Adam Adaptive, bias-corrected, optional L2 weight decay
AdamW Decoupled weight decay (Loshchilov & Hutter 2019)
clip_grad_norm_(params, max_norm) Global L2 gradient clipping

LR Schedulers

Scheduler Description
StepLR(opt, step_size, gamma) Multiplicative decay every N epochs
CosineAnnealingLR(opt, T_max, eta_min) Cosine schedule with warm restart support

Data

Component Description
Dataset Abstract base class
TensorDataset(*tensors) Wrap numpy arrays / Tensors
DataLoader(dataset, batch_size, shuffle, drop_last) Minibatch iterator
Subset(dataset, indices) Index-based dataset view
random_split(dataset, lengths, seed) Deterministic train/val/test split

Infrastructure

Feature Description
Module.save(path) / Module.load(path) Model checkpointing (.npz)
Opt.state_dict() / Opt.load_state_dict() Optimizer checkpointing with index-stable keys
Module.register_forward_pre_hook(hook) Pre-forward hooks
Module.register_forward_hook(hook) Post-forward hooks
Module.train() / Module.eval() Training/eval mode toggle
Module.zero_grad() Zero all parameter gradients

Verified Against PyTorch — 35/35 Tests

Feature Test Parity
Linear + ReLU autograd test_linear_relu_autograd ≤ 1e⁻⁵
LayerNorm (2D, 3D, Module) test_layer_norm_* ≤ 1e⁻⁴
Multi-Head Attention test_mha_autograd_parity ≤ 1e⁻⁴
Softmax Cross-Entropy test_cross_entropy_parity ≤ 1e⁻⁴
Causal Masking test_causal_mask_parity ≤ 1e⁻⁴
Embedding scatter-add test_embedding_parity ≤ 1e⁻⁴
Batched MatMul test_batched_matmul_parity ≤ 1e⁻⁴
Reshape test_reshape_parity ≤ 1e⁻⁵
Dropout test_dropout_parity ≤ 1e⁻⁶
BatchNorm1d test_batchnorm1d_parity ≤ 1e⁻⁵
Negation test_neg_parity ≤ 1e⁻⁶
Division test_truediv_parity ≤ 1e⁻⁵
Sigmoid test_sigmoid_parity ≤ 1e⁻⁶
Tanh test_tanh_parity ≤ 1e⁻⁶
GELU (tanh approx) test_gelu_parity ≤ 1e⁻⁵
Exp test_exp_parity ≤ 1e⁻⁵
Log test_log_parity ≤ 1e⁻⁵
Sequential test_sequential_parity ≤ 1e⁻⁴
MSE Loss test_mse_loss_parity ≤ 1e⁻⁶
AdamW test_adamw_parity ≤ 1e⁻⁶
Gradient Clipping test_clip_grad_norm_parity ≤ 1e⁻⁶
Conv2d test_conv2d_parity ≤ 1e⁻⁴
MaxPool2d test_maxpool2d_parity ≤ 1e⁻⁶
RMSNorm test_rmsnorm_parity ≤ 1e⁻⁴
RoPE test_rope_parity ≤ 1e⁻⁵
StepLR test_step_lr_parity ≤ 1e⁻⁸
CosineAnnealingLR test_cosine_lr_parity ≤ 1e⁻⁶
Random Split test_random_split
Optimizer state_dict test_optimizer_state_dict
Forward Hooks test_forward_hooks

Architecture

class MicroGPT(Module):
    def __init__(self, vocab_size, d_model, num_heads):
        super().__init__()
        self.token_emb = Embedding(vocab_size, d_model)
        self.transformer = TransformerEncoderBlock(d_model, num_heads)
        self.lm_head = Linear(d_model, vocab_size)

    def forward(self, idx, mask=None):
        x = self.token_emb(idx)
        x = self.transformer(x, mask=mask)
        return self.lm_head(x)

See examples/ for full demos: XOR, MNIST (MLP + ConvNet), Transformer, Micro-GPT, Adam vs SGD, CartPole RL.


License

MIT — see LICENSE.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

baretensor-0.3.0.tar.gz (518.3 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

baretensor-0.3.0-py3-none-any.whl (17.7 kB view details)

Uploaded Python 3

File details

Details for the file baretensor-0.3.0.tar.gz.

File metadata

  • Download URL: baretensor-0.3.0.tar.gz
  • Upload date:
  • Size: 518.3 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.1

File hashes

Hashes for baretensor-0.3.0.tar.gz
Algorithm Hash digest
SHA256 e25e7725acda1aea847e74ef748610ce2b1ef23d43c9aeeab4a630fea7a6d7bb
MD5 6e8b37be6adc85c7d42d35878a15beaf
BLAKE2b-256 25bfbfb28fc11235125c7fa124e3af6397d304946ef5cbe29782e87fad21f8a8

See more details on using hashes here.

File details

Details for the file baretensor-0.3.0-py3-none-any.whl.

File metadata

  • Download URL: baretensor-0.3.0-py3-none-any.whl
  • Upload date:
  • Size: 17.7 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.1

File hashes

Hashes for baretensor-0.3.0-py3-none-any.whl
Algorithm Hash digest
SHA256 4b5506aa0297349b1805fe697d179140c55bad7938076dcdd9f58fdc7fd99aee
MD5 f99d361a84497b400199480b2f71e817
BLAKE2b-256 291a67290e6163c5364808d593c299e257af94e5fdff72aa68325744323dacd1

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page