Mojo 中的科学计算：Endia

mojo_hub

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}
}