Endia is a dynamic Machine Learning Library, featuring:
- Algorithmic Differentiation: Compute derivatives of arbitrary order for training Neural Networks and beyond.
- Signal Processing: Complex Numbers, Fourier Transforms, Convolution, and more.
- JIT Compilation: Leverage the MAX Engine to speed up your code.
- Dual API: Choose between a PyTorch-like imperative or a JAX-like functional interface.
Checkout the Endia documentation for more information.
Prerequisites: Mojo 24.6 🔥 and Magic
-
Add the Modular Community channel to your project's
.tomlfile and install the package:magic project channel add "https://repo.prefix.dev/modular-community" magic add endia
-
Clone the Repository
git clone https://github.com/endia-ai/Endia.git cd Endia magic shell -
Run the Examples, Tests and Benchmarks
mojo run_all.mojo
Recommended: Go to the
run_all.mojofile and adjust the execution to your liking.
In this guide, we'll demonstrate how to compute the value, gradient, and the Hessian (i.e. the second-order derivative) of a simple function. First by using Endia's Pytorch-like API and then by using a more Jax-like functional API. In both examples, we initially define a function foo that takes an array and returns the sum of the squares of its elements.
When using Endia's imperative (PyTorch-like) interface, we compute the gradient of a function by calling the backward method on the function's output. This imperative style requires explicit management of the computational graph, including setting requires_grad=True for the input arrays (i.e. leaf nodes) and using create_graph=True in the backward method when computing higher-order derivatives.
from endia import Tensor, sum, arange
import endia.autograd.functional as F
# Define the function
def foo(x: Tensor) -> Tensor:
return sum(x ** 2)
def main():
# Initialize variable - requires_grad=True needed!
x = arange(1.0, 4.0, requires_grad=True) # [1.0, 2.0, 3.0]
# Compute result, first and second order derivatives
y = foo(x)
y.backward(create_graph=True)
dy_dx = x.grad()
d2y_dx2 = F.grad(outs=sum(dy_dx), inputs=x)[0]
# Print results
print(y) # 14.0
print(dy_dx) # [2.0, 4.0, 6.0]
print(d2y_dx2) # [2.0, 2.0, 2.0]
When using Endia's functional (JAX-like) interface, the computational graph is handled implicitly. By calling the grad or jacobian function on foo, we create a Callable which computes the full Jacobian matrix. This Callable can be passed to the grad or jacobian function again to compute higher-order derivatives.
from endia import grad, jacobian
from endia.numpy import sum, arange, ndarray
def foo(x: ndarray) -> ndarray:
return sum(x**2)
def main():
# create Callables for the first and second order derivatives
foo_jac = grad(foo)
foo_hes = jacobian(foo_jac)
x = arange(1.0, 4.0) # [1.0, 2.0, 3.0]
print(foo(x)) # 14.0
print(foo_jac(x)[ndarray]) # [2.0, 4.0, 6.0]
print(foo_hes(x)[ndarray]) # [[2.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 2.0]]And there is so much more! Endia can handle complex valued functions, can perform both forward and reverse-mode automatic differentiation, it even has a builtin JIT compiler to make things go brrr. Explore the full list of features in the documentation.
Contributions to Endia are welcome! If you'd like to contribute, please follow the contribution guidelines in the CONTRIBUTING.md file in the repository.
If you use Endia in your research or project, please cite it as follows:
@software{Fehrenbach_Endia_2025,
author = {Fehrenbach, Tillmann},
license = {Apache-2.0 with LLVM Exceptions},
doi = {10.5281/zenodo.12810766},
month = {01},
title = {{Endia}},
url = {https://github.com/endia-ai/Endia},
version = {24.6.0},
year = {2025}
}Endia is licensed under the Apache-2.0 license with LLVM Exeptions.
