A lightweight automatic differentiation library with PyTorch-like API, supporting higher-order gradients and eager memory management. Designed as a drop-in replacement for NumPy with readability and clarity as a top priority.
This codebase does assume the reader has some existing requisite knowledge about multivariable calculus and linear algebra, but it's meant to be as readable as possible with that in mind. You can find great resources for both of those topics pretty much anywhere on the internet!
Finding partial derivatives for almost any function is pretty simple, just:
Define some variables
x = md.Tensor([[0, 2, -2, 1], [-1, -1, -2, -2]], allow_grad=True)
y = md.Tensor([[2, 3, 4, 5], [0, -1, -3, 2]], allow_grad=True)
Run a calculation
f = 2 * y * md.sin(x) - x**2
And call backward()
f.backward(allow_higher_order=True)
# first order partial derivatives
# df/dx
print(x.grad)
# df/dy
print(y.grad)
If you want to go a little deeper and compute higher order partial derivatives, it's as simple as just doing a backward pass on the gradient tensors
x.grad.backward()
print("second order")
# now d^2f/dx^2
print(x.grad)
# now d^2f/dxdy
print(y.grad)
Implementing your own differentiable operations is also pretty simple.
Helper functions like create_op_func() and as_minidiff() automatically wrap backend functions as minidiff functions and manage computation graphs, respectively.
See minidiff/ops/definitions.py for some examples.
Feel free to explore around the repo, it's structured as follows:
minidiff/tensor.py: Contains all the basic tensor code which includes a few convenient tensor creation functions and important computational graph properties
minidiff/topology.py: The OpNode class is defined here; it represents some function performed on tensors which represent incoming edges. Handles gradient accumulation and backpropagation
minidiff/caching.py: Functions and context manager to allow for caching indices of the topologically-sorted graph constructed during the forward pass
minidiff/typing.py: Just extra types for convenience to keep type hinting clean and readable
minidiff/utils.py: A few utilities used throughout the codebase ranging from finite-difference testing to a graph visualization tool
minidiff/ops/definitions.py: Every operation natively exported by minidiff is defined here, including their forward and backward functions
minidiff/ops/wrapping.py: This file mostly contains code that allows wrapping backend functions as Tensor functions and automatic differentiability to those resulting Tensor functions
minidiff/backend/: Submodule containing logic for selecting a suitable backend and exporting functions from the backend. Currently contains CuPy, NumPy, and MLX