Autograd
Grilly provides automatic differentiation through the Variable class. Operations build a computation graph during the forward pass; calling .backward() traverses the graph in reverse to compute gradients.
Variable
Variable wraps a numpy array and tracks operations for gradient computation:
from grilly.nn import Variable, tensor
x = Variable(tensor([1.0, 2.0, 3.0]), requires_grad=True)
y = (x * x).sum()
y.backward()
print(x.grad) # [2.0, 4.0, 6.0]
Set requires_grad=True to track gradients. Intermediate results inherit gradient tracking from their inputs.
Operator Overloading
Variable supports natural Python syntax through operator overloading:
a = Variable(tensor([2.0, 3.0]), requires_grad=True)
b = Variable(tensor([4.0, 5.0]), requires_grad=True)
# Arithmetic
c = a + b
c = a - b
c = a * b
c = a / b
c = a ** 2
# Reductions
s = c.sum()
m = c.mean()
# Chain operations
loss = ((a * b - 1.0) ** 2).mean()
loss.backward()
print(a.grad)
print(b.grad)
Supported Operations
Arithmetic
add, sub, mul, div, neg, pow, matmul
Activations
relu, sigmoid, tanh, gelu, silu, leaky_relu, elu, softplus, softmax
Reductions
sum, mean, max, min, var, std, norm
Shape Operations
reshape, transpose, squeeze, unsqueeze, flatten, view, expand, repeat, permute, contiguous, clone, concat, stack
Trigonometric
sin, cos, tan, asin, acos, atan, atan2
Math
exp, log, sqrt, abs, clamp
Comparisons
eq, ne, lt, le, gt, ge, where
Loss Functions
cross_entropy, mse_loss, l1_loss, smooth_l1_loss, bce_loss, bce_with_logits_loss, nll_loss, kl_div_loss
Context Managers
from grilly.nn import no_grad, enable_grad, is_grad_enabled
# Disable gradient tracking (inference, evaluation)
with no_grad():
output = model(x)
# No computation graph built
# Re-enable inside a no_grad block
with no_grad():
with enable_grad():
y = Variable(tensor([1.0]), requires_grad=True)
z = y * 2
z.backward() # Works
Factory Functions
from grilly.nn import tensor, zeros, ones, randn, rand, linspace, arange, eye, full
x = tensor([1.0, 2.0, 3.0])
z = zeros(3, 4)
o = ones(2, 3)
r = randn(5, 5)
e = eye(4)
Custom Functions
Extend autograd with custom forward/backward operations:
from grilly.nn import Function, FunctionCtx
class MyReLU(Function):
@staticmethod
def forward(ctx: FunctionCtx, x):
ctx.save_for_backward(x)
return x * (x > 0)
@staticmethod
def backward(ctx: FunctionCtx, grad_output):
x, = ctx.saved_tensors
return grad_output * (x > 0)
# Use it
x = Variable(tensor([-1.0, 0.5, 2.0]), requires_grad=True)
y = MyReLU.apply(x)
y.sum().backward()
print(x.grad) # [0.0, 1.0, 1.0]
Backend Autograd (GradientTape)
The backend also provides a lower-level GradientTape for recording operations at the Vulkan dispatch level:
from grilly.backend.autograd_core import GradientTape
with GradientTape() as tape:
tape.watch(x)
y = backend.fnn.linear(x, w)
grads = tape.gradient(y, x)
This is used internally by nn.Module.backward(). Most users should use Variable instead.