Tensor

The core autograd engine. Wraps a numpy array, tracks every op into a graph, and computes gradients on .backward() via reverse topological sort.

x = Tensor([1.0, 2.0], requires_grad=True)
y = (x * 2).sum()
y.backward()
x.grad  # [2. 2.]

Arithmetic

# +
(Tensor([1.0,2.0]) + Tensor([3.0,4.0])).data   # [4. 6.]

# - / neg
(Tensor([5.0]) - Tensor([2.0])).data              # [3.]

# *
(Tensor([2.0,3.0]) * Tensor([4.0,5.0])).data   # [8. 15.]

# /
(Tensor([6.0]) / Tensor([2.0])).data              # [3.]

# **
(Tensor([3.0]) ** 2).data                      # [9.]

# matmul / @
x = Tensor(np.random.randn(2,3))
w = Tensor(np.random.randn(3,4))
(x.matmul(w)).shape   # (2, 4)
(x @ w).shape          # same, (2, 4)

Reduction

# sum
Tensor([[1.0,2.0],[3.0,4.0]]).sum(axis=1).data   # [3. 7.]

# mean
Tensor([2.0,4.0,6.0]).mean().data   # 4.0

# max
Tensor([1.0,5.0,3.0]).max().data   # 5.0

# argmax - no grad, inference only
Tensor([[0.1,0.7,0.2]]).argmax(axis=1)   # array([1])

Activation

# tanh
Tensor([0.0,1.0]).tanh().data

# relu
Tensor([-1.0,2.0]).relu().data   # [0. 2.]

# softmax - sums to 1 along axis
Tensor([[1.0,2.0,3.0]]).softmax(axis=-1).data

# gelu - smooth, nonzero grad for negative x
Tensor([-1.0,0.0,1.0]).gelu().data

Elementwise math

# log
Tensor([1.0, np.e]).log().data   # [0. 1.]

# sqrt
Tensor([4.0]).sqrt().data   # [2.]

# clip
Tensor([0.5,5.0,-1.0]).clip(0.0,1.0).data   # [0.5 1.  0. ]

# log_softmax - stable log(softmax(x))
Tensor([[1.0,2.0]]).log_softmax(axis=-1).data

Shape

# reshape / view
Tensor(np.zeros((2,3,4))).reshape(2,12).shape   # (2, 12)
Tensor(np.zeros((2,3,4))).view(2,12).shape     # same as reshape

# flatten
Tensor(np.zeros((2,3,4,4))).flatten(start_dim=1).shape   # (2, 48)

# transpose / permute
Tensor(np.zeros((2,3,4))).transpose(0,2,1).shape   # (2, 4, 3)
Tensor(np.zeros((2,3))).permute(1,0).shape         # (3, 2)

Indexing / gather-scatter

# __getitem__ - gather rows (used by Embedding)
w = Tensor(np.arange(9.0).reshape(3,3), requires_grad=True)
w[[0,2]].shape   # (2, 3)

# topk - returns (values, indices)
vals, idx = Tensor([[1.0,5.0,3.0,2.0]]).topk(2, axis=1)
vals.data   # [[5. 3.]]
idx         # [[1 2]]

# scatter - write src into copy of self at index positions
base = Tensor(np.zeros((3,4)), requires_grad=True)
idx = np.array([[0],[1],[2]])
src = Tensor(np.ones((3,1)), requires_grad=True)
base.scatter(1, idx, src)

# index_add - scatter-add (used to combine MOE expert outputs)
base2 = Tensor(np.zeros(3), requires_grad=True)
src2 = Tensor(np.array([1.0,2.0]), requires_grad=True)
base2.index_add(0, np.array([0,0]), src2).data   # [3. 0. 0.]

Combine / split

# cat - static method
a = Tensor(np.zeros((2,3)), requires_grad=True)
b = Tensor(np.zeros((2,3)), requires_grad=True)
Tensor.cat([a,b], axis=1).shape   # (2, 6)

# stack - static method, adds a new axis
Tensor.stack([a,b], axis=0).shape   # (2, 2, 3)

# split - inverse of cat
c = Tensor(np.zeros((2,6)), requires_grad=True)
parts = c.split(3, axis=1)   # two tensors, each (2, 3)

Sampling (no grad)

# multinomial - sample class indices from probabilities
Tensor([[0.0,1.0,0.0]]).multinomial(num_samples=1)   # array([[1]])

Convolution and pooling

# conv2d - (N,C,H,W) x (out_c,C,kh,kw)
x = Tensor(np.random.randn(1,3,8,8), requires_grad=True)
k = Tensor(np.random.randn(4,3,3,3), requires_grad=True)
x.conv2d(k, stride=1, padding=1).shape   # (1, 4, 8, 8)
x.conv2d(k, padding=1, pad_mode="reflect").shape   # same shape, mirrored border

# maxpool2d
Tensor(np.random.randn(1,2,4,4)).maxpool2d(2).shape   # (1, 2, 2, 2)

# avgpool2d
Tensor(np.ones((1,2,4,4))).avgpool2d(2).data   # all 1.0
Non-constant pad_mode values (reflect, replicate, circular) are exact in the forward pass, but backward crops the gradient rather than scattering it back through the padding rule - a documented approximation. Indexing ops use np.add.at for backward, so repeated indices (e.g. duplicate tokens in an embedding lookup) accumulate correctly instead of overwriting.

Layers

Every layer holds a params dict of Tensor(requires_grad=True). No manual backward needed - autograd handles it once you call .backward() on the loss.

Linear

layer = Linear(4, 2)
x = Tensor(np.random.randn(5,4))
layer.forward(x).shape   # (5, 2)

Tanh

Tanh().forward(Tensor([-1.0,0.0,1.0])).data

ReLU

ReLU().forward(Tensor([-1.0,2.0])).data   # [0. 2.]

GELU

GELU().forward(Tensor([-1.0,0.0,1.0])).data

Conv2D

conv = Conv2D(in_channels=3, out_channels=8,
              kernel_size=3, padding=1, pad_mode="reflect")
x = Tensor(np.random.randn(2,3,16,16), requires_grad=True)
conv.forward(x).shape   # (2, 8, 16, 16)

pad_mode accepts constant (default, exact gradient), reflect, replicate, or circular.

Embedding

emb = Embedding(vocab_size=1000, embed_dim=32)
emb.forward([1, 5, 9]).shape   # (3, 32) - idx is a plain list, not a Tensor

TiedLinear

emb = Embedding(vocab_size=1000, embed_dim=32)
lm_head = TiedLinear(emb)   # shares emb's weight, transposed

hidden = emb.forward([1, 5, 9])
lm_head.forward(hidden).shape   # (3, 1000)

Weight tying shares the same Tensor object - gradients from both uses accumulate into one param, and the optimizer only updates it once.

LayerNorm

ln = LayerNorm(dim=16)
x = Tensor(np.random.randn(4,16), requires_grad=True)
ln.forward(x).shape   # (4, 16)

BatchNorm2D

bn = BatchNorm2D(num_features=3)
x = Tensor(np.random.randn(4,3,8,8), requires_grad=True)
bn.forward(x)        # training mode - uses batch stats, updates running stats
bn.eval()
bn.forward(x)        # inference mode - uses running_mean / running_var
bn.train()

Dropout

do = Dropout(p=0.5)
x = Tensor(np.ones((10,10)))
do.forward(x)   # training - random mask, survivors scaled by 1/(1-p)
do.eval()
do.forward(x)   # identity
do.train()

MaxPool2D

MaxPool2D(kernel_size=2).forward(Tensor(np.random.randn(1,3,8,8))).shape   # (1, 3, 4, 4)

AvgPool2D

AvgPool2D(kernel_size=2).forward(Tensor(np.ones((1,2,4,4)))).data   # all 1.0

Flatten

Flatten().forward(Tensor(np.zeros((2,3,4,4)))).shape   # (2, 48)

GlobalAvgPool

GlobalAvgPool().forward(Tensor(np.ones((2,3,5,5)))).shape   # (2, 3)

Loss functions

All return a scalar Tensor - call .backward() on the result.

MSE

pred = Tensor([1.0,2.0], requires_grad=True)
actual = Tensor([0.0,0.0])
loss = MSE().loss(pred, actual)
loss.backward()
pred.grad   # [2. 4.]

BinaryCrossEntropy

pred = Tensor([0.9,0.1], requires_grad=True)
actual = Tensor([1.0,0.0])
loss = BinaryCrossEntropy().loss(pred, actual)
loss.backward()

NLLLoss

log_probs = Tensor(np.log([[0.7,0.3],[0.2,0.8]]), requires_grad=True)
targets = np.array([0,1])   # plain int array, not a Tensor
loss = NLLLoss().loss(log_probs, targets)
loss.backward()

CrossEntropyLoss

logits = Tensor(np.random.randn(8,10), requires_grad=True)
targets = np.array([3,0,7,7,1,9,2,5])
loss = CrossEntropyLoss().loss(logits, targets)   # raw logits, not probabilities
loss.backward()
CrossEntropyLoss runs log_softmax internally - pass raw logits. Use NLLLoss only if you already have log-probabilities.

Optimizers

SGD

optimizer = SGD(lr=0.01)
loss.backward()
optimizer.step(net)   # param.data -= lr * param.grad

RMSprop

optimizer = RMSprop(lr=0.001, decay=0.9)
loss.backward()
optimizer.step(net)   # divides grad by running avg of squared grad

Adam

optimizer = Adam(lr=0.001, beta1=0.9, beta2=0.999)
for step in range(100):
    net.zero_grad()
    loss = MSE().loss(net.forward(x), y)
    loss.backward()
    optimizer.step(net)   # momentum + adaptive scaling, bias-corrected

AdamW

optimizer = AdamW(lr=0.001, weight_decay=0.01)
loss.backward()
optimizer.step(net)   # Adam + decoupled weight decay

NeuralNet

net = NeuralNet([Linear(2, 8), Tanh(), Linear(8, 1)])

net.forward(x)
list(net.params())   # yields every layer's params, in order
net.zero_grad()
net.save("model.picodl")  
net.load("model.picodl")

Data

iterator = BatchIterator(batch_size=32, shuffle=True)
for batch in iterator(inputs, targets):
    batch.inputs, batch.targets  # raw numpy arrays

Train loop

train(net, inputs, targets,
      num_epochs=50,
      iterator=BatchIterator(batch_size=16),
      loss=MSE(),
      optimizer=Adam(lr=0.01))

net.save("model.picodl")

Accepts either a Tensor or raw numpy array for inputs/targets - unwrapped automatically before batching.

Targets are checked by dtype each batch: int-dtype targets (class indices, for CrossEntropyLoss / NLLLoss) are passed through raw; anything else (float targets, for MSE / BinaryCrossEntropy) gets wrapped in a Tensor. This matters because the classification losses do fancy indexing on the target array directly - passing a Tensor there instead of a raw int array breaks it.
# classification: y_train is an int array
train(net, X_train, y_train, loss=CrossEntropyLoss(), optimizer=AdamW())

# regression: y_train is a float array, wrapped in Tensor automatically
train(net, X_train, y_train, loss=MSE(), optimizer=Adam())