optimistix

<h1 align='center'>Optimistix</h1>

Optimistix 是一个用于非线性求解器的 JAX 库：用于求根、最小化、不动点和最小二乘法。

特点包括：

可互操作的求解器：例如，将求根问题自动转换为最小二乘问题，然后使用最小化算法求解。
模块化优化器：例如，使用带有信赖域更新的狗腿下降路径的 BFGS 二次碗。
使用 PyTree 作为状态。
快速编译和运行时间。
与 Optax 的互操作性。
使用 JAX 的所有优势：自动微分、自动并行、GPU/TPU 支持等。

安装

pip install optimistix

需要 Python 3.9+ 和 JAX 0.4.14+ 以及 Equinox 0.11.0+。

文档

可在 https://docs.kidger.site/optimistix 获取。

快速示例

import jax.numpy as jnp
import optimistix as optx

# 让我们用隐式欧拉法求解微分方程 dy/dt=tanh(y(t))。
# 我们需要找到 y1，使得 y1 = y0 + tanh(y1)dt。

y0 = jnp.array(1.)
dt = jnp.array(0.1)

def fn(y, args):
    return y0 + jnp.tanh(y) * dt

solver = optx.Newton(rtol=1e-5, atol=1e-5)
sol = optx.fixed_point(fn, solver, y0)
y1 = sol.value  # 满足 y1 == fn(y1)

引用

如果您在学术工作中发现这个库很有用，请引用：（arXiv 链接）

@article{optimistix2024,
    title={Optimistix: modular optimisation in JAX and Equinox},
    author={Jason Rader and Terry Lyons and Patrick Kidger},
    journal={arXiv:2402.09983},
    year={2024},
}