开源地址:https://github.com/endia-org/Endia
Endia是用于加速科学计算的动态数组库。它提供:
自动微分:计算任意阶导数。
复数支持:使用 Endia 进行高级科学应用。
双重 API:在类似 PyTorch 的命令式或类似 JAX 的功能接口之间进行选择。
JIT 编译:利用 MAX 加速训练和推理。

安装
安装Mojo 和 MAX 🔥 (v24.4.0)
克隆存储库:
git clone https://github.com/endia-org/Endia.git
cd Endia
设置环境:
chmod +x setup.sh
./setup.sh
所需依赖项:torch
,,numpy
。graphviz
这些将由安装脚本自动安装。
一个小例子
在本指南中,我们将演示如何计算简单函数的值、梯度和Hessian(即二阶导数)。首先使用 Endia 的类似 Pytorch 的 API,然后使用更类似 Jax 的函数式 API。在这两个示例中,我们首先定义一个函数foo,该函数接受一个数组并返回其元素平方和。
Pytorch方式
使用 Endia 的命令式(类似 PyTorch)接口时,我们通过对函数的输出调用后向方法来计算函数的梯度。这种命令式风格需要明确管理计算图,包括设置输入数组(即叶节点)以及在计算高阶导数时requires_grad=True
使用后向方法。retain_graph=True
import endia as nd
# Define the function
def foo(x: nd.Array) -> nd.Array:
return nd.sum(x ** 2)
# Initialize variable - requires_grad=True needed!
x = nd.array('[1.0, 2.0, 3.0]', requires_grad=True)
# Compute result, first and second order derivatives
y = foo(x)
y.backward(retain_graph=True)
dy_dx = x.grad()
d2y_dx2 = nd.grad(outs=dy_dx, inputs=x)[nd.Array]
# Print results
print(y) # out: [14.0]
print(dy_dx) # out: [2.0, 4.0, 6.0]
print(d2y_dx2) # out: [2.0, 2.0, 2.0]
JAX方式
使用 Endia 的功能性(类似 JAX)接口时,计算图是隐式处理的。通过grad
在 foo 上调用函数,我们创建一个Callable
计算梯度的函数。这Callable
可以grad
再次传递给函数以计算高阶导数。
import endia as nd
# Define the function
def foo(x: nd.Array) -> nd.Array:
return nd.sum(x ** 2)
# Create callables for the jacobian and hessian
foo_jac = nd.grad(foo)
foo_hes = nd.grad(foo_jac)
# Initialize variable - no requires_grad=True needed
x = nd.array('[1.0, 2.0, 3.0]')
# Compute result and derivatives
y = foo(x)
dy_dx = foo_jac(x)[nd.Array]
dy2_dx2 = foo_hes(x)[nd.Array]
# Print results
print(y) # out: [14.0]
print(dy_dx) # out: [2.0, 4.0, 6.0]
print(dy2_dx2) # out: [2.0, 2.0, 2.0]
还有更多!Endia 可以处理复值函数,可以执行正向和反向模式的自动微分,它甚至有一个内置的 JIT 编译器,可以让事情变得快速。在文档中探索完整的功能列表。
我们的任务
- 🧠先进的人工智能和科学计算:通过清晰易懂的算法突破界限
- 🚀 Mojo 驱动的 Clarity:高性能开源代码,始终保持可读性和 Python 风格
- 📐可解释性:优先考虑清晰度和教育价值,而不是详尽的功能
贡献
欢迎为 Endia 做出贡献!如果您愿意做出贡献,请遵循存储库中CONTRIBUTING.md文件中的贡献指南。
引用
如果您在研究或项目中使用 Endia,请按以下方式引用:
@software{Fehrenbach_Endia_-_Scientific_2024,
author = {Fehrenbach, Tillmann},
license = {Apache-2.0},
month = jul,
title = {{Endia - Scientific Computing in Mojo}},
url = {https://github.com/endia-org/Endia},
version = {24.4.0},
year = {2024}
}