Needle

Needle is a compact deep-learning framework built around a reverse-mode automatic differentiation engine. It includes tensor operations, neural-network layers, optimizers, data utilities, and native ndarray backends for CPU and CUDA.

Features

• Reverse-mode automatic differentiation
• NumPy reference, native CPU, and optional CUDA ndarray devices
• Convolutional, recurrent, LSTM, attention, and Transformer layers
• SGD and Adam optimizers
• Dataset, data-loader, image-transform, and language-model utilities

Project layout

src/needle/    Python package
native/        C++ and CUDA ndarray backends
examples/      End-to-end model examples
tests/         Project-level unit and integration tests

Setup

Python 3.10 or newer, a C++ compiler, CMake, and pybind11 are required. CUDA is optional.

python3 -m pip install -e '.[dev]'
make build

To build without probing for a CUDA compiler:

cmake -S . -B build -DNEEDLE_BUILD_CUDA=OFF
cmake --build build --parallel

Usage

import needle as ndl
import needle.nn as nn

device = ndl.cpu()
x = ndl.Tensor([[1.0, 2.0], [3.0, 4.0]], device=device)
model = nn.Sequential(nn.Linear(2, 8, device=device), nn.ReLU(), nn.Linear(8, 1, device=device))
y = model(x)
y.sum().backward()

Use ndl.cpu_numpy() when a NumPy reference device is useful for debugging. The native CPU device is used by initializers unless another device is supplied.

Examples

The examples follow the model/loss/optimizer workflow commonly used in PyTorch tutorials, while running entirely on Needle:

PYTHONPATH=src:. python -m examples.linear_regression
PYTHONPATH=src:. python -m examples.mlp_classifier
PYTHONPATH=src:. python -m examples.cnn_classifier
PYTHONPATH=src:. python -m examples.sequence_models

• linear_regression demonstrates reverse-mode autograd and SGD.
• mlp_classifier trains a Sequential network with SoftmaxLoss and Adam.
• cnn_classifier runs a complete Conv/BatchNorm/ReLU training step.
• sequence_models compares RNN, LSTM, and Transformer output dimensions and verifies gradients through each architecture.

These scripts expose reusable functions and are also exercised by tests/test_examples.py, so the documentation examples cannot silently drift away from the public API.

Tests

make test

CUDA tests should only be enabled on a machine with a compatible CUDA driver and runtime.