Docs
Autograd

Autograd

Endia provides automatic differentiation capabilities through its autograd system. This module contains functions for computing gradients of scalar-valued functions with respect to input arrays.

Imperative API

backward

def backward(arg: Array, retain_graph: Bool) -> None
Computes the gradients of the output Array with respect to the input Arrays.
Parameters
  • arg (Array): The output Array to compute gradients for.
  • retain_graph (Bool, optional): Whether to retain the computation graph after backward pass. Default is False. This is useful when computing higher-order derivatives.
Note

It follows the same semantics as PyTorch's autograd.backward function. An ancestor of the input Array must have requires_grad=True in order to compute the derivative of the output Array with respect to it.

Example
import endia as nd
 
x = nd.array('[1.0, 2.0, 3.0]', requires_grad=True)
y = nd.array('[4.0, 5.0, 6.0]', requires_grad=True)
z = nd.sum(x * y)
 
backward(z, retain_graph=False)
dx = x.grad()
dy = y.grad()

grad

def grad(outs: List[Array], inputs: List[Array], retain_grads: Bool = True, retain_graph: Bool = False) -> Variant[Array, List[Array]]
Computes the gradient of outs with respect to inputs.
Parameters
  • outs (List[Array]): The output Arrays to compute gradients for.
  • inputs (List[Array]): The input Arrays to compute gradients with respect to.
  • retain_grads (Bool, optional): Whether to retain gradients on intermediate operations. Default is True.
  • retain_graph (Bool, optional): Whether to retain the computation graph after backward pass. Default is False.
Returns
  • Variant[Array, List[Array]]: The computed gradients. Returns a single Array if there's only one input, otherwise returns a List of Arrays.
Note

It follows the same semantics as PyTorch's autograd.grad function.

Example
import endia as nd
 
x = nd.array('[1.0, 2.0, 3.0]', requires_grad=True)
y = nd.array('[4.0, 5.0, 6.0]', requires_grad=True)
z = nd.sum(x * y)
 
dx, dy = nd.grad(inputs=z, outputs=List(x, y))

Functional API

grad

def grad(f: Variant[Callable, def (List[Array]) -> Array], argnums: List[Int] = List(-1)) -> Callable
Returns a function that computes the derivative of f.
Parameters
  • f (Variant[Callable, Callable[List[Array], Array]]): A function that takes one or more Arrays as input and returns an Array or a Callable.
  • argnums (List[Int], optional): Indices of the arguments to differentiate with respect to. Default is [-1], which corresponds to computing the derivative of the output of f with respect to all arguments.
Returns
  • Callable: A function that computes the gradient of f.
Note

It follows the same semantics as JAX's grad function.

Example
import endia as nd
 
def f(args: List[Array]):
    return nd.sum(args[0] * args[1])
 
grad_f = nd.grad(f)
x = nd.array('[1.0, 2.0, 3.0]')
y = nd.array('[4.0, 5.0, 6.0]')
dx, dy = grad_f(List(x, y))

value_and_grad

def value_and_grad(arg: Variant[Callable, def (List[Array]) -> Array], argnums: List[Int] = List(-1))-> Callable
Returns a function that computes both the value and the gradient of f.
Parameters
  • f (Callable): A function that takes one or more Arrays as input and returns an Array.
  • argnums (List[Int], optional): Indices of the arguments to differentiate with respect to. Default is [-1], which corresponds to computing the gradient with respect to all arguments.
Returns
  • Callable: A function that computes both the value and the gradient of f.
Example
import endia as nd
 
def f(x, y):
    return nd.sum(x * y)
 
val_grad_f = nd.value_and_grad(f)
x = nd.array('[1.0, 2.0, 3.0]')
y = nd.array('[4.0, 5.0, 6.0]')
val, grads = val_grad_f(x, y)[List[List[Array]]]

Automatic Differentiation Explained

What is Autograd?

Endia uses reverse-mode automatic differentiation to compute gradients. When you call backward() on an Array or use the grad() function, it builds a computation graph and performs a backward pass to compute gradients.

The backward() function propagates gradients from the output to the inputs, accumulating gradients for any intermediate Arrays that require gradients (requires_grad=True).

Learn more about Automatic Differentiation in this excellent video (opens in a new tab).

WIP: Forward-mode AutoDiff
Spacial OpsCustom Ops