Architecture

This document explains how Gorch’s automatic differentiation works and how to extend it. The design is intentionally simple: every differentiable operation stores its inputs and a hand-written local gradient rule, and backward() walks the resulting graph in reverse.

Package layout

Package	Responsibility
gorch/tensor/	Tensor, the reverse-mode graph, and all differentiable ops
gorch/nn/	Module, layers, activations, losses, functional helpers
gorch/optim/	Gradient optimizers, Kalman filter, RLS, Levenberg–Marquardt
gorch/utils/	Datasets and data loading

The autograd graph

Tensor

A Tensor wraps a NumPy array in .value and exposes:

• requires_grad: bool
• grad: Tensor | None — populated by .backward() (None initially)
• _grad_fn: Backward | None — the node that created this tensor (None for leaves)

*Backward nodes

Each differentiable operation creates a node, e.g. AddBackward(x, y), which:

1. Stores its inputs in self.input = [x, y, ...].
2. Implements backward(gradient) -> list, returning one gradient per input (None entries are allowed for inputs that don’t need a gradient).

When Tensor.backward() runs, it seeds the output with gradient 1, then walks _grad_fn backward, chaining gradients:

# tensor.py (simplified)
self.grad = gradient.transpose()
if self._grad_fn is not None:
    grads = self._grad_fn.backward(gradient)
    for input, g in zip(self._grad_fn.input, grads):
        if g is not None and input.requires_grad:
            input.backward(g)

The gradient convention (important)

Gradients are carried transposed through the graph:

• Tensor.backward stores self.grad = gradient.transpose(), so on a leaf the stored .grad.value matches the leaf’s .value shape.
• Inside a node’s backward(gradient), gradient is the upstream gradient in transposed space — its shape is the output shape reversed (output.shape[::-1]).
• The node must therefore return each input’s gradient in transposed space too, i.e. with shape input.shape[::-1].

Example — SqrtBackward:

class SqrtBackward:
    def __init__(self, tensor):
        self.input = [tensor]

    def backward(self, gradient):
        (tensor,) = self.input
        # dy/dx = 0.5 / sqrt(x); transpose so the result matches
        # the input's transposed shape, then wrap in a Tensor.
        grad_x = 0.5 / tensor.value.sqrt()
        return [gorch.Tensor((gradient.transpose().value * grad_x).transpose())]

Nodes that must deal with broadcasting (e.g. AddBackward, MulBackward) reduce the gradient with sum(axis=...) back onto the input’s shape — still in transposed space.

Adding a new operation

To add a differentiable operation foo, four places must be touched:

1. gorch/tensor/backward.py — add a FooBackward node implementing backward(gradient), following the transposed-gradient convention above, and add it to __all__.

2. gorch/tensor/methods.py — add a module-level foo(tensor, ...) that computes the forward NumPy result, builds a Tensor with requires_grad set, and attaches the node:

   def foo(tensor):
       if not isinstance(tensor, gorch.Tensor):
           raise ValueError("Input must be a Tensor")
       out = gorch.Tensor(np.foo(tensor.value), requires_grad=tensor.requires_grad)
       if tensor.requires_grad:
           out._grad_fn = FooBackward(tensor)
       return out
   

   Also add foo to methods.__all__ and — if it is a reduction with axis/keepdims — make sure np.* is used inside so Python builtins are not shadowed (methods.py deliberately redefines sum, max, min, …).

3. gorch/tensor/tensor.py — add a convenience method on Tensor that delegates to the module-level function.
4. gorch/tensor/__init__.py — add foo to __all__ so it is re-exported at gorch.foo.

Non-differentiable ops (e.g. diag, var, median, argmax) simply return requires_grad=False and never attach a node.

Then verify with a finite-difference test:

from tests._helpers import check_gradient, check_binary_gradient

check_gradient(gorch.foo, np.random.standard_normal((3, 4)))

nn and optim

• Module.parameters() / children() discover sub-modules via inspect.getmembers, so parameters and sub-modules must be plain instance attributes set in __init__ (this is how Sequential stores its layers as 0, 1, …).
• Optimizer subclasses implement step() and inherit zero_grad() from the base class; new optimizers must be exported from gorch/optim/__init__.py.
• The control tools use the functional helpers:
  • jacobian(model, x) returns d(model(x))/d(parameters), of shape (batch * out_dim, n_params).
  • vectorize_parameters / devectorize_parameters convert between a model and a flat parameter vector (devectorize_parameters adds its argument).

Limitations

• backward() only works on scalar outputs unless an explicit gradient Tensor is supplied.
• Grads are stored transposed during the pass (see above); the load-bearing convention is that a leaf’s .grad.value matches the leaf’s .value shape.
• No GPU support, no sparse tensors — this is a teaching library.